mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
300 lines
11 KiB
300 lines
11 KiB
# Colossal-AI
|
|
<div id="top" align="center">
|
|
|
|
[](https://www.colossalai.org/)
|
|
|
|
Colossal-AI: A Unified Deep Learning System for Big Model Era
|
|
|
|
<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>
|
|
|
|
[](https://github.com/hpcaitech/ColossalAI/actions/workflows/build.yml)
|
|
[](https://colossalai.readthedocs.io/en/latest/?badge=latest)
|
|
[](https://www.codefactor.io/repository/github/hpcaitech/colossalai)
|
|
[](https://huggingface.co/hpcai-tech)
|
|
[](https://join.slack.com/t/colossalaiworkspace/shared_invite/zt-z7b26eeb-CBp7jouvu~r0~lcFzX832w)
|
|
[](https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/WeChat.png)
|
|
|
|
|
|
| [English](README.md) | [中文](README-zh-Hans.md) |
|
|
|
|
</div>
|
|
|
|
## Table of Contents
|
|
<ul>
|
|
<li><a href="#Why-Colossal-AI">Why Colossal-AI</a> </li>
|
|
<li><a href="#Features">Features</a> </li>
|
|
<li>
|
|
<a href="#Parallel-Training-Demo">Parallel Training Demo</a>
|
|
<ul>
|
|
<li><a href="#ViT">ViT</a></li>
|
|
<li><a href="#GPT-3">GPT-3</a></li>
|
|
<li><a href="#GPT-2">GPT-2</a></li>
|
|
<li><a href="#BERT">BERT</a></li>
|
|
<li><a href="#PaLM">PaLM</a></li>
|
|
<li><a href="#OPT">OPT</a></li>
|
|
</ul>
|
|
</li>
|
|
<li>
|
|
<a href="#Single-GPU-Training-Demo">Single GPU Training Demo</a>
|
|
<ul>
|
|
<li><a href="#GPT-2-Single">GPT-2</a></li>
|
|
<li><a href="#PaLM-Single">PaLM</a></li>
|
|
</ul>
|
|
</li>
|
|
<li>
|
|
<a href="#Inference-Energon-AI-Demo">Inference (Energon-AI) Demo</a>
|
|
<ul>
|
|
<li><a href="#GPT-3-Inference">GPT-3</a></li>
|
|
</ul>
|
|
</li>
|
|
<li>
|
|
<a href="#Installation">Installation</a>
|
|
<ul>
|
|
<li><a href="#PyPI">PyPI</a></li>
|
|
<li><a href="#Install-From-Source">Install From Source</a></li>
|
|
</ul>
|
|
</li>
|
|
<li><a href="#Use-Docker">Use Docker</a></li>
|
|
<li><a href="#Community">Community</a></li>
|
|
<li><a href="#contributing">Contributing</a></li>
|
|
<li><a href="#Quick-View">Quick View</a></li>
|
|
<ul>
|
|
<li><a href="#Start-Distributed-Training-in-Lines">Start Distributed Training in Lines</a></li>
|
|
<li><a href="#Write-a-Simple-2D-Parallel-Model">Write a Simple 2D Parallel Model</a></li>
|
|
</ul>
|
|
<li><a href="#Cite-Us">Cite Us</a></li>
|
|
</ul>
|
|
|
|
## Why Colossal-AI
|
|
<div align="center">
|
|
<a href="https://youtu.be/KnXSfjqkKN0">
|
|
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/JamesDemmel_Colossal-AI.png" width="600" />
|
|
</a>
|
|
|
|
Prof. James Demmel (UC Berkeley): Colossal-AI makes training AI models efficient, easy, and scalable.
|
|
</div>
|
|
|
|
<p align="right">(<a href="#top">back to top</a>)</p>
|
|
|
|
## Features
|
|
|
|
Colossal-AI provides a collection of parallel components for you. We aim to support you to write your
|
|
distributed deep learning models just like how you write your model on your laptop. We provide user-friendly tools to kickstart
|
|
distributed training and inference in a few lines.
|
|
|
|
- Parallelism strategies
|
|
- Data Parallelism
|
|
- Pipeline Parallelism
|
|
- 1D, [2D](https://arxiv.org/abs/2104.05343), [2.5D](https://arxiv.org/abs/2105.14500), [3D](https://arxiv.org/abs/2105.14450) Tensor Parallelism
|
|
- [Sequence Parallelism](https://arxiv.org/abs/2105.13120)
|
|
- [Zero Redundancy Optimizer (ZeRO)](https://arxiv.org/abs/1910.02054)
|
|
|
|
- Heterogeneous Memory Management
|
|
- [PatrickStar](https://arxiv.org/abs/2108.05818)
|
|
|
|
- Friendly Usage
|
|
- Parallelism based on configuration file
|
|
|
|
- Inference
|
|
- [Energon-AI](https://github.com/hpcaitech/EnergonAI)
|
|
|
|
<p align="right">(<a href="#top">back to top</a>)</p>
|
|
|
|
## Parallel Training Demo
|
|
### ViT
|
|
<p align="center">
|
|
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/ViT.png" width="450" />
|
|
</p>
|
|
|
|
- 14x larger batch size, and 5x faster training for Tensor Parallelism = 64
|
|
|
|
### GPT-3
|
|
<p align="center">
|
|
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/GPT3-v5.png" width=700/>
|
|
</p>
|
|
|
|
- Save 50% GPU resources, and 10.7% acceleration
|
|
|
|
### GPT-2
|
|
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/GPT2.png" width=800/>
|
|
|
|
- 11x lower GPU memory consumption, and superlinear scaling efficiency with Tensor Parallelism
|
|
|
|
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/(updated)GPT-2.png" width=800>
|
|
|
|
- 24x larger model size on the same hardware
|
|
- over 3x acceleration
|
|
### BERT
|
|
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/BERT.png" width=800/>
|
|
|
|
- 2x faster training, or 50% longer sequence length
|
|
|
|
### PaLM
|
|
- [PaLM-colossalai](https://github.com/hpcaitech/PaLM-colossalai): Scalable implementation of Google's Pathways Language Model ([PaLM](https://ai.googleblog.com/2022/04/pathways-language-model-palm-scaling-to.html)).
|
|
|
|
### OPT
|
|
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/OPT.png" width=800/>
|
|
|
|
- [Open Pretrained Transformer (OPT)](https://github.com/facebookresearch/metaseq), a 175-Billion parameter AI language model released by Meta, which stimulates AI programmers to perform various downstream tasks and application deployments because public pretrained model weights.
|
|
- 40% speedup fine-tuning OPT at low cost in lines. [[Example]](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/language/opt)
|
|
|
|
Please visit our [documentation](https://www.colossalai.org/) and [examples](https://github.com/hpcaitech/ColossalAI-Examples) for more details.
|
|
|
|
<p align="right">(<a href="#top">back to top</a>)</p>
|
|
|
|
## Single GPU Training Demo
|
|
|
|
### GPT-2
|
|
<p id="GPT-2-Single" align="center">
|
|
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/GPT2-GPU1.png" width=450/>
|
|
</p>
|
|
|
|
- 20x larger model size on the same hardware
|
|
|
|
<p id="GPT-2-NVME" align="center">
|
|
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/GPT2-NVME.png" width=800/>
|
|
</p>
|
|
|
|
- 120x larger model size on the same hardware (RTX 3080)
|
|
|
|
### PaLM
|
|
<p id="PaLM-Single" align="center">
|
|
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/PaLM-GPU1.png" width=450/>
|
|
</p>
|
|
|
|
- 34x larger model size on the same hardware
|
|
|
|
<p align="right">(<a href="#top">back to top</a>)</p>
|
|
|
|
|
|
## Inference (Energon-AI) Demo
|
|
|
|
### GPT-3
|
|
<p id="GPT-3-Inference" align="center">
|
|
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/inference_GPT-3.jpg" width=800/>
|
|
</p>
|
|
|
|
- [Energon-AI](https://github.com/hpcaitech/EnergonAI): 50% inference acceleration on the same hardware
|
|
|
|
<p align="right">(<a href="#top">back to top</a>)</p>
|
|
|
|
## Installation
|
|
|
|
### Download From Official Releases
|
|
|
|
You can visit the [Download](https://www.colossalai.org/download) page to download Colossal-AI with pre-built CUDA extensions.
|
|
|
|
|
|
### Download From Source
|
|
|
|
> 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. :)
|
|
|
|
```shell
|
|
git clone https://github.com/hpcaitech/ColossalAI.git
|
|
cd ColossalAI
|
|
|
|
# install dependency
|
|
pip install -r requirements/requirements.txt
|
|
|
|
# install colossalai
|
|
pip install .
|
|
```
|
|
|
|
If you don't want to install and enable CUDA kernel fusion (compulsory installation when using fused optimizer):
|
|
|
|
```shell
|
|
NO_CUDA_EXT=1 pip install .
|
|
```
|
|
|
|
<p align="right">(<a href="#top">back to top</a>)</p>
|
|
|
|
## Use Docker
|
|
|
|
### Pull from DockerHub
|
|
|
|
You can directly pull the docker image from our [DockerHub page](https://hub.docker.com/r/hpcaitech/colossalai). The image is automatically uploaded upon release.
|
|
|
|
|
|
### Build On Your Own
|
|
|
|
Run the following command to build a docker image from Dockerfile provided.
|
|
|
|
> Building Colossal-AI from scratch requires GPU support, you need to use Nvidia Docker Runtime as the default when doing `docker build`. More details can be found [here](https://stackoverflow.com/questions/59691207/docker-build-with-nvidia-runtime).
|
|
> We recommend you install Colossal-AI from our [project page](https://www.colossalai.org) directly.
|
|
|
|
|
|
```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
|
|
```
|
|
|
|
<p align="right">(<a href="#top">back to top</a>)</p>
|
|
|
|
## 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](https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/WeChat.png "qrcode") to share your suggestions, feedback, and questions with our engineering team.
|
|
|
|
## Contributing
|
|
|
|
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!
|
|
|
|
<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.*
|
|
|
|
<p align="right">(<a href="#top">back to top</a>)</p>
|
|
|
|
## Quick View
|
|
|
|
### Start Distributed Training in Lines
|
|
|
|
```python
|
|
parallel = dict(
|
|
pipeline=2,
|
|
tensor=dict(mode='2.5d', depth = 1, size=4)
|
|
)
|
|
```
|
|
|
|
### Start Heterogeneous Training in Lines
|
|
|
|
```python
|
|
zero = dict(
|
|
model_config=dict(
|
|
tensor_placement_policy='auto',
|
|
shard_strategy=TensorShardStrategy(),
|
|
reuse_fp16_shard=True
|
|
),
|
|
optimizer_config=dict(initial_scale=2**5, gpu_margin_mem_ratio=0.2)
|
|
)
|
|
|
|
```
|
|
|
|
<p align="right">(<a href="#top">back to top</a>)</p>
|
|
|
|
## Cite Us
|
|
|
|
```
|
|
@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}
|
|
}
|
|
```
|
|
|
|
<p align="right">(<a href="#top">back to top</a>)</p>
|