|
|
# Colossal-AI
|
|
|
<div id="top" align="center">
|
|
|
|
|
|
[](https://www.colossalai.org/)
|
|
|
|
|
|
一个整合高效并行技术的 AI 大模型训练系统。
|
|
|
|
|
|
<h3> <a href="https://arxiv.org/abs/2110.14883"> 论文 </a> |
|
|
|
<a href="https://www.colossalai.org/"> 文档 </a> |
|
|
|
<a href="https://github.com/hpcaitech/ColossalAI-Examples"> 例程 </a> |
|
|
|
<a href="https://github.com/hpcaitech/ColossalAI/discussions"> 论坛 </a> |
|
|
|
<a href="https://medium.com/@hpcaitech"> 博客 </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>
|
|
|
|
|
|
|
|
|
## 目录
|
|
|
<ul>
|
|
|
<li><a href="#为何选择-Colossal-AI">为何选择 Colossal-AI</a> </li>
|
|
|
<li><a href="#特点">特点</a> </li>
|
|
|
<li>
|
|
|
<a href="#并行样例展示">并行样例展示</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>
|
|
|
</ul>
|
|
|
</li>
|
|
|
<li>
|
|
|
<a href="#单GPU样例展示">单GPU样例展示</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="#安装">安装</a>
|
|
|
<ul>
|
|
|
<li><a href="#PyPI">PyPI</a></li>
|
|
|
<li><a href="#从源代码安装">从源代码安装</a></li>
|
|
|
</ul>
|
|
|
</li>
|
|
|
<li><a href="#使用-Docker">使用 Docker</a></li>
|
|
|
<li><a href="#社区">社区</a></li>
|
|
|
<li><a href="#做出贡献">做出贡献</a></li>
|
|
|
<li><a href="#快速预览">快速预览</a></li>
|
|
|
<ul>
|
|
|
<li><a href="#几行代码开启分布式训练">几行代码开启分布式训练</a></li>
|
|
|
<li><a href="#构建一个简单的2维并行模型">构建一个简单的2维并行模型</a></li>
|
|
|
</ul>
|
|
|
<li><a href="#引用我们">引用我们</a></li>
|
|
|
</ul>
|
|
|
|
|
|
## 为何选择 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>
|
|
|
|
|
|
James Demmel 教授 (加州大学伯克利分校): Colossal-AI 让分布式训练高效、易用、可扩展。
|
|
|
</div>
|
|
|
|
|
|
<p align="right">(<a href="#top">返回顶端</a>)</p>
|
|
|
|
|
|
## 特点
|
|
|
|
|
|
Colossal-AI 为您提供了一系列并行训练组件。我们的目标是让您的分布式 AI 模型训练像普通的单 GPU 模型一样简单。我们提供的友好工具可以让您在几行代码内快速开始分布式训练。
|
|
|
|
|
|
- 并行化策略
|
|
|
- 数据并行
|
|
|
- 流水线并行
|
|
|
- 1维, [2维](https://arxiv.org/abs/2104.05343), [2.5维](https://arxiv.org/abs/2105.14500), [3维](https://arxiv.org/abs/2105.14450) 张量并行
|
|
|
- [序列并行](https://arxiv.org/abs/2105.13120)
|
|
|
- [零冗余优化器 (ZeRO)](https://arxiv.org/abs/2108.05818)
|
|
|
- 异构内存管理
|
|
|
- [PatrickStar](https://arxiv.org/abs/2108.05818)
|
|
|
- 使用友好
|
|
|
- 基于参数文件的并行化
|
|
|
<p align="right">(<a href="#top">返回顶端</a>)</p>
|
|
|
|
|
|
## 并行样例展示
|
|
|
### ViT
|
|
|
<p align="center">
|
|
|
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/ViT.png" width="450" />
|
|
|
</p>
|
|
|
|
|
|
- 14倍批大小和5倍训练速度（张量并行=64）
|
|
|
|
|
|
### GPT-3
|
|
|
<p align="center">
|
|
|
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/GPT3.png" width=700/>
|
|
|
</p>
|
|
|
|
|
|
- 释放 50% GPU 资源占用, 或 10.7% 加速
|
|
|
|
|
|
### GPT-2
|
|
|
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/GPT2.png" width=800/>
|
|
|
|
|
|
- 降低11倍 GPU 显存占用，或超线性扩展（张量并行）
|
|
|
|
|
|
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/(updated)GPT-2.png" width=800>
|
|
|
|
|
|
- 用相同的硬件条件训练24倍大的模型
|
|
|
- 超3倍的吞吐量
|
|
|
|
|
|
### BERT
|
|
|
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/BERT.png" width=800/>
|
|
|
|
|
|
- 2倍训练速度，或1.5倍序列长度
|
|
|
|
|
|
### PaLM
|
|
|
- [PaLM-colossalai](https://github.com/hpcaitech/PaLM-colossalai): 可扩展的谷歌 Pathways Language Model ([PaLM](https://ai.googleblog.com/2022/04/pathways-language-model-palm-scaling-to.html)) 实现。
|
|
|
|
|
|
请访问我们的[文档和教程](https://www.colossalai.org/)以了解详情。
|
|
|
|
|
|
<p align="right">(<a href="#top">返回顶端</a>)</p>
|
|
|
|
|
|
## 单GPU样例展示
|
|
|
|
|
|
### 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>
|
|
|
|
|
|
- 用相同的硬件条件训练20倍大的模型
|
|
|
|
|
|
### 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>
|
|
|
|
|
|
- 用相同的硬件条件训练34倍大的模型
|
|
|
|
|
|
<p align="right">(<a href="#top">back to top</a>)</p>
|
|
|
|
|
|
## 安装
|
|
|
|
|
|
### 从官方安装
|
|
|
|
|
|
您可以访问我们[下载](/download)页面来安装Colossal-AI，在这个页面上发布的版本都预编译了CUDA扩展。
|
|
|
|
|
|
### 从源安装
|
|
|
|
|
|
> 此文档将与版本库的主分支保持一致。如果您遇到任何问题，欢迎给我们提 issue :)
|
|
|
|
|
|
```shell
|
|
|
git clone https://github.com/hpcaitech/ColossalAI.git
|
|
|
cd ColossalAI
|
|
|
|
|
|
# install dependency
|
|
|
pip install -r requirements/requirements.txt
|
|
|
|
|
|
# install colossalai
|
|
|
pip install .
|
|
|
```
|
|
|
|
|
|
如果您不想安装和启用 CUDA 内核融合（使用融合优化器时强制安装）：
|
|
|
|
|
|
```shell
|
|
|
NO_CUDA_EXT=1 pip install .
|
|
|
```
|
|
|
|
|
|
<p align="right">(<a href="#top">返回顶端</a>)</p>
|
|
|
|
|
|
## 使用 Docker
|
|
|
|
|
|
运行以下命令从我们提供的 docker 文件中建立 docker 镜像。
|
|
|
|
|
|
```bash
|
|
|
cd ColossalAI
|
|
|
docker build -t colossalai ./docker
|
|
|
```
|
|
|
|
|
|
运行以下命令从以交互式启动 docker 镜像.
|
|
|
|
|
|
```bash
|
|
|
docker run -ti --gpus all --rm --ipc=host colossalai bash
|
|
|
```
|
|
|
|
|
|
<p align="right">(<a href="#top">返回顶端</a>)</p>
|
|
|
|
|
|
## 社区
|
|
|
欢迎通过[论坛](https://github.com/hpcaitech/ColossalAI/discussions),
|
|
|
[Slack](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 "qrcode")加入 Colossal-AI 社区，与我们分享你的建议和问题。
|
|
|
|
|
|
|
|
|
## 做出贡献
|
|
|
|
|
|
欢迎为该项目做出贡献，请参阅[贡献指南](./CONTRIBUTING.md)。
|
|
|
|
|
|
真诚感谢所有贡献者！
|
|
|
|
|
|
<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>
|
|
|
|
|
|
*贡献者头像的展示顺序是随机的。*
|
|
|
|
|
|
<p align="right">(<a href="#top">返回顶端</a>)</p>
|
|
|
|
|
|
## 快速预览
|
|
|
|
|
|
### 几行代码开启分布式训练
|
|
|
|
|
|
```python
|
|
|
parallel = dict(
|
|
|
pipeline=2,
|
|
|
tensor=dict(mode='2.5d', depth = 1, size=4)
|
|
|
)
|
|
|
```
|
|
|
|
|
|
### 几行代码开启异构训练
|
|
|
|
|
|
```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">返回顶端</a>)</p>
|
|
|
|
|
|
## 引用我们
|
|
|
|
|
|
```
|
|
|
@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">返回顶端</a>)</p> |
