ColossalAI
An integrated large-scale model training system with efficient parallelization techniques
Installation
PyPI
pip install colossalai
Install From Source
git clone [email protected]: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)
pip install -v --no-cache-dir --global-option="--cuda_ext" .
Documentation
Quick View
Start Distributed Training in Lines
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.
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.
![[BUG]: Memory consumption by fp16 is not normal](https://avatars.githubusercontent.com/u/24365567?v=4)
![[BUG]: RuntimeError of](https://avatars.githubusercontent.com/u/27288110?v=4)
![[BUG]: type object 'ChunkManager' has no attribute 'search_chunk_size'](https://avatars.githubusercontent.com/u/67773022?v=4)
![[BUG]: colossalai/kernel/cuda_native/csrc/moe_cuda_kernel.cu:5:10: fatal error: cub/cub.cuh: No such file or directory (update: now with more build errors!)](https://avatars.githubusercontent.com/u/7685285?v=4)
![[BUG]: Issue with Colossal-AI on Cuda 11.4 and Docker ?](https://avatars.githubusercontent.com/u/63513596?v=4)
![[BUG]: examples/images/diffusion ran failed](https://avatars.githubusercontent.com/u/24288375?v=4)
![[BUG]: RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one.](https://avatars.githubusercontent.com/u/69774354?v=4)
![[BUG]: CUDA extension build skipped when installing from source](https://avatars.githubusercontent.com/u/50511903?v=4)
![[workflow]New version: Create workflow files for examples' auto check](https://avatars.githubusercontent.com/u/99854690?v=4)
![[BUG]: FusedScaleMaskSoftmax last dimension does not sum to 1](https://avatars.githubusercontent.com/u/22529082?v=4)
![[BUG] [DIFFUSION]: Sampling Fails to produce output](https://avatars.githubusercontent.com/u/110270449?v=4)
The only minute error that I have that even comes up is that I don't have a validation set so using the default metrics throws a warning about the metrics not being passed. I did though train a new model with new metrics that threw no error and got the same result. If there are things you want me to test I can do that.
I don't know if you'd prefer I open another issue or just tack it on here but is using Triton fully supported? I got a proper Triton install to stop a warning from being thrown but I was wondering if I should've ignored it. As well I couldn't get any models to properly work when trying to do training with the -r call. Is there a specific model that is compatible with this new version? The error it throws is
RuntimeError: Error(s) in loading state_dict for GeminiDDP:
Missing keys in state dict: _forward_modulelvlb_weight, _forward_module.cond_state_model.attn_mask
Unexpected Keys in state dict: _forward_module.model_ema.decay _foward_model.model_ema.num_updates
48 Dec 30, 2022
45 Jan 1, 2023
185 Dec 21, 2022
607 Dec 31, 2022
72 Dec 16, 2022
130 Dec 2, 2022
432 Dec 16, 2022
33 Jun 27, 2021
36 Oct 30, 2022
137 Dec 15, 2022
128 Nov 29, 2022
43 Dec 26, 2022
290 Dec 29, 2022
294 Dec 12, 2022
37 Dec 8, 2022