Chapter 09

DeepSpeed 与 sequence parallel 集成

📌 commit 964f70e多卡 / 多机 / 长上下文配方

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 GBbaseline
ZeRO-2~16 GB
ZeRO-3~12 GB稍慢
ZeRO-2 + SP=4~14 GBSP overhead 5%
seqlen=32K + SP=4~22 GB

9.5Variable-length Attention

use_varlen_attn=Truepack_to_max_length=True 时,XTuner 用 flash-attn 的 cu_seqlens 让 attention 在 packed 序列内部不串。这是 packing 必须开的开关。

9.6这章你需要带走的

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=8sft.py:16-21 / dpo.py:12-17 自动调用上述工具。可与 DeepSpeed ZeRO + DP 三层叠加。

9.8 这章你需要带走的