HybridParallelPlugin：一份配置开多种并行

HybridParallel 是 ColossalAI 在 ZeRO-3 之外的"另一极"—— 显存够时它能给到接近 Megatron 的吞吐，且只用一行 plugin 配置。本章拆解它的 mesh、API 与典型配置。

6.15 种并行一锅炖

HybridParallel 的设计是：先把 world_size 拆成 DP × TP × PP 三维 mesh，再在 DP 维度上叠 ZeRO，TP 维度上选要不要 SP。

6.2Mesh 是什么

$$\text{world\_size} = \text{dp\_size} \times \text{tp\_size} \times \text{pp\_size}$$

以 16 卡为例，HybridParallelPlugin(tp_size=2, pp_size=2) 隐式 dp_size = 16/(2×2) = 4。 ColossalAI 会建 4 个 DP group、2 个 TP group、2 个 PP group。每张 GPU 同时属于这三个 group，分别走对应的通信。

6.3关键参数详解

看 colossalai/booster/plugin/hybrid_parallel_plugin.py，重要参数：

plugin = HybridParallelPlugin(
    # —— 并行维度 ——
    tp_size=2,
    pp_size=2,
    sp_size=1,                          # 通常 sp_size 沿用 tp_size
    ep_size=1,                          # MoE 才用
    # —— 序列并行类型 ——
    sequence_parallelism_mode="split_gather",  # split_gather / ring_attn / all_to_all
    enable_sequence_parallelism=False,
    # —— ZeRO ——
    zero_stage=1,                       # 0 / 1 / 2，不支持 3（用 Gemini 代替）
    cpu_offload=False,
    # —— 精度 ——
    precision="bf16",
    # —— Pipeline schedule ——
    num_microbatches=None,              # micro-batch 数
    microbatch_size=None,               # 或显式 size
    pp_style="1f1b",                    # 1f1b / interleaved / zbv
    num_model_chunks=1,                 # interleaved 用
    # —— ShardFormer 优化 ——
    enable_fused_normalization=True,
    enable_flash_attention=True,
    enable_jit_fused=False,
    # —— Gradient ——
    enable_gradient_accumulation=False,
    initial_scale=2**16,
    # —— Checkpointing ——
    gradient_checkpoint_config=None,
)

pp_style 三种

其中 $p$ = PP stage 数，$m$ = micro-batch 数，$v$ = num_model_chunks（仅 interleaved）。 新项目建议先用 1f1b 跑通，再试 zbv（实现位置 colossalai/pipeline/schedule/zero_bubble_pp.py）。

6.4典型配置三连

① 7B / 单机 8 卡 / 显存够

plugin = HybridParallelPlugin(
    tp_size=2, pp_size=1,
    zero_stage=1,
    precision="bf16",
    enable_flash_attention=True,
)

纯 TP=2、ZeRO-1，跑 LLaMA-7B 在 8×A100-40GB 上约 50% MFU。

② 13B / 单机 8 卡 / 紧张

plugin = HybridParallelPlugin(
    tp_size=4, pp_size=1,
    zero_stage=1,
    precision="bf16",
    enable_fused_normalization=True,
    enable_flash_attention=True,
)

TP=4 切 attention/MLP 权重，加 ZeRO-1 切 optimizer。

③ 70B / 多机 4×8 / PP 必须开

plugin = HybridParallelPlugin(
    tp_size=4, pp_size=4,            # DP = 32/(4·4) = 2
    zero_stage=1,
    precision="bf16",
    num_microbatches=8,
    pp_style="interleaved",
    num_model_chunks=2,
    enable_sequence_parallelism=True,
    sequence_parallelism_mode="ring_attn",
)

4D 并行（DP×TP×PP×SP），32 张 A100 训 70B 的实际配置。

6.5execute_pipeline：pp_size>1 时的特殊调用

开 PP 后训练循环要改写。看 colossalai/booster/plugin/pp_plugin_base.py：

# pp_size == 1 时（普通）
for batch in dl:
    loss = model(**batch).loss
    booster.backward(loss, optimizer)
    optimizer.step()

# pp_size > 1 时（必须）
for batch in dl:
    outputs = booster.execute_pipeline(
        data_iter=iter([batch]),     # 单个 batch 内部会被切成 micro-batch
        model=model,
        criterion=lambda outputs, batch: outputs.loss,
        optimizer=optimizer,
        return_loss=True,
    )
    optimizer.step()
    # backward 已经在 execute_pipeline 内部完成了

execute_pipeline 内部按 pp_style 选 schedule，处理 send/recv、调度 micro-batch、累计梯度，所以用户不能再写 loss.backward()。

6.6Sequence Parallelism 三种模式

短序列（<8K）建议关 SP；长序列（32K+）用 ring_attn。

6.7性能基线（参考）

LLaMA-7B、8×A100-80GB、bs=1024、seqlen=2048、bf16：

6.8这章你需要带走的

• HybridParallelPlugin = DP × TP × PP × SP × ZeRO 一键配置；
• 三个核心 size：tp_size / pp_size / sp_size；dp_size = world_size / (tp×pp) 自动推；
• 开 PP 必须用 booster.execute_pipeline(...)，不能用普通的 backward；
• 短序列用 split_gather SP；长上下文用 ring_attn；
• pp_style 优先 1f1b 跑通，激进可试 zbv；
• 显存够时 HybridParallel 比 ZeRO/Gemini 吞吐高 50%–100%。