mirror of https://github.com/hpcaitech/ColossalAI
156 lines
6.4 KiB
Markdown
156 lines
6.4 KiB
Markdown
# 序列并行
|
||
|
||
作者: Mingyan Jiang
|
||
|
||
**前置教程**
|
||
- [并行技术](../concepts/paradigms_of_parallelism.md)
|
||
- [Booster API](../basics/booster_api.md)
|
||
- [Shardformer](../features/shardformer.md)
|
||
- [Booster 插件](../basics/booster_plugins.md)
|
||
|
||
**示例代码**
|
||
- [使用序列并行策略](https://github.com/hpcaitech/ColossalAI/blob/main/examples/language/llama/benchmark.py)
|
||
|
||
**相关论文**
|
||
[Reducing Activation Recomputation in Large Transformer Models](https://arxiv.org/pdf/2205.05198)
|
||
[DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models](https://arxiv.org/abs/2309.14509)
|
||
[Ring Attention with Blockwise Transformers for Near-Infinite Context](https://arxiv.org/pdf/2310.01889)
|
||
|
||
## 快速预览
|
||
|
||
在本教程中，你将学习如何使用序列并行。在 Colossal-AI 中, 我们实现了包括TP+SP， DeepSpeed-Ulysses， ring attention等多种序列并行. 我们下面将介绍如何使用这几种序列并行。
|
||
|
||
## 目录
|
||
|
||
在本教程中，我们将介绍三种序列并行的使用:
|
||
|
||
1. 使用TP+SP；
|
||
2. 使用DeepSpeed-Ulysses；
|
||
3. 使用ring attention
|
||
|
||
|
||
## Colossal-AI中的实现
|
||
|
||
在 Colossal-AI 中，shardformer实现了序列并行，并通过`HybridParallelPlugin`和`MoeHybridParallelPlugin`接口可进行调用。相关plugin的介绍请参考plugin的[使用文档](../basics/booster_plugins.md)。
|
||
|
||
### 使用`HybridParallelPlugin`的序列并行
|
||
`HybridParallelPlugin`的序列支持了TP+SP， DeepSpeed-Ulysses， ring attention三种实现，相关序列并行的结束可参考[并行技术介绍文档](../concepts/paradigms_of_parallelism.md)，`HybridParallelPlugin`中的序列并行[例子](https://github.com/hpcaitech/ColossalAI/blob/main/examples/language/llama/benchmark.py)
|
||
|
||
#### 定义模型相关组件
|
||
|
||
```python
|
||
from tqdm import tqdm
|
||
from transformers import AutoModelForCausalLM
|
||
from transformers.models.llama.configuration_llama import LlamaConfig
|
||
from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
|
||
import torch.distributed as dist
|
||
from colossalai.booster import Booster
|
||
config = LlamaConfig(max_position_embeddings=4096)
|
||
from colossalai.booster.plugin import HybridParallelPlugin
|
||
|
||
# 定义数据集
|
||
class RandomDataset(Dataset):
|
||
def __init__(self, num_samples: int = 1000, max_length: int = 2048, vocab_size: int = 32000):
|
||
self.num_samples = num_samples
|
||
self.max_length = max_length
|
||
self.input_ids = torch.randint(
|
||
0, vocab_size, (num_samples, max_length), device=get_accelerator().get_current_device()
|
||
)
|
||
self.attention_mask = torch.ones_like(self.input_ids)
|
||
|
||
def __len__(self):
|
||
return self.num_samples
|
||
|
||
def __getitem__(self, idx):
|
||
return {
|
||
"input_ids": self.input_ids[idx],
|
||
"attention_mask": self.attention_mask[idx],
|
||
"labels": self.input_ids[idx],
|
||
}
|
||
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("-b", "--batch_size", type=int, default=2, help="Batch size")
|
||
parser.add_argument("-s", "--num_steps", type=int, default=5, help="Number of steps to run")
|
||
parser.add_argument("-l", "--max_length", type=int, default=4096, help="Max sequence length")
|
||
parser.add_argument("--tp", type=int, default=1, help="Tensor parallel size")
|
||
parser.add_argument("--sp", type=int, default=1, help="Sequence parallel size")
|
||
args = parser.parse_args()
|
||
|
||
model = AutoModelForCausalLM.from_config(
|
||
config,
|
||
trust_remote_code=True,
|
||
attn_implementation="flash_attention_2",
|
||
torch_dtype=torch.bfloat16,
|
||
)
|
||
optimizer = HybridAdam(model.parameters())
|
||
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.1)
|
||
# usually, num_samples=args.batch_size * args.num_steps * dp_size
|
||
dataset = RandomDataset(
|
||
num_samples=10000, max_length=args.max_length, vocab_size=config.vocab_size
|
||
)
|
||
```
|
||
### 使用TP+SP
|
||
定义plugin,使用该序列并行，`sp_size`会被设置为`tp_size`一致，且tp group 与sp group是重叠的。
|
||
```python
|
||
plugin = HybridParallelPlugin(
|
||
tp_size=4,
|
||
sp_size=1,
|
||
enable_all_optimization=True,
|
||
enable_sequence_parallelism=True,
|
||
sequence_parallelism_mode="split_gather",
|
||
)
|
||
```
|
||
|
||
#### 使用DeepSpeed-Ulysses
|
||
定义plugin， 在DeepSpeed-Ulysses的序列并行种，tp group与sp group 是正交的，
|
||
```python
|
||
plugin = HybridParallelPlugin(
|
||
tp_size=2,
|
||
sp_size=2,
|
||
enable_all_optimization=True,
|
||
enable_sequence_parallelism=True,
|
||
sequence_parallelism_mode="all_to_all",
|
||
)
|
||
```
|
||
|
||
#### 使用ring attention
|
||
定义plugin， 在ring attention的序列并行种，tp group与sp group 是正交的，sp_size必须传入准确的并行大小。
|
||
```python
|
||
plugin = HybridParallelPlugin(
|
||
tp_size=2,
|
||
sp_size=2,
|
||
enable_all_optimization=True,
|
||
enable_sequence_parallelism=True,
|
||
sequence_parallelism_mode="ring_attn",
|
||
)
|
||
```
|
||
#### 使用booster
|
||
```python
|
||
booster = Booster(plugin=plugin)
|
||
dataloader = plugin.prepare_dataloader(dataset, batch_size=args.batch_size, shuffle=True, drop_last=True, seed=42)
|
||
model, optimizer, _, dataloader, _ = booster.boost(model, optimizer, dataloader=dataloader)
|
||
```
|
||
|
||
#### 训练模型
|
||
```python
|
||
for step, batch in enumerate(tqdm(dataloader, desc="Step", disable=not dist.get_rank()==0)):
|
||
outputs = model(**batch)
|
||
loss = outputs[0]
|
||
del outputs # free memory
|
||
|
||
if dist.get_rank() == dist.get_world_size() - 1:
|
||
print(f"Step {step} loss: {loss}")
|
||
booster.backward(loss, optimizer)
|
||
optimizer.step()
|
||
optimizer.zero_grad()
|
||
```
|
||
### 使用`MoeHybridParallelPlugin`的序列并行
|
||
`MoeHybridParallelPlugin`中的序列并行暂时只支持DeepSpeed-Ulysses类型,使用方法与`HybridParallelPlugin`类似，具体可参考[例子](https://github.com/hpcaitech/ColossalAI/blob/main/examples/language/deepseek/benchmark.py)
|
||
|
||
|
||
|
||
### 结论
|
||
在上述序列并行方法中，ring attention对head number没有要求，可训练超长文本，但是由于细分了计算，计算性能会有所下降。TP+SP， DeepSpeed-Ulysses对于head number有要求，需要可被sp group size 整除。这些序列并行都可与其他高性能注意力兼容，如flash attention。sp可与Gemini一起使用训练超大规模模型，也可以与TP，PP，DP等组成4D并行。
|
||
|
||
<!-- doc-test-command: torchrun --standalone --nproc_per_node=4 sequence_parallelism.py -->
|