Chapter 09
DeepSpeed 与 sequence parallel 集成
XTuner 多卡靠 DeepSpeed,长上下文用 Ulysses 风格 sequence parallel。本章给推荐配方。
9.1DeepSpeed 内置 preset
XTuner 自带四个 DS 配置:
| preset | 对应 |
|---|---|
| deepspeed_zero1 | 切 optim |
| deepspeed_zero2 | + grad |
| deepspeed_zero2_offload | + CPU offload |
| deepspeed_zero3 | + param |
| deepspeed_zero3_offload | + CPU offload param |
启动加 --deepspeed deepspeed_zero2。
9.2多卡命令
# 单机 8 卡
NPROC_PER_NODE=8 xtuner train ./cfg.py --deepspeed deepspeed_zero2
# 多机
NPROC_PER_NODE=8 \
NNODES=4 \
NODE_RANK=$RANK \
PORT=12345 \
ADDR=10.0.0.5 \
xtuner train ./cfg.py --deepspeed deepspeed_zero2
9.3Sequence Parallel(长上下文专用)
XTuner 集成 Ulysses 风格 SP(DeepSpeed-Ulysses),把 attention 沿 head 切,沿 seq 做 all-to-all。能把 8K 上下文扩到 128K+。
# config 里
from xtuner.parallel.sequence import SequenceParallelSampler
sequence_parallel_size = 4
model = dict(
type=SupervisedFinetune,
use_varlen_attn=True,
...
)
train_dataloader = dict(
...,
sampler=dict(type=SequenceParallelSampler, shuffle=True),
)
启动加:
NPROC_PER_NODE=8 xtuner train ./cfg_with_sp.py \
--deepspeed deepspeed_zero2 \
--seq-parallel-size 4
对 8 卡训 32K 长度的 Llama-3-8B,sp=4 + tp=1 + dp=2 是常见组合。
9.4显存预算
InternLM2-7B、8×A100-80GB、QLoRA、seqlen=8K:
| 配置 | 显存/卡 | 速度 |
|---|---|---|
| ZeRO-1 | ~20 GB | baseline |
| ZeRO-2 | ~16 GB | 同 |
| ZeRO-3 | ~12 GB | 稍慢 |
| ZeRO-2 + SP=4 | ~14 GB | SP overhead 5% |
| seqlen=32K + SP=4 | ~22 GB | — |
9.5Variable-length Attention
use_varlen_attn=True 配 pack_to_max_length=True 时,XTuner 用 flash-attn 的 cu_seqlens 让 attention 在 packed 序列内部不串。这是 packing 必须开的开关。
9.6这章你需要带走的
- XTuner 多卡用 DeepSpeed preset:zero1/2/3 / 加 offload;
- 长上下文用 sequence parallel(Ulysses 风格);
- 启动加
NPROC_PER_NODE、多机加NNODES/NODE_RANK/ADDR/PORT; - 开 packing 时
use_varlen_attn=True必备。
9.7 sequence parallel 源码定位
XTuner 的 SP 走 Ulysses(all-to-all),实现都在 xtuner/parallel/sequence/:
| 文件 | 关键函数 | 用途 |
|---|---|---|
setup_distributed.py | init_sequence_parallel()get_sequence_parallel_group() | 建 2D(DP × SP)process group |
comm.py | all_to_all()split_for_sequence_parallel()gather_forward_split_backward() | Ulysses 通信原语 |
attention.py | sequence_parallel_wrapper() | 把 attention 包成 SP 版 |
sampler.py | SequenceParallelSampler | 给 dataloader 用,按 SP 组切 batch |
reduce_loss.py | reduce_sequence_parallel_loss() | SP 组内 loss 求平均 |
启用:config 写 sequence_parallel_size=8;sft.py:16-21 / dpo.py:12-17 自动调用上述工具。可与 DeepSpeed ZeRO + DP 三层叠加。
9.8 这章你需要带走的
- SP 工具集中在
xtuner/parallel/sequence/,6 文件 ~500 行; - SequenceParallelSampler 是切 batch 的关键;
- DeepSpeed strategy 在
xtuner/engine/_strategy/deepspeed.py。