mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-22 09:59:38 +00:00
[doc] clean up outdated docs (#4765)
* [doc] clean up outdated docs * [doc] fix linking * [doc] fix linking
This commit is contained in:
Hongxin Liu
GitHub
@@ -1,113 +0,0 @@
|
||||
# 添加你自己的并行模式
|
||||
|
||||
作者: Shenggui Li, Yongbin Li
|
||||
|
||||
**前置教程**
|
||||
- [定义配置文件](../basics/define_your_config.md)
|
||||
- [并行配置](../basics/configure_parallelization.md)
|
||||
|
||||
## 引言
|
||||
|
||||
为了使研究人员和工程师能够以更少的努力将我们的系统扩展到其他新颖的大规模分布式训练算法，我们已经将训练生命周期中的各种组件解耦。你可以通过简单地继承基类来实现你自己的并行模式。
|
||||
|
||||
主要组件有:
|
||||
|
||||
1. `ProcessGroupInitializer`
|
||||
2. `GradientHandler`
|
||||
3. `Schedule`
|
||||
|
||||
**目前这需要对源代码进行一些改动，因此我们建议你用`-e`标志从源代码安装。`-e`标志使得安装是可编辑的，因此，你的代码变化将反映在你的Python运行时中。我们将在这方面努力，以避免在未来的版本中改变源代码。**
|
||||
|
||||
|
||||
## 进程组初始化器
|
||||
|
||||
并行通常由进程组来管理，参与相同并行算法的进程被置于同一进程组。对于不同的并行算法，需要创建不同的进程组。
|
||||
Colossal-AI 为用户提供了一个全局 context，使他们能够轻松地管理进程组。如果你想添加新的进程组，你可以很容易地定义一个新的类并在你的配置文件中设置它。为了定义你自己的进程组创建方式，你可以按照下面的步骤来创建一个新的分布式初始化。
|
||||
|
||||
1. 在 `colossalai.legacy.context.parallel_mode.ParallelMode` 中添加你自己的并行模式。
|
||||
```python
|
||||
class ParallelMode(Enum):
|
||||
GLOBAL = 'global'
|
||||
DATA = 'data'
|
||||
PIPELINE = 'pipe'
|
||||
...
|
||||
|
||||
NEW_MODE = 'new_mode' # define your mode here
|
||||
```
|
||||
|
||||
2. 创建一个 `ProcessGroupInitializer`。 你可以参考 `colossalai.context.dist_group_initializer` 中给出的例子，前六个参数是固定的。
|
||||
`ParallelContext` 将为你传入这些参数。如果你需要设置其他参数，可以像下面的例子中的 `arg1, arg2` 一样，在后面添加它。
|
||||
最后，通过添加装饰器 `@DIST_GROUP_INITIALIZER.register_module` 将你的初始化程序注册到注册表。
|
||||
```python
|
||||
# sample initializer class
|
||||
@DIST_GROUP_INITIALIZER.register_module
|
||||
class MyParallelInitializer(ProcessGroupInitializer):
|
||||
|
||||
def __init__(self,
|
||||
rank: int,
|
||||
world_size: int,
|
||||
config: Config,
|
||||
data_parallel_size: int,
|
||||
pipeline_parallel_size: int,
|
||||
tensor_parallel_size: int,
|
||||
arg1,
|
||||
arg2):
|
||||
super().__init__(rank, world_size, config)
|
||||
self.arg1 = arg1
|
||||
self.arg2 = arg2
|
||||
# ... your variable init
|
||||
|
||||
def init_parallel_groups(self):
|
||||
# initialize your process groups
|
||||
pass
|
||||
|
||||
```
|
||||
然后，你可以将你的新初始化器插入到 `colossalai.constants.INITIALIZER_MAPPING` 当前的模式与初始化映射中。你可以修改该文件或动态插入新的键值对。
|
||||
|
||||
```python
|
||||
colossalai.constants.INITIALIZER_MAPPING['new_mode'] = 'MyParallelInitializer'
|
||||
```
|
||||
|
||||
3. 在你的配置文件中设置你的初始化器。你可以传入你的自定义参数。这允许
|
||||
`ParallelContext` 创建你的初始化器并初始化你期望的进程组。
|
||||
|
||||
```python
|
||||
parallel = dict(
|
||||
pipeline=dict(size=1),
|
||||
tensor=dict(size=x, mode='new_mode') # this is where you enable your new parallel mode
|
||||
)
|
||||
```
|
||||
|
||||
## 梯度 Handler
|
||||
|
||||
梯度 handler 是对参数的梯度执行 all-reduce 操作的对象。由于不同的 all-reduce 策略或许在不同的并行中被执行，用户可以继承
|
||||
`colossalai.legacy.engine.gradient_handler.BaseGradientHandler` 来实现其策略。目前，Colossal-AI 使用普通的数据并行梯度 handler 在数据并行的 rank 间 all-reduce 梯度。
|
||||
如果数据并行被检测到，梯度 handler 会被自动添加进 engine。
|
||||
|
||||
你可以添加你自己的梯度 handler，如下所示：
|
||||
|
||||
```python
|
||||
from colossalai.legacy.registry import GRADIENT_HANDLER
|
||||
from colossalai.legacy.engine import BaseGradientHandler
|
||||
|
||||
@GRADIENT_HANDLER.register_module
|
||||
class YourGradientHandler(BaseGradientHandler):
|
||||
|
||||
def handle_gradient(self):
|
||||
do_something()
|
||||
|
||||
```
|
||||
|
||||
之后，你可以在配置文件中指定你要使用的梯度 handler。
|
||||
|
||||
```python
|
||||
gradient_handlers = [
|
||||
dict(type='YourGradientHandler'),
|
||||
]
|
||||
```
|
||||
|
||||
## Schedule
|
||||
|
||||
Schedule 包含了如何执行前向和后向计算。目前， Colossal-AI 提供了流水和非流水的 schedule。
|
||||
如果你想修改前向和后向计算的执行方式，你可以继承 `colossalai.legacy.engine.schedule.BaseSchedule` 并实现 `forward_back_step` 函数。
|
||||
<!-- doc-test-command: echo -->
|
||||
@@ -1,31 +0,0 @@
|
||||
# 定义你自己的并行模型
|
||||
|
||||
作者: Zhengda Bian, Yongbin Li
|
||||
|
||||
> ⚠️ 我们正在编写此文档以使其更加详细。 我们将介绍不同并行的机制以及如何使用它们来编写模型。
|
||||
|
||||
假设您有一个具有数十亿参数的巨大 MLP 模型，其极大的隐藏层大小使其无法直接被单个 GPU 容纳。别担心，Colossal-AI 可以帮你解决这个问题。
|
||||
在 Colossal-AI 的帮助下，您可以用所熟悉的为单个 GPU 编写模型的方式编写大模型，而 Colossal-AI 会自动拆分您的模型权重，并将它们完美地分配到一组 GPU 中。我们给出一个简单的示例，展示如何在 Colossal-AI 中编写简单的 2D 并行模型。
|
||||
|
||||
## 写一个简单的2D并行模型
|
||||
|
||||
```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
|
||||
```
|
||||
|
||||
## 使用预定义的模型
|
||||
|
||||
为了方便您的使用，我们在 Colossal-AI 的 Model Zoo 中提供一些流行的模型，如*BERT*, *ViT*, *MoE* 和 *GPT*，请自由地将它们定制为不同的尺寸，以满足您的特殊需求。
|
||||
@@ -1,179 +0,0 @@
|
||||
# 使用ColoTensor让串行程序像Megatron-LM一样并行
|
||||
|
||||
Author: [Haichen Huang](https://github.com/1SAA) and [Jiarui Fang](https://github.com/feifeibear)
|
||||
|
||||
**Prerequisite:**
|
||||
- [ColoTensor Concepts](../basics/colotensor_concept.md)
|
||||
|
||||
## 介绍
|
||||
|
||||
在新版本中，我们引入了ColoTensor。ColoTensor为用户使用并行训练提供了极大的便利，使得用户可以在原本的串行代码上，通过较小的修改将训练改为并行。在本教程中，我们将说明如何修改训练模型以自动使代码采取像 Megatron-LM 一样的方式并行训练。我们以 HuggingFace 提供的 GPT-2 模型为例，并提供一种方式让你可以在单个GPU上预训练GPT-2模型。
|
||||
|
||||
Megatron-LM 提供了一个具有影响力的并行化范式，这个范式主要应用于Transformer大模型的训练。然而，为了大规模训练 Transformer 语言大模型，用户必须使用Megatron-LM提供的特殊模块来构建他们的模型。这给用户带来了一些困难的工作，例如从预先训练的模型中加载权重，或是构建自己的并行训练模型。为了减轻用户的麻烦，我们提供 ColoTensor 类，以完成自动启用张量模型并行。
|
||||
|
||||
## 定义模型和损失函数
|
||||
|
||||
首先，我们直接调用 HuggingFace 库中的 GPTModel 和 GPTLoss。
|
||||
|
||||
```python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from transformers import GPT2Config, GPT2LMHeadModel
|
||||
|
||||
class GPTLMModel(nn.Module):
|
||||
def __init__(self, hidden_size=768, num_layers=12, num_attention_heads=12, max_seq_len=1024, vocab_size=50257, checkpoint=False):
|
||||
super().__init__()
|
||||
self.checkpoint = checkpoint
|
||||
self.model = GPT2LMHeadModel(GPT2Config(n_embd=hidden_size, n_layer=num_layers,
|
||||
n_head=num_attention_heads, n_positions=max_seq_len, n_ctx=max_seq_len, vocab_size=vocab_size))
|
||||
if checkpoint:
|
||||
self.model.gradient_checkpointing_enable()
|
||||
|
||||
def forward(self, input_ids, attention_mask):
|
||||
# Only return lm_logits
|
||||
return self.model(input_ids=input_ids, attention_mask=attention_mask, use_cache=not self.checkpoint)[0]
|
||||
|
||||
|
||||
class GPTLMLoss(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.loss_fn = nn.CrossEntropyLoss()
|
||||
|
||||
def forward(self, logits, labels):
|
||||
shift_logits = logits[..., :-1, :].contiguous()
|
||||
shift_labels = labels[..., 1:].contiguous()
|
||||
# Flatten the tokens
|
||||
return self.loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
|
||||
```
|
||||
|
||||
## 对GPT-2的简短回顾
|
||||
|
||||
现在，我们回顾一下 GPT-2 模型的结构。每个 GPT-2 模型都可以表示为一个 DAG。如下图所示，每个圆圈代表一个算子，每个方块代表一个权重。每个箭头表示输入数据的流向，而箭头旁边的符号表示输入数据的形状。
|
||||
|
||||
然后，让我们深入了解一下这个 GPT-2 模型。它由三部分组成，分别是**嵌入模块**、**转换器层**和**分类头**。
|
||||
|
||||
嵌入模块包含两个权重，符号嵌入权重和位置嵌入权重。在嵌入模块的前向操作之后，原始输入数据的所有序列中的每个单词都会被嵌入到隐藏状态。
|
||||
|
||||
<figure style={{textAlign: "center"}}>
|
||||
<img src="https://s2.loli.net/2022/08/17/omfkIEN6ui5jcL3.png"/>
|
||||
<figcaption>嵌入模块</figcaption>
|
||||
</figure>
|
||||
|
||||
每个转换器层包含两个块。自注意操作在第一个块中调用，同时一个双层感知器位于第二个块中。
|
||||
|
||||
<figure style={{textAlign: "center"}}>
|
||||
<img src="https://s2.loli.net/2022/08/17/LAVzDlpRcj4dYeb.png"/>
|
||||
<figcaption>转换器层</figcaption>
|
||||
</figure>
|
||||
|
||||
最后，分类头只是一个不加偏差的线性模块，里面只有一个线性权重。
|
||||
|
||||
## 应用ColoTensor
|
||||
|
||||
两个步骤使您的串行代码采取 Megatron-LM 张量并行风格。
|
||||
1. 在ColoInitContext的上下文中初始化模型。
|
||||
2. 为每个参数设置 ColoTensorSpec。
|
||||
|
||||
### 使用 ColoInitContext 初始化
|
||||
|
||||
我们应该在 ColoInitContext 中构建模型。在该种上下文中，任何初始化的参数都将转换为 ColoParameter 并自动移动到相应的设备上。
|
||||
|
||||
```python
|
||||
from colossalai.utils.model.colo_init_context import ColoInitContext
|
||||
|
||||
with ColoInitContext(device=torch.device('cpu')):
|
||||
model = GPTLMModel()
|
||||
```
|
||||
|
||||
### 为每个参数设置 ColoTensorSpec
|
||||
|
||||
模型创建完成后，我们通过ProcessGroup建立分布式环境。这里，我们将张量并行度指定为所有GPU的数量，即数据并行度为一。
|
||||
|
||||
```python
|
||||
import torch.distributed as dist
|
||||
from colossalai.tensor import ProcessGroup
|
||||
|
||||
pg = ProcessGroup(tp_degree=dist.get_world_size())
|
||||
```
|
||||
|
||||
现在，我们需要一些辅助函数为下一步做准备。我们定义了两个函数来切分参数。Megatron-LM张量并行需要沿参数的第一维或最后一维切分参数张量。
|
||||
|
||||
```python
|
||||
from colossalai.tensor import ShardSpec, ComputeSpec, ComputePattern, ColoParameter, ProcessGroup
|
||||
|
||||
def split_param_single_dim_tp1d(dim: int, param: ColoParameter, pg: ProcessGroup):
|
||||
spec = (ShardSpec([dim], [pg.tp_world_size()]), ComputeSpec(ComputePattern.TP1D))
|
||||
if param.process_group.tp_world_size() == 1:
|
||||
param.set_process_group(pg)
|
||||
param.set_tensor_spec(*spec)
|
||||
|
||||
|
||||
def split_param_row_tp1d(param: ColoParameter, pg: ProcessGroup):
|
||||
split_param_single_dim_tp1d(0, param, pg)
|
||||
|
||||
|
||||
def split_param_col_tp1d(param: ColoParameter, pg: ProcessGroup):
|
||||
split_param_single_dim_tp1d(-1, param, pg)
|
||||
```
|
||||
|
||||
然后我们使模型采用张量并行。根据 Megatron 中使用的张量并行，应该沿着张量的最后一个维度进行切片，包括符号嵌入的权重，位置嵌入的权重，自注意力块中的所有线性权重和偏差，以及每个双层感知器中的第一个线性权重和偏差。且需要沿第一个维度切分双层感知器中的第二个线性权重。
|
||||
|
||||
```python
|
||||
for mn, module in model.named_modules():
|
||||
for pn, param in module.named_parameters(recurse=False):
|
||||
# set process group for all parameters
|
||||
param.set_process_group(pg)
|
||||
|
||||
if 'mlp.c_fc' in mn:
|
||||
if 'weight' in pn or 'bias' in pn:
|
||||
split_param_col_tp1d(param, pg) # column slice
|
||||
# keep the shape of the output from c_fc
|
||||
param.compute_spec.set_output_replicate(False)
|
||||
elif 'mlp.c_proj' in mn:
|
||||
if 'weight' in pn:
|
||||
split_param_row_tp1d(param, pg) # row slice
|
||||
elif 'wte' in mn or 'wpe' in mn:
|
||||
split_param_col_tp1d(param, pg) # column slice
|
||||
elif 'c_attn' in mn or 'c_proj' in mn:
|
||||
split_param_col_tp1d(param, pg) # column slice
|
||||
```
|
||||
|
||||
修改后的模型如下图所示。
|
||||
|
||||
嵌入模块:
|
||||
|
||||
<figure style={{textAlign: "center"}}>
|
||||
<img src="https://s2.loli.net/2022/08/17/Yu2xzXEabHV7pwe.png"/>
|
||||
<figcaption>修改后的嵌入模块</figcaption>
|
||||
</figure>
|
||||
|
||||
转换器层:
|
||||
|
||||
<figure style={{textAlign: "center"}}>
|
||||
<img src="https://s2.loli.net/2022/08/17/4HWsA2xz51IhPFO.png"/>
|
||||
<figcaption>修改后的转换器层</figcaption>
|
||||
</figure>
|
||||
|
||||
一旦用户指定了每个参数的在并行中的分布模式，ColoTensor 就能够推断出所有算子的计算模式，包括矩阵乘法、线性函数、torch.nn.functional 中的其他逐元素函数，以及其他的一些常用函数。这样，用户可以像往常一样训练他们的模型。
|
||||
|
||||
在我们最新示例中还定义了一个Gemini + ZeRO DDP 的模型从而减小开销，提升效率。这一部分的详细内容可以参考[ZeRO](../features/zero_with_chunk.md)，你可以将这两部分内容结合起来看从而理解我们整个训练流程：
|
||||
|
||||
```python
|
||||
def gemini_zero_dpp(model: torch.nn.Module, pg: ProcessGroup, placement_policy: str = "auto"):
|
||||
from colossalai.zero import GeminiDDP
|
||||
model = GeminiDDP(model,
|
||||
device=get_current_device(),
|
||||
placement_policy=placement_policy,
|
||||
pin_memory=True,
|
||||
search_range_m=32)
|
||||
return model
|
||||
```
|
||||
|
||||
## 在单个GPU上预训练GPT-2
|
||||
|
||||
我们做的上述优化让我们可以在单GPU上训练GPT-2模型，只需要将`run.sh`中设置参数`GPUNUM`=1，再运行文件时就可以在单个GPU上完成模型的训练。
|
||||
|
||||
GPT-2 示例在[Train GPT with Colossal-AI](https://github.com/hpcaitech/ColossalAI/tree/main/examples/language/gpt). 获得。
|
||||
|
||||
|
||||
<!-- doc-test-command: torchrun --standalone --nproc_per_node=1 parallelize_your_training_like_Megatron.py -->
|
||||
Reference in New Issue
Block a user