PP+CP+FSDP+TP 与单卡梯度对齐踩坑汇总：CP head_dim/seq 防呆缺失 + PP-stage-root fully_shard 梯度偏差

该问题是怎么引起的？

host 8×910B3 / torch_npu 2.7.1，做 PP(2)×DP-FSDP(2)×CP(2)×TP(1)（mesh=(pp,dp,cp,tp)=(2,2,2,1)，4 层 DenseFFN）与单卡逐参数梯度对齐时失配。定位出 4 项，每项给出根因与如何修改（hp 库侧 / 用户脚本+PP 侧）。文末给出最小可跑通集与优化项。①② 偏库侧防呆+文档，③ 示例工具，④ 真实库 bug。

① [库侧防呆 / 文档] ContextParallel 对 BNSD 的 head_dim 误配无早期校验

• 现象 / 根因：BNSD [B, N, S, D] 的 attention core 用 ContextParallel(seq_dim=2, head_dim=3)（误指每头特征维 D）；head_dim 应是 num_heads 轴（BNSD = 1，core/context_parallel/context_parallel.py:449）。_scatter_seq_to_head 的 divisibility 校验用 shape[head_dim] = 64（:394），64 % 2 == 0 蒙混过关，all-to-all 切掉收缩轴 → q @ kᵀ 退化为 Partial:['sum'] → div 崩；SkipDTensorDispatch 下静默算错。
• 如何修改 — hp 库侧：apply / pre-hook 增加 head_dim != seq_dim 校验，divisibility 针对真正 head 轴，parallel_ops.py 的 _check_partial_inputs 报错补「Partial 多因 ContextParallel head_dim 指到注意力收缩维」提示。
• 如何修改 — 用户脚本：cp_plan = ContextParallel(seq_dim=2, head_dim=1, ulysses_degree=cp_size)。

② [文档 / 防呆] Ulysses 要求按 seq 分片喂，库无显式校验

• 现象 / 根因：模型整栈跑整 seq、只在 attention 包 CP、两个 cp rank 喂同一条整序列 → 库 _to_cp_dtensor 把整 seq 当 Shard(seq_dim) 拼成 S·cp 重复伪全序列；head_dim 修对后 scores 变 [., ., S·cp, S·cp]，与 [S, S] 的 mask 不匹配，因果 mask 被 SDPACore 门控静默丢。契约见 tests/torch/context_parallel/_test_context_parallel.py:498-500（先 slice 再喂）。
• 如何修改 — hp 库侧：README / 示例补 BNSD / BSHD 的 seq_dim / head_dim 对照表 + 「每 rank 喂 seq 分片」醒目说明；可选在 pre-hook 支持传入全局 seq 做一致性校验。
• 如何修改 — 用户脚本 + PP 模型侧（端到端真用 CP）：
  1. 数据按 cp 切 seq（送进 schedule.run 前）：ls = seq // cp_size; x = x[:, cp_rank*ls:(cp_rank+1)*ls]，labels 同样切（参照官方 example 第 242 行 tokens_local = global_tokens[:, cp_slice]）。
  2. PanGuMoETransformerPP.forward 不变：embedding / RMSNorm / DenseFFN 都是 pointwise，对 local seq 透明；CP-wrapped sdpa_core 内部 all-to-all 把 local seq 重组成整 seq 算注意力、再切回；mask 传整 [S, S]（buffer 已是整 seq，all-to-all 后 scores = [., ., S, S] 匹配）。
  3. cp 维参数梯度规约（关键）：CP 参数在 cp 维是复制的，各 cp rank 只见 S/cp 个 token，对同一参数的梯度需 cp 维 all-reduce（sum）；loss 同理需 cp 维规约。建议 fully_shard 的 mesh 纳入 cp 维（(dp, cp)），或在 optimizer.step() 前对参数梯度在 cp 维 all-reduce。官方 example examples/torch/llama3/pp_fsdp_tp_cp_sp_example.py 只验证「能跑 + loss 有限」、未做与单卡的逐参数梯度对齐，也无显式 cp 维梯度规约 —— 此对齐需在 demo 里额外处理并实测。

③ [示例改进] PP stage 层名重映射应按 pp_rank 偏移

• 现象 / 根因：权重同步 / 梯度对比用 key.replace("layers.0.", f"layers.{rank // dp_size}.")，层偏移混了 dp / pp 且只替换字面 layers.0.。8-rank 验证：rank 0/1 对，2/3 第 0 层错，4/5 第 1 层错，6/7 两层全错 → 多数 rank 装错初始权重、配错梯度层。
• 如何修改 — 用户脚本（权重同步处 & 梯度对比处都改）：

import re
layers_per_stage = (num_layers + pp_size - 1) // pp_size
offset = pp_rank * layers_per_stage          # 只按 pp_rank，与 dp / cp 无关

def remap(key):
    return re.sub(r"layers\.(\d+)\.", lambda m: f"layers.{offset + int(m.group(1))}.", key)

# 权重同步: s_key  = remap(key)
# 梯度对比: s_name = remap(name2)

④ [库 bug] PP-stage-root fully_shard + GPipe 多 microbatch 下 root 级参数梯度偏差

• 现象 / 根因：fully_shard(chunk_model)（分片 stage root）+ GPipe 多 microbatch 时，root-FSDP 单元里的非 block 参数（末 stage 的 final_norm.weight / lm_head.weight）梯度算偏（稀疏 ~0.4% 元素超 0.002 容差，个别 ~28×）；各 block 单独 fully_shard(block) 全对。与 CP 无关（cp=1 仍失配）、与 mb 无关（mb=1 仍失配）；自 !790 [feature] torch backend pp support FSDP metastep 上线即在，非回归。怀疑 platform/torch/fully_shard/scheduler.py + platform/torch/pipeline_parallel/stage.py。
• 如何修改 — hp 库侧：修 root-FSDP grad reduce under PP metastep，并在 fully_shard docstring 明确 warn「PP stage root 不要 fully_shard」。
• 如何修改 — 用户脚本（规避）：删掉 fully_shard(chunk_model, mesh=dp_mesh)；改为对每个子模块（每个 block + final_norm + lm_head）逐个 fully_shard、不分片 stage root（参照 example _shard_blocks 的 "not the stage root"）。

---

✅ 最小可跑通集（改这些 → demo 与单卡梯度对齐）

1. ③ remap 改成 pp_rank * layers_per_stage 偏移 + 正则覆盖所有本地层（必须，否则权重 / 梯度配错层）。
2. ④ 删掉 fully_shard(chunk_model)，改为逐子模块 fully_shard（每个 block + final_norm + lm_head）、不 shard stage root（必须：规避库 bug，且让 all_gather 还原 full grad 的对比逻辑成立）。
3. CP 二选一：
   • （最省事·确定对齐）去掉 cp 维：mesh = (pp, dp, tp)、不 apply ContextParallel。
   • （保留 CP）head_dim=1 + 数据按 cp 切 seq + cp 维参数梯度规约（见 ②，需实测）。

🔧 优化项（分开，非跑通必需）

• hp 库侧：① head_dim 防呆校验；② seq 分片文档 / 校验；④ root-FSDP grad 根治 + fully_shard docstring warn。
• demo 增强：把上面「去 cp」升级为真正启用 CP（端到端 seq 切 + cp 维 grad 规约）并验证逐参数对齐。

重现步骤

环境：host 8×910B3，torch_npu 2.7.1，hyper_parallel editable 安装。

1. CP 正向 2×2 消融（cp=2，BNSD = [1, 32, 32, 64]，SkipDTensorDispatch 下对照单卡 SDPA）：head_dim ∈ {3, 1} × 喂法 ∈ {整 seq, seq 分片} 共 4 组，torchrun --nproc_per_node=2 比对与单卡 local 切片的 maxdiff。
2. ③ 重映射：对 mesh = (2, 2, 2, 1) 枚举 8 个 rank，分别按 rank // dp_size 与 pp_rank * layers_per_stage 计算层映射，对照是否命中正确全局层。
3. ④ root-FSDP：PP = 2 stage 上 fully_shard(layers[0]) + fully_shard(chunk_model)，GPipe 多 microbatch 训练后 all_gather 比对 final_norm / lm_head 梯度与单卡。
4. 最小集验证：应用「最小可跑通集」后，逐参数梯度对齐应通过。

报错信息

CP 正向 2×2 消融实测（与单卡 maxdiff）：

head_dim=3 + 整seq    : 3.481    MISMATCH   ← 现配置
head_dim=3 + seq分片  : 1.770    MISMATCH
head_dim=1 + 整seq    : 3.524    MISMATCH
head_dim=1 + seq分片  : 0.000000 MATCH      ← 唯一与单卡逐位相等

real DTensor 路径下 head_dim=3 直接崩（收缩轴被切 → 部分和）：

ValueError: For div, input 0 with Layout Configuration:
  Mesh shape: (2,)
  Alias Names: ('cp',)
  Partial: ['sum']
  Tensor Map: ('None', 'None', 'None', 'None')
  Rank List: (0, 1) has Partial status which is not allowed.

④ root-FSDP：final_norm.weight / lm_head.weight 梯度稀疏偏差，~0.4% 元素超 (rtol = atol = 0.002)，个别元素差约 28×；各 block 参数梯度全对。

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 241
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/241

Start with the CP matrix and minimal gradient-alignment reproduction described in the issue, then inspect core/context_parallel/context_parallel.py:394,449 and tests/torch/context_parallel/_test_context_parallel.py:498-500. For the FSDP discrepancy, read platform/torch/fully_shard/scheduler.py, platform/torch/pipeline_parallel/stage.py, and the example’s _shard_blocks path. Done means focused regression coverage and matching per-parameter gradients for the supported configurations.