README.md 4.96 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
10
   <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> | 
   <a href="https://medium.com/@hpcaitech"> Blog </a></h3>
11
   
12
   [![Build](https://github.com/hpcaitech/ColossalAI/actions/workflows/PR_CI.yml/badge.svg)](https://github.com/hpcaitech/ColossalAI/actions/workflows/PR_CI.yml)
13
   [![Documentation](https://readthedocs.org/projects/colossalai/badge/?version=latest)](https://colossalai.readthedocs.io/en/latest/?badge=latest)
Frank Lee's avatar
Frank Lee committed
14
   [![codebeat badge](https://codebeat.co/badges/bfe8f98b-5d61-4256-8ad2-ccd34d9cc156)](https://codebeat.co/projects/github-com-hpcaitech-colossalai-main)
15
</div>
ver217's avatar
ver217 committed
16
17
An integrated large-scale model training system with efficient parallelization techniques.

zbian's avatar
zbian committed
18
19
## Installation

ver217's avatar
ver217 committed
20
21
22
23
24
25
### PyPI

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

ver217's avatar
ver217 committed
27
28
29
30
31
32
33
34
35
36
37
38
39
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

> The documentation 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
40
41

```shell
42
git clone https://github.com/hpcaitech/ColossalAI.git
zbian's avatar
zbian committed
43
44
45
46
47
48
49
50
cd ColossalAI
# install dependency
pip install -r requirements/requirements.txt

# install colossalai
pip install .
```

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

```shell
ver217's avatar
ver217 committed
54
pip install --global-option="--no_cuda_ext" .
zbian's avatar
zbian committed
55
56
```

Frank Lee's avatar
Frank Lee committed
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
## 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
```

zbian's avatar
zbian committed
72
73
74
75
76
77
## Quick View

### Start Distributed Training in Lines

```python
import colossalai
78
79
80
81
82
83
84
85
86
87
88
89
90
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
91
)
92
93

# build your model
94
model = ...
95

96
# build you dataset, the dataloader will have distributed data
97
# sampler by default
98
train_dataset = ...
99
train_dataloader = get_dataloader(dataset=dataset,
100
                                shuffle=True
101
                                )
102
103


104
105
# build your
optimizer = ...
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127

# build your loss function
criterion = ...

# build your lr_scheduler
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
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
```

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

156
Colossal-AI provides a collection of parallel training components for you. We aim to support you to write your
zbian's avatar
zbian committed
157
158
159
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.

160
161
162
163
164
165
166
167
168
- Data Parallelism
- Pipeline Parallelism
- 1D, 2D, 2.5D, 3D and sequence parallelism
- Friendly trainer and engine
- Extensible for new parallelism
- Mixed Precision Training
- Zero Redundancy Optimizer (ZeRO)

Please visit our [documentation and tutorials](https://www.colossalai.org/) for more details.
zbian's avatar
zbian committed
169

170
## Cite Us
zbian's avatar
zbian committed
171

172
173
174
175
176
177
178
179
```
@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}
}
```