README.md 6.56 KB
Newer Older
1
# Colossal-AI
2

3
[![logo](./docs/images/Colossal-AI_logo.png)](https://www.colossalai.org/)
4

5
<div align="center">
6
7
8
9
   <h3> <a href="https://arxiv.org/abs/2110.14883"> Paper </a> | 
   <a href="https://www.colossalai.org/"> Documentation </a> | 
   <a href="https://github.com/hpcaitech/ColossalAI-Examples"> Examples </a> |   
   <a href="https://github.com/hpcaitech/ColossalAI/discussions"> Forum </a> | 
10
11
12
   <a href="https://medium.com/@hpcaitech"> Blog </a></h3> 
   <br/>

13
   [![Build](https://github.com/hpcaitech/ColossalAI/actions/workflows/PR_CI.yml/badge.svg)](https://github.com/hpcaitech/ColossalAI/actions/workflows/PR_CI.yml)
14
   [![Documentation](https://readthedocs.org/projects/colossalai/badge/?version=latest)](https://colossalai.readthedocs.io/en/latest/?badge=latest)
Frank Lee's avatar
Frank Lee committed
15
   [![codebeat badge](https://codebeat.co/badges/bfe8f98b-5d61-4256-8ad2-ccd34d9cc156)](https://codebeat.co/projects/github-com-hpcaitech-colossalai-main)
binmakeswell's avatar
binmakeswell committed
16
17
   [![slack badge](https://img.shields.io/badge/Slack-join-blueviolet?logo=slack&amp)](https://join.slack.com/t/colossalaiworkspace/shared_invite/zt-z7b26eeb-CBp7jouvu~r0~lcFzX832w)
   [![WeChat badge](https://img.shields.io/badge/微信-加入-green?logo=wechat&amp)](./docs/images/WeChat.png)
binmakeswell's avatar
binmakeswell committed
18
19

   | [English](README.md) | [中文](README-zh-Hans.md) |
20
</div>
ver217's avatar
ver217 committed
21
22
An integrated large-scale model training system with efficient parallelization techniques.

binmakeswell's avatar
binmakeswell committed
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

## Features

Colossal-AI 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
- Pipeline Parallelism
- 1D, 2D, 2.5D, 3D tensor parallelism
- Sequence parallelism
- Friendly trainer and engine
- Extensible for new parallelism
- Mixed Precision Training
- Zero Redundancy Optimizer (ZeRO)

## Examples
### ViT

42
<img src="./docs/images/ViT.png" width="450" />
binmakeswell's avatar
binmakeswell committed
43

44
- 14x larger batch size, and 5x faster training for Tensor Parallel = 64
binmakeswell's avatar
binmakeswell committed
45

46
### GPT-3
binmakeswell's avatar
binmakeswell committed
47

48
<img src="./docs/images/GPT3.png" width=700/>
binmakeswell's avatar
binmakeswell committed
49

50
- Free 50% GPU resources, or 10.7% acceleration
51
52

### GPT-2
53
<img src="./docs/images/GPT2.png" width=800/>
54

55
- 11x lower GPU RAM, or superlinear scaling
56

binmakeswell's avatar
binmakeswell committed
57
### BERT
58
<img src="./docs/images/BERT.png" width=800/>
binmakeswell's avatar
binmakeswell committed
59

60
- 2x faster training, or 50% longer sequence length
binmakeswell's avatar
binmakeswell committed
61
62
63
64

Please visit our [documentation and tutorials](https://www.colossalai.org/) for more details.


zbian's avatar
zbian committed
65
66
## Installation

ver217's avatar
ver217 committed
67
68
69
70
71
72
### PyPI

```bash
pip install colossalai
```
This command will install CUDA extension if your have installed CUDA, NVCC and torch. 
73

ver217's avatar
ver217 committed
74
75
76
77
78
79
80
81
82
83
84
85
If you don't want to install CUDA extension, you should add `--global-option="--no_cuda_ext"`, like:
```bash
pip install colossalai --global-option="--no_cuda_ext"
```

If you want to use `ZeRO`, you can run:
```bash
pip install colossalai[zero]
```

### Install From Source

binmakeswell's avatar
binmakeswell committed
86
> The version of Colossal-AI will be in line with the main branch of the repository. Feel free to raise an issue if you encounter any problem. :)
zbian's avatar
zbian committed
87
88

```shell
89
git clone https://github.com/hpcaitech/ColossalAI.git
zbian's avatar
zbian committed
90
91
92
93
94
95
96
97
cd ColossalAI
# install dependency
pip install -r requirements/requirements.txt

# install colossalai
pip install .
```

ver217's avatar
ver217 committed
98
If you don't want to install and enable CUDA kernel fusion (compulsory installation when using fused optimizer):
zbian's avatar
zbian committed
99
100

```shell
ver217's avatar
ver217 committed
101
pip install --global-option="--no_cuda_ext" .
zbian's avatar
zbian committed
102
103
```

binmakeswell's avatar
binmakeswell committed
104

Frank Lee's avatar
Frank Lee committed
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
## Use Docker

Run the following command to build a docker image from Dockerfile provided.

```bash
cd ColossalAI
docker build -t colossalai ./docker
```

Run the following command to start the docker container in interactive mode.

```bash
docker run -ti --gpus all --rm --ipc=host colossalai bash
```

binmakeswell's avatar
binmakeswell committed
120
121
122
123
124
125
126
127
128

## Community

Join the Colossal-AI community on [Forum](https://github.com/hpcaitech/ColossalAI/discussions),
[Slack](https://join.slack.com/t/colossalaiworkspace/shared_invite/zt-z7b26eeb-CBp7jouvu~r0~lcFzX832w),
and [WeChat](./docs/images/WeChat.png "qrcode") to share your suggestions, advice, and questions with our engineering team.



129
130
## Contributing

binmakeswell's avatar
binmakeswell committed
131
132
133
If you wish to contribute to this project, please follow the guideline in [Contributing](./CONTRIBUTING.md).

Thanks so much to all of our amazing contributors!
134

binmakeswell's avatar
binmakeswell committed
135
136
137
<a href="https://github.com/hpcaitech/ColossalAI/graphs/contributors"><img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/contributor_avatar.png" width="800px"></a>

*The order of contributor avatars is randomly shuffled.*
138

zbian's avatar
zbian committed
139
140
141
142
143
144
## Quick View

### Start Distributed Training in Lines

```python
import colossalai
145
146
147
148
149
150
151
152
153
154
155
156
157
from colossalai.utils import get_dataloader


# my_config can be path to config file or a dictionary obj
# 'localhost' is only for single node, you need to specify
# the node name if using multiple nodes
colossalai.launch(
    config=my_config,
    rank=rank,
    world_size=world_size,
    backend='nccl',
    port=29500,
    host='localhost'
zbian's avatar
zbian committed
158
)
159
160

# build your model
161
model = ...
162

163
# build you dataset, the dataloader will have distributed data
164
# sampler by default
165
train_dataset = ...
166
train_dataloader = get_dataloader(dataset=dataset,
167
                                shuffle=True
168
                                )
169
170


binmakeswell's avatar
binmakeswell committed
171
# build your optimizer
172
optimizer = ...
173
174
175
176

# build your loss function
criterion = ...

binmakeswell's avatar
binmakeswell committed
177
# initialize colossalai
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
engine, train_dataloader, _, _ = colossalai.initialize(
    model=model,
    optimizer=optimizer,
    criterion=criterion,
    train_dataloader=train_dataloader
)

# start training
engine.train()
for epoch in range(NUM_EPOCHS):
    for data, label in train_dataloader:
        engine.zero_grad()
        output = engine(data)
        loss = engine.criterion(output, label)
        engine.backward(loss)
        engine.step()

zbian's avatar
zbian committed
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
```

### 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

```

221

zbian's avatar
zbian committed
222

223
## Cite Us
zbian's avatar
zbian committed
224

225
226
227
228
229
230
231
232
```
@article{bian2021colossal,
  title={Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training},
  author={Bian, Zhengda and Liu, Hongxin and Wang, Boxiang and Huang, Haichen and Li, Yongbin and Wang, Chuanrui and Cui, Fan and You, Yang},
  journal={arXiv preprint arXiv:2110.14883},
  year={2021}
}
```