forked from hpcaitech/ColossalAI
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
2ebaefc
commit 404ecbd
Showing
409 changed files
with
35,853 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,144 @@ | ||
# Byte-compiled / optimized / DLL files | ||
__pycache__/ | ||
*.py[cod] | ||
*$py.class | ||
|
||
# C extensions | ||
*.so | ||
|
||
# Distribution / packaging | ||
.Python | ||
build/ | ||
develop-eggs/ | ||
dist/ | ||
downloads/ | ||
eggs/ | ||
.eggs/ | ||
lib/ | ||
lib64/ | ||
parts/ | ||
sdist/ | ||
var/ | ||
wheels/ | ||
pip-wheel-metadata/ | ||
share/python-wheels/ | ||
*.egg-info/ | ||
.installed.cfg | ||
*.egg | ||
MANIFEST | ||
|
||
# PyInstaller | ||
# Usually these files are written by a python script from a template | ||
# before PyInstaller builds the exe, so as to inject date/other infos into it. | ||
*.manifest | ||
*.spec | ||
|
||
# Installer logs | ||
pip-log.txt | ||
pip-delete-this-directory.txt | ||
|
||
# Unit test / coverage reports | ||
htmlcov/ | ||
.tox/ | ||
.nox/ | ||
.coverage | ||
.coverage.* | ||
.cache | ||
nosetests.xml | ||
coverage.xml | ||
*.cover | ||
*.py,cover | ||
.hypothesis/ | ||
.pytest_cache/ | ||
|
||
# Translations | ||
*.mo | ||
*.pot | ||
|
||
# Django stuff: | ||
*.log | ||
local_settings.py | ||
db.sqlite3 | ||
db.sqlite3-journal | ||
|
||
# Flask stuff: | ||
instance/ | ||
.webassets-cache | ||
|
||
# Scrapy stuff: | ||
.scrapy | ||
|
||
# Sphinx documentation | ||
docs/_build/ | ||
docs/.build/ | ||
|
||
# PyBuilder | ||
target/ | ||
|
||
# Jupyter Notebook | ||
.ipynb_checkpoints | ||
|
||
# IPython | ||
profile_default/ | ||
ipython_config.py | ||
|
||
# pyenv | ||
.python-version | ||
|
||
# pipenv | ||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. | ||
# However, in case of collaboration, if having platform-specific dependencies or dependencies | ||
# having no cross-platform support, pipenv may install dependencies that don't work, or not | ||
# install all needed dependencies. | ||
#Pipfile.lock | ||
|
||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow | ||
__pypackages__/ | ||
|
||
# Celery stuff | ||
celerybeat-schedule | ||
celerybeat.pid | ||
|
||
# SageMath parsed files | ||
*.sage.py | ||
|
||
# Environments | ||
.env | ||
.venv | ||
env/ | ||
venv/ | ||
ENV/ | ||
env.bak/ | ||
venv.bak/ | ||
|
||
# Spyder project settings | ||
.spyderproject | ||
.spyproject | ||
|
||
# Rope project settings | ||
.ropeproject | ||
|
||
# mkdocs documentation | ||
/site | ||
|
||
# mypy | ||
.mypy_cache/ | ||
.dmypy.json | ||
dmypy.json | ||
|
||
# Pyre type checker | ||
.pyre/ | ||
|
||
# IDE | ||
.idea/ | ||
.vscode/ | ||
|
||
# macos | ||
.DS_Store | ||
#data/ | ||
|
||
# launcher setting | ||
tests/launcher/log | ||
tests/launcher/personal | ||
|
||
docs/.build |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,4 @@ | ||
include *.txt README.md | ||
recursive-include requirements *.txt | ||
recursive-include colossalai *.cpp *.h *.cu *.tr *.cuh *.cc | ||
recursive-include csrc *.cpp *.h *.cu *.tr *.cuh *.cc |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,104 @@ | ||
# ColossalAI | ||
|
||
An integrated large-scale model training framework with efficient parallelization techniques | ||
|
||
## Installation | ||
|
||
### PyPI | ||
|
||
```bash | ||
pip install colossalai | ||
``` | ||
|
||
### Install From Source | ||
|
||
```shell | ||
git clone git@github.com:hpcaitech/ColossalAI.git | ||
cd ColossalAI | ||
# install dependency | ||
pip install -r requirements/requirements.txt | ||
|
||
# install colossalai | ||
pip install . | ||
``` | ||
|
||
Install and enable CUDA kernel fusion (compulsory installation when using fused optimizer) | ||
|
||
```shell | ||
pip install -v --no-cache-dir --global-option="--cuda_ext" . | ||
``` | ||
|
||
## Documentation | ||
|
||
- [Documentation](https://www.colossalai.org/) | ||
|
||
## Quick View | ||
|
||
### Start Distributed Training in Lines | ||
|
||
```python | ||
import colossalai | ||
from colossalai.engine import Engine | ||
from colossalai.trainer import Trainer | ||
from colossalai.core import global_context as gpc | ||
|
||
model, train_dataloader, test_dataloader, criterion, optimizer, schedule, lr_scheduler = colossalai.initialize() | ||
engine = Engine( | ||
model=model, | ||
criterion=criterion, | ||
optimizer=optimizer, | ||
lr_scheduler=lr_scheduler, | ||
schedule=schedule | ||
) | ||
|
||
trainer = Trainer(engine=engine, | ||
hooks_cfg=gpc.config.hooks, | ||
verbose=True) | ||
trainer.fit( | ||
train_dataloader=train_dataloader, | ||
test_dataloader=test_dataloader, | ||
max_epochs=gpc.config.num_epochs, | ||
display_progress=True, | ||
test_interval=5 | ||
) | ||
``` | ||
|
||
### Write a Simple 2D Parallel Model | ||
|
||
Let's say we have a huge MLP model and its very large hidden size makes it difficult to fit into a single GPU. We can | ||
then distribute the model weights across GPUs in a 2D mesh while you still write your model in a familiar way. | ||
|
||
```python | ||
from colossalai.nn import Linear2D | ||
import torch.nn as nn | ||
|
||
|
||
class MLP_2D(nn.Module): | ||
|
||
def __init__(self): | ||
super().__init__() | ||
self.linear_1 = Linear2D(in_features=1024, out_features=16384) | ||
self.linear_2 = Linear2D(in_features=16384, out_features=1024) | ||
|
||
def forward(self, x): | ||
x = self.linear_1(x) | ||
x = self.linear_2(x) | ||
return x | ||
|
||
``` | ||
|
||
## Features | ||
|
||
ColossalAI provides a collection of parallel training components for you. We aim to support you to write your | ||
distributed deep learning models just like how you write your single-GPU model. We provide friendly tools to kickstart | ||
distributed training in a few lines. | ||
|
||
- [Data Parallelism](./docs/parallelization.md) | ||
- [Pipeline Parallelism](./docs/parallelization.md) | ||
- [1D, 2D, 2.5D, 3D and sequence parallelism](./docs/parallelization.md) | ||
- [friendly trainer and engine](./docs/trainer_engine.md) | ||
- [Extensible for new parallelism](./docs/add_your_parallel.md) | ||
- [Mixed Precision Training](./docs/amp.md) | ||
- [Zero Redundancy Optimizer (ZeRO)](./docs/zero.md) | ||
|
||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,4 @@ | ||
from .initialize import init_dist, initialize | ||
from .nn import * | ||
|
||
__version__ = '0.0.1' |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,2 @@ | ||
from .builder import * | ||
from .pipeline import ModelInitializer |
Oops, something went wrong.