Trainer 接入 Pipeline Parallel 设计方案

1. 背景与目标

当前 Trainer 以完整 Hugging Face 模型执行普通 forward。启用 Pipeline
Parallel（PP）后，用户仍提供完整模型，框架自动按 decoder layer 切分为逻辑
stage；每个物理 PP rank 只保留本地 stage，并在 stage 内继续应用 TP、CP、EP
和 FSDP。Trainer 保留数据准备、loss 统计、梯度裁剪、优化器和学习率调度。

当前交付覆盖 Qwen 风格 decoder-only causal language model，已验证裁剪的
Qwen3 dense 和 Qwen3-MoE。其他模型需要满足 stage forward 输入契约；跨层共享
可微状态、多模态输入或 encoder-decoder 结构需要额外适配。

2. 明确的对外配置

当前 VPP 配置如下：

accelerator:
  pp_size: 4

pipeline:
  enabled: true
  schedule: interleaved1f1b
  pp_microbatch_size: 1
  split: [[1, 1, 1, 1], [1, 2, 2, 1]]
  p2p_overlap: true
  p2p_multi_stream: false

该示例表示 4 个物理 PP rank、2 个 virtual chunk。二维 split 的每一行是一个
virtual chunk，每一列对应一个物理 PP rank。10 层模型的逻辑 stage 层数为
[1,1,1,1,1,2,2,1]，rank 0 持有 S0/S4，rank 1 持有 S1/S5，rank 2 持有
S2/S6，rank 3 持有 S3/S7。

普通 PP 使用一维 split：

pipeline:
  enabled: true
  schedule: 1f1b       # 或 gpipe
  pp_microbatch_size: 1
  split: [2, 3, 3, 2]
  p2p_overlap: true
  p2p_multi_stream: false

interleaved1f1b 是 VPP 的一种交错 1F1B 调度，VPP 本身指一个物理 rank 持有多个
逻辑 stage。两者不是同一个概念，但本接口当前只通过 interleaved1f1b 开放 VPP。
普通一维 PP 仍使用 1f1b 或 gpipe；二维 split 配这两个普通 schedule 会报错。
因此，原配置中的 schedule: 1f1b 在这个二维 split 示例中改为 interleaved1f1b，
其余字段保持不变；不再根据 split 静默改选调度器。

具体分配如下（layer 下标从 0 开始）：

逻辑 stage	PP rank	chunk	原始 layer	附加模块
S0	0	0	0	embedding
S1	1	0	1	—
S2	2	0	2	—
S3	3	0	3	—
S4	0	1	4	—
S5	1	1	5、6	—
S6	2	1	7、8	—
S7	3	1	9	final norm、LM head

逻辑 stage 按行展开，stage_id = chunk_id * pp_size + pp_rank。前向顺序为
S0 → S1 → … → S7；S3 → S4 需要从 rank 3 通信回 rank 0，不能把同 rank 的
S0 和 S4 直接串成一个 Sequential。反向沿相反的逻辑 stage 顺序运行。

3. 配置字段和约束

配置	含义与约束
accelerator.pp_size	物理 PP rank 数量；VPP 不增加设备数量。
pipeline.enabled	必须与 pp_size 大于 1 一致。
pipeline.schedule	普通 PP 为 gpipe 或 1f1b；VPP 为 interleaved1f1b。
pipeline.pp_microbatch_size	每个 pipeline microbatch 的样本数，必须为正数。
pipeline.split	空列表自动均衡；一维列表表示普通 PP 各 rank 层数；二维列表表示 VPP 的 virtual_chunk 与 pp_rank 层数。
pipeline.p2p_overlap	是否延迟 P2P 等待到计算消费前，默认 true。
pipeline.p2p_multi_stream	是否使用多流 P2P，默认 false；开启时必须同时开启 p2p_overlap。

二维 split 必须至少两行，每行长度等于 pp_size，元素为正整数，且总层数等于
模型 decoder layer 数。VPP 数量由二维 split 行数推导，不另设 vpp_size。空行、
空 stage 和层数不匹配会在构建阶段报错。

Trainer 的 batch 关系为：

B_local = global_batch_size / dp_size
num_microbatches = B_local / pp_microbatch_size

两次除法都必须整除；pp_microbatch_size 与数据加载器的
training.micro_batch_size 是不同配置。
以上 YAML 是合并到已有模型、数据、优化器配置中的片段，不是独立训练配置。
例如 DP=1、global_batch_size=4 时，本例产生 4 个 PP microbatch；VPP chunk 数
不参与 batch 的再次切分。所有元素表示 layer 数量，而不是 layer ID 或边界下标。
空列表只用于普通 PP 自动均衡；interleaved1f1b 配空列表、一维列表或只有一行的
二维列表均拒绝。关闭 PP 时同时设置 enabled=false 和 pp_size=1。

4. 对外 Python 接口

device = torch.device("npu", torch.npu.current_device())
pipeline = AutoPipeline(world_mesh=mesh, config=config.pipeline)
pipeline.build(
    model=full_model,
    parallelize_fn=apply_stage_parallelism,
    batch_size=dp_local_batch_size,
)
optimizer = torch.optim.SGD(pipeline.parameters(), lr=1e-3)
optimizer.zero_grad()

losses = pipeline.step(
    input_ids if pipeline.info.has_first_stage else None,
    target=labels if pipeline.info.has_last_stage else None,
    attention_mask=attention_mask,
)
optimizer.step()

AutoPipeline.build 接收完整模型，完成切分、裁剪、stage forward 改写、
stage-local 并行包装、PipelineStage 构建和 schedule 构建。Trainer 不直接对
单个 stage 调用普通 forward，而是调用 pipeline.step 驱动完整调度。
AutoPipeline 本身就是返回给 Trainer 的 nn.Module；parts 注册当前 rank 全部
本地 chunk，info 保存 schedule、stages 及首尾 stage 标志。当前接口不存在
pipeline.local_model，不使用早期 PipelineRuntime 设计。上述直接调用示例要求
已经初始化分布式 mesh，完整模型位于正确设备且参数有效；Trainer 场景自动构建。

5. 构建顺序

完整模型（默认 Meta 构建）
  -> 发现 decoder.layers 并生成逻辑 stage 分配
  -> 当前 PP rank 选择本地 stage
  -> 深拷贝并删除非本地 layer，保持原始 FQN
  -> 改写首、中、末 stage forward
  -> 保持 stage 参数为 Meta，不提前分配 NPU 参数存储
  -> 对每个本地 stage 应用 TP、CP、EP、activation 和 FSDP
  -> 构造 PipelineStage(stage_index=逻辑 stage id)，通信设备取当前 NPU
  -> 按 schedule 构造 GPipe、1F1B 或 Interleaved1F1B
  -> 外层统一 model.to_empty(device)，再使用已有流程读取权重或初始化参数
  -> 重建必要的 buffer
  -> Trainer 构建包含全部本地 chunk 参数的 optimizer

FQN 保持是 splitter 的硬性契约。TP/CP/EP planner、FSDP 参数分组和 checkpoint
依赖原始名称；外层 parts 前缀通过 checkpoint_fqn_aliases 映射。

PP 与非 PP 使用同一套 materialize 生命周期：Meta 模型先完成切分以及 stage-local
TP/CP/EP/FSDP 包装，FSDP 记录参数的 DTensor 和 shard 元数据；随后由外层
_materialize_and_load_model() 对整个 AutoPipeline 调用 model.to_empty(device)，
再加载 checkpoint 或执行随机初始化。PipelineStage 的 P2P 通信设备取当前
NPU 加速器，不从仍处于 Meta 状态的参数推导通信设备。

该顺序不是“先把完整模型权重加载到每张卡”：切分和 stage-local sharding 仍先完成，
to_empty 只为当前 rank 的本地参数分配存储，随后才写入 checkpoint 或初始化值。
不过单个 stage、多个 VPP chunk、临时 buffer、激活和原始实参模型副本仍会影响峰值，
不能承诺任意模型不 OOM。首次调度执行时，FSDP 的 lazy initialization 会重新读取
materialize 后的参数，刷新内部 shard 引用，并检查参数已不在 Meta 上。

关键实现伪代码（device 仅为 NPU）：

device = torch.device("npu", torch.npu.current_device())
with torch.device("meta"):
    full_model = Model(config)

def parallelize_stage(stage):
    stage = apply_tp_cp_ep(stage)
    stage = apply_activation_features(stage)
    return apply_fsdp(stage)             # 参数仍为 Meta

model = AutoPipeline(pp_mesh, pipeline_config).build(
    full_model,
    parallelize_fn=parallelize_stage,
    batch_size=dp_local_batch_size,
)
# 以下由已有 _materialize_and_load_model / _move_model_to_device 统一执行。
model.to_empty(device=device)           # 为本地参数分配存储，不初始化值
load_or_initialize(model)
rebuild_non_persistent_buffers(model)
optimizer = Optimizer(model.parameters())
model.step(...)                         # HSDP lazy init/reset，随后执行调度

原问题在于 device=next(part.parameters()).device 将 Meta 参数设备作为 P2P
buffer 设备，导致 dist.irecv(meta_tensor, ...) 失败；不是 FSDP 禁止 Meta 构建，
也不是外层 to_empty 必然破坏 FSDP 引用。因此删除 stage 内提前 to_empty 和
already_materialized 跳过分支，修复运行设备与参数存储设备混用的问题，复用原有
非 PP 初始化流程。AutoPipeline.build 不新增运行设备参数；在明确使用 NPU 的
Trainer 路径中，PipelineStage 直接取当前 NPU 作为 P2P 通信设备，无需新增 YAML
配置。

6. Stage forward 和输入

• 首 stage 接收本地 input_ids，执行 embedding 和自己的 decoder layers。
• 中间 stage 接收上游 hidden_states，并在本 rank 使用本地 attention_mask、
  position_ids 等辅助输入。
• 末 stage 执行 final norm、LM head 和 loss。
• P2P 只传输 stage 边界所需 Tensor；辅助输入由 Trainer 在各 rank 重建或提供。

当前 Qwen stage forward 的反向边界是 hidden_states。下游 backward 得到输入
hidden_states.grad，通过 PP runtime 传给上游；每个 stage 只更新自己的本地参数。
每个 stage 是本地 autograd 图，跨进程梯度连接由显式 P2P 实现，不是共享整张图。
当前 patch 本地构建 causal mask 和 rotary embedding，use_cache=False；输入 batch
及辅助 kwargs 的 sample/microbatch 顺序必须在同一 PP line 上一致。
compressed KV、index key、candidate blocks、4D mHC residual 等模型专有可微
状态需要扩展 StageState 输入、输出和梯度协议，不能仅通过 split 配置自动支持。

7. Schedule 行为

普通 gpipe：所有 microbatch 先完成前向，再按逆序完成反向，显存占用较高。

普通 1f1b：warmup 后交替执行前向和反向，降低激活峰值。

interleaved1f1b：每个物理 rank 拥有多个 virtual chunk，调度器按
logical_stage_id 和 microbatch_id 交错执行本地 chunk 的前向、反向和 P2P。
多个 chunk 必须由同一个 Interleaved1F1B 调度器管理。

Trainer 一个 step 的核心流程：

batch = concat_dp_local_batches()
optimizer.zero_grad()
losses = pipeline.step(
    batch.input_ids if pipeline.info.has_first_stage else None,
    target=batch.labels if pipeline.info.has_last_stage else None,
    attention_mask=batch.attention_mask,
)
grad_norm = trainer.prepare_optimizer_step()
trainer.step_optimizers_and_schedulers()

pipeline.step 已完成本 step 的 forward、backward、P2P 和本地梯度累计；Trainer
不再对返回 loss 额外调用 backward。
只有最后逻辑 stage 产生 loss；Trainer 将各 rank 指标求和，再除以
world_size / pp_size，使用物理 PP 数而不是逻辑 stage 总数。FSDP 的 unshard、
reshard、梯度归约由底层 schedule 针对各本地 stage 管理。

注意 loss 指标归约不等于训练梯度归一化。当前默认 loss 是每个 microbatch 的
mean CE，反向累计不自动等价于完整 batch 的 token mean；严格精度比较必须
固定 loss reduction、有效 token 权重和归一化。跨 stage MoE auxiliary loss
汇总不能从普通 CE 路径推断为已实现，需要单独适配和测试。

8. 混合并行、checkpoint 和边界

PP 使用独立的一维 pp_mesh。每个物理 PP rank 的本地 stage 使用对应的
stage-local mesh 应用 TP、CP、EP 和 FSDP；VPP chunk 不增加设备网格维度。
所有 rank 必须使用一致的 pp_size、split、schedule 和 P2P 配置。
拓扑满足 world_size = pp_size * tp_size * cp_size * dp_size；EP 依赖现有
专家布局和 planner 的嵌套关系，不能简单作为额外维度再乘一次。

p2p_overlap 透传 overlap_p2p；开启多流时 p2p_transport=multi_stream，否则为
auto。多流是底层实验性传输，不等同于 zero-bubble 或 B/F 计算重叠。调度器
消费接收 buffer 前必须等待完成，发送 buffer 也不能提前释放或复用。

保存和加载使用完整模型的原始 FQN；每个 PP rank 只处理自己负责的 layer，再按
TP/CP/EP/FSDP layout 读写本地 shard。完整断点续训、跨 PP size 转换、跨层 tied
weight 和特殊跨层状态仍需单独验收。

当前明确支持 Qwen 风格 decoder-only 模型、普通 GPipe、普通 1F1B、二维 split
的 VPP Interleaved1F1B、stage-local TP/CP/EP/FSDP，以及 P2P overlap 和可选
多流 P2P。不支持空 stage、二维 split 配普通 schedule，未适配模型的特殊
forward contract、encoder-decoder 和多模态输入。

9. 测试和交付状态

已运行本地回归：

tests/ut/auto_models/_transformers/test_pipeline.py
tests/ut/auto_models/_transformers/test_deferred_pretrained.py
tests/ut/trainer/test_distributed_runtime.py
tests/ut/trainer/test_text_trainer.py
合计 20 passed，60 subtests passed

四卡 NPU ST 已覆盖真实 apply_model_infrastructure 路径：Qwen3 普通 1F1B、
VPP interleaved1f1b、FSDP comm_fusion=false/true 均完成两步训练；Meta 参数在
FSDP 包装后再由外层 to_empty 成功 materialize，loss、梯度和参数更新与完整模型
对齐。独立四层 Qwen3 诊断记录中，最大梯度误差约 1.9e-6，最大更新后参数误差
约 3.7e-9。本轮另已重跑四卡十层 Qwen3 的 GPipe、普通 1F1B、VPP
以及 P2P 关闭/overlap/多流的九种组合，均完成两步 loss、梯度、参数更新对照。
本轮未重新验证八卡 Qwen3-MoE 的 TP/CP/EP 混合并行及断点续训。

最终接口不新增 runtime_device 参数，直接复用 PipelineStage 当前加速器默认值，
在本环境中为当前 NPU。移除该参数后重新独立启动上述两项四卡 ST，分别通过
（32.48 秒、35.24 秒）；本轮未再次发生此前的 HCCL/SDMA 内部错误。

对应四卡 ST：

• tests/torch/pipeline_parallel/test_auto_pipeline.py::test_qwen3_meta_pipeline_fsdp_matches_reference
• tests/torch/pipeline_parallel/test_auto_pipeline.py::test_qwen3_pipeline_splits_match_reference

验收要求：覆盖 split 合法/非法组合、首/中/末 stage、所有本地 chunk 并行包装
和 FQN alias；同初始化、batch、dtype、loss 归一化下比较非 PP/PP/VPP 的 loss、
各参数梯度及更新后参数；执行多 microbatch、多 optimizer step 以检查缓存残留
和死锁；覆盖 P2P 关闭/overlap/多流、混合并行及 PP 关闭回归。Checkpoint 需
另外验收保存、恢复 optimizer 状态、继续训练；性能用固定硬件与 batch 实测
吞吐、step time、峰值内存，不承诺 VPP 一定加速。

主要实现文件：

• hyper_parallel/models/_transformers/pipeline.py
• hyper_parallel/models/_transformers/model_builder.py
• hyper_parallel/trainer/runtime/distributed.py
• hyper_parallel/trainer/text_trainer.py
• hyper_parallel/core/pipeline_parallel/scheduler.py

代码基线位于 GitHub trainer_pp_master 分支，PR：
https://github.com/mindspore-ai/hyper-parallel/pull/1280

截至 2026-10-26，本正文中的 Meta/materialize 生命周期修正及当前 NPU 默认通信
设备修订已实现并通过上述 UT/ST，提交 07155cc1 已推送到 PR 的
trainer_pp_master 源分支。