Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,9 @@

An integrated large-scale model training system with efficient parallelization techniques.

arXiv: [Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training](https://arxiv.org/abs/2110.14883)
Paper: [Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training](https://arxiv.org/abs/2110.14883)

Blog: [Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training](https://www.hpcaitech.com/blog)

## Installation

Expand Down
23 changes: 7 additions & 16 deletions docs/add_your_parallel_zh.md
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
# 添加新的并行技术

为了方便科研人员和工程师们更方便地拓展我们的框架来兼容一些新的大规模分布式训练算法,我们对训练过程中的几个组件进行了解耦,您可以通过继承基类的方式
来实现新的并行技术。
为了方便科研人员和工程师们更方便地拓展我们的系统来兼容一些新的大规模分布式训练算法,我们对训练过程中的几个组件进行了解耦,您可以通过继承基类的方式来实现新的并行技术。

主要的组件如下所示:

Expand All @@ -11,9 +10,7 @@

## 进程组初始化器

并行化一般是通过进程组来进行管理的,同属于一个并行化算法的进程将被分到一个进程组中,如果系统中存在多种不同的并行化技术,那么需要创建多个不同的进程组。
ColossalAI为用户提供了一个全局上下文变量来便捷地管理他们的进程组。如果您希望增加新的进程组,您可以定义一个新的类并且在您的配置文件中进行设置。下方的
代码块中介绍了如果在系统中加入您的新并行技术以及如何进行初始化。
并行化一般是通过进程组来进行管理的,同属于一个并行化算法的进程将被分到一个进程组中,如果系统中存在多种不同的并行化技术,那么需要创建多个不同的进程组。Colossal-AI为用户提供了一个全局上下文变量来便捷地管理他们的进程组。如果您希望增加新的进程组,您可以定义一个新的类并且在您的配置文件中进行设置。下方的代码块介绍了如何在系统中加入您的新颖并行技术以及如何进行初始化。

1. 在`colossalai.context.parallel_mode.ParallelMode`中添加新的并行模式。
```python
Expand All @@ -28,9 +25,7 @@ class ParallelMode(Enum):
NEW_MODE = 'new_mode' # define your mode here
```

2. 创建一个`ProcessGroupInitializer`的子类,您可以参考`colossalai.context.dist_group_initializer`中给出的例子。前六个参数将由`ParallelContext`
决定。如果您需要设置新的参数,您可以用新的参数替换下面例子中的`arg1`与`arg2`。最后,您需要使用`@DIST_GROUP_INITIALIZER.register_module`装饰器
在我们的注册表注册您的初始化器。
2. 创建一个`ProcessGroupInitializer`的子类,您可以参考`colossalai.context.dist_group_initializer`中给出的例子。前六个参数将由`ParallelContext`决定。如果您需要设置新的参数,您可以用新的参数替换下面例子中的`arg1`与`arg2`。最后,您需要使用`@DIST_GROUP_INITIALIZER.register_module`装饰器在我们的注册表中注册您的初始化器。
```python
# sample initializer class
@DIST_GROUP_INITIALIZER.register_module
Expand All @@ -55,14 +50,13 @@ class MyParallelInitializer(ProcessGroupInitializer):
pass
```

在此之后,您可以将您的初始化器插入到当前的mode-to-initialize映射`colossalai.constants.INITIALIZER_MAPPING`中,您也可以通过更改该文件来动态变更名称与
并行模式的映射。
在此之后,您可以将您的初始化器插入到当前的mode-to-initialize映射`colossalai.constants.INITIALIZER_MAPPING`中,您也可以通过更改该文件来动态变更名称与并行模式的映射。

```python
colossalai.constants.INITIALIZER_MAPPING['new_mode'] = 'MyParallelInitializer'
```

3. 在配置文件中设置您的初始化器如果您的初始化器需要参数,您可以自行传入下面的代码可以让`ParallelContext`来创建您的初始化器并初始化您需要的进程组。
3. 在配置文件中设置您的初始化器如果您的初始化器需要参数,您可以自行传入下面的代码可以让`ParallelContext`来创建您的初始化器并初始化您需要的进程组。

```python
parallel = dict(
Expand All @@ -73,9 +67,7 @@ parallel = dict(

## 梯度处理器

梯度处理器的功能是对模型参数的梯度进行all-reduce操作。由于不同的并行技术可能需要不同的all-reduce操作,用户们可以通过继承
`colossalai.engine.gradient_handler.BaseGradientHandler`来执行其个性化操作。目前,ColossalAI使用普通的数据并行梯度处理器,该处理器在所有的数据
并行rank上执行all-reduce操作,且当ColossalAI监测到当前系统使用了数据并行时,该处理器会被自动创建。您可以使用下方代码块中的代码添加您自定义的梯度处理器:
梯度处理器的功能是对模型参数的梯度进行all-reduce操作。由于不同的并行技术可能需要不同的all-reduce操作,用户们可以通过继承`colossalai.engine.gradient_handler.BaseGradientHandler`来执行其个性化操作。目前,Colossal-AI使用普通的数据并行梯度处理器,该处理器在所有的数据并行rank上执行all-reduce操作,且当Colossal-AI监测到当前系统使用了数据并行时,该处理器会被自动创建。您可以使用下方代码块中的代码添加您自定义的梯度处理器:

```python
from colossalai.registry import GRADIENT_HANDLER
Expand All @@ -99,5 +91,4 @@ dist_initializer = [

## 调度器

调度器中指定了在前向传播和后向传播时需要执行哪些操作,ColossalAI提供了支持流水线和不支持流水线的调度器。如果您想要修改前向传播和后向传播的执行方式,您可以
继承`colossalai.engine.BaseSchedule`并实现您想要的操作。您也可以在训练模型之前将您的调度器添加到我们的引擎中来。
调度器中指定了在前向传播和后向传播时需要执行哪些操作,Colossal-AI提供了流水线和非流水线的调度器。如果您想要修改前向传播和后向传播的执行方式,您可以继承`colossalai.engine.BaseSchedule`并实现您想要的操作。您也可以在训练模型之前将您的调度器添加到我们的引擎中来。
15 changes: 6 additions & 9 deletions docs/amp_zh.md
Original file line number Diff line number Diff line change
@@ -1,21 +1,19 @@
# 混合精度训练

ColossalAI可以使用如下三种不同的混合精度训练方式
Colossal-AI可以使用如下三种不同的混合精度训练方式
1. torch.cuda.amp
2. apex.amp
3. 张量并行AMP

前两种混合精度训练方式依赖于[PyTorch](https://pytorch.org/docs/stable/amp.html)的原生实现(1.6或以上版本)以及
[Nvidia Apex](https://github.com/NVIDIA/apex),但这两种方法与张量并行并不兼容,因为在张量并行中我们需要将张量进行切分并保存在不同的设备上,
因此,实现兼容张量并行的混合精度训练需要在不同进程之间不断通信来交流`inf`以及`nan`是否存在于模型参数中,因此我们才用了
[Nvidia Apex](https://github.com/NVIDIA/apex),但这两种方法与张量并行并不兼容,因为在张量并行中我们需要将张量进行切分并保存在不同的设备上。所以,实现兼容张量并行的混合精度训练需要在不同进程之间不断通信来交流`inf`以及`nan`是否存在于模型参数中,因此,我们才用了
[Megatron-LM](https://github.com/NVIDIA/Megatron-LM)的实现方式。

您可以简单地将配置文件中的`fp16`字段设置为True来使用混合精度训练。目前,PyTorch与Apex的amp不能保证与张量和流水线并行兼容,因此,我们推荐您使用
最后一种混合精度训练方式。
您可以简单地将配置文件中的`fp16`字段设置为True来使用混合精度训练。目前,PyTorch与Apex的amp不能保证与张量并行、流水线并行兼容,因此,我们推荐您使用最后一种混合精度训练方式。

## PyTorch AMP

PyTorch在1.6及以上版本中提供了混合精度训练,其可以在保持一些操作的精度为`fp32`的同时,将数据转换成`fp16`格式,您可以在配置文件中配置使用。
PyTorch在1.6及以上版本中提供了混合精度训练,它可以在保持一些操作的精度为`fp32`的同时,将数据转换成`fp16`格式,您可以在配置文件中配置使用。

```python
from colossalai.engine import AMP_TYPE
Expand All @@ -33,8 +31,7 @@ fp16=dict(

## Apex AMP

我们使用了[Apex](https://nvidia.github.io/apex/)中的混合精度训练,因为该模式提供了细粒度的混合精度控制,例如,`O2`级(第二级优化器)将会保持
批标准化在`fp32`上进行。下面的代码块展示了使用Apex AMP的配置文件。
我们使用了[Apex](https://nvidia.github.io/apex/)中的混合精度训练,因为该模式提供了细粒度的混合精度控制,例如,`O2`级(第二级优化器)将会保持批标准化在`fp32`上进行。下面的代码块展示了使用Apex AMP的配置文件。

```python
from colossalai.engine import AMP_TYPE
Expand All @@ -59,7 +56,7 @@ fp16 = dict(

## 张量并行AMP

我们借鉴了Megatron-LM的混合精度训练实现,该实现方式与张量并行与流水线并行相兼容。下面的代码块展示了使用张量并行AMP的配置文件。
我们借鉴了Megatron-LM的混合精度训练实现,该实现方式与张量并行、流水线并行相兼容。下面的代码块展示了使用张量并行AMP的配置文件。

```python
from colossalai.engine import AMP_TYPE
Expand Down
2 changes: 1 addition & 1 deletion docs/config_zh.md
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
# 配置文件

下方代码块中的示例展示了如何在CIFAR10数据集上使用ColossalAI训练ViT模型
下方代码块中的示例展示了如何在CIFAR10数据集上使用Colossal-AI训练ViT模型

```python
# build train_dataset and train_dataloader from this dictionary
Expand Down
2 changes: 1 addition & 1 deletion docs/model.md
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
# Define your own parallel model

Let's say that you have a huge MLP model with billions of parameters and its extremely large hidden layer size makes it
impossible to fit into a single GPU. Don't worry, ColossalAI is here to help you sort things out. With the help of ColossalAI,
impossible to fit into a single GPU directly. Don't worry, ColossalAI is here to help you sort things out. With the help of ColossalAI,
you can write your model in the familiar way in which you used to write models for a single GPU, while ColossalAI automatically
splits your model weights and fit them perfectly into a set of GPUs. We give a simple example showing how to write a simple
2D parallel model in the ColossalAI context.
Expand Down
6 changes: 2 additions & 4 deletions docs/model_zh.md
Original file line number Diff line number Diff line change
@@ -1,8 +1,6 @@
# 定义符合您需求的并行模型

如果您在训练一个拥有数亿级参数的巨大MLP模型,那么该模型一定无法在单个GPU上进行训练,不用担心,ColossalAI可以帮您解决这一问题。您仍旧可以像写单GPU模型
那样来写您的模型,ColossalAI会按照您的并行设置自动将模型参数进行切割,并将它们均匀地存入一组GPU中。下面是一个简单的例子,来向您展示如何在ColossalAI
环境下写一个2D张量并行的模型。
如果您在训练一个拥有数亿级参数的巨大MLP模型,那么该模型一定无法在单个GPU上直接进行训练,不用担心,Colossal-AI可以帮您解决这一问题。您仍旧可以像写单GPU模型那样来写您的模型,Colossal-AI会按照您的并行设置自动将模型参数进行切割,并将它们均匀地存入一组GPU中。下面是一个简单的例子,来向您展示如何在Colossal-AI环境下写一个2D张量并行的模型。

## 简单的2D张量并行模型

Expand All @@ -25,4 +23,4 @@ class MLP_2D(nn.Module):

## 使用事先定义好的模型

为了您使用的方便,我们事先在我们的Model Zoo中定义好了一些现在流行的模型,比如*BERT*、*VIT*以及*MLP-Mixer*,您可以根据您的需求来自定义这些模型的规模。
为了您使用的方便,我们事先在我们的Model Zoo中定义好了一些现在流行的模型,比如*BERT*、*VIT*以及*MLP-Mixer*,您可以根据您的需求来自定义这些模型的规模。
Loading