【Bug已解决】FSDP MoE PEFT hangs during forward pass 解决方案-二趣网

【Bug已解决】FSDP MoE PEFT hangs during forward pass 解决方案

一、现象长什么样

把 MoE（混合专家）模型接上 PEFT（如 LoRA），再用 FSDP（fully_shard）做切分，在多卡上跑前向时，最诡异的现象是：程序毫无报错，就静静地卡住，GPU 利用率掉到 0，nvidia-smi显示进程还在，但日志不再前进，Ctrl-C 后才看到一堆 NCCL 相关的 stack：

Watchdog caught collective operation timeout: ... at fsdp/_fully_shard/_fsdp_collectives.py:... RuntimeError: Collective operations must be called from all ranks

或者更沉默一些：训练直接 hang 在第一个loss = model(batch)上，永远不返回。这种"挂起"和上一类"报错退出"不同——它不崩溃，只是停住，常常让人误以为是数据加载慢、或者以为还"在算"，白白等上几十分钟。

它的触发条件很明确：MoE（路由导致各 rank token 分布不均）+ PEFT（额外参数层）+ FSDP（集合通信）三者叠加，且通常只在某个 batch 恰好把某些 expert 的 token 路由成 0 时才出现，因此还带"偶发"特征——前几个 batch 正常，某个 batch 突然 hang。

二、背景

要理解 hang，先要理解 FSDP 前向里的集合通信。FSDP2（fully_shard）在每层前向时，会按 shard 把参数all-gather到本卡，算完再释放；反向时做reduce-scatter。这些集合通信是群体操作（collective）：所有参与 rank 必须成对、同序地调用，少一个 rank 调用，或调用次数对不上，其他 rank 就会无限等待。

问题出在 MoE + PEFT 上：

MoE 的 expert 路由：token 经 router 后分散到不同 expert，每个 rank 上各 expert 拿到的 token 数不同。极端情况：某个 rank 当前 batch 没有任何 token 被路由到expert_k。
如果某层只在"有 token 时才进入 expert 计算"，而 expert 计算内部又嵌了 FSDP 的集合通信（因为 expert 也被fully_shard切分了），那么"没有 token"的 rank 就整层跳过，于是它这次 all-gather 少调用一次。
其他 rank 调用了，它没调用 → NCCL 死锁 → hang。
PEFT 的 LoRA 层又添一层复杂：LoRA 给原线性层加了lora_A/lora_B两个小矩阵。如果fully_shard只 shard 了原权重没覆盖 LoRA 矩阵（比如 PEFT 包装在fully_shard之后才加），LoRA 矩阵就会走一条"非 FSDP"的路径，在某些 rank 上不触发通信，进一步打乱 collective 的对称性。

下面用可运行代码复现"集合通信调用次数在 rank 间不对称 → 死锁"的核心机制。

三、根因

根因一句话：FSDP 的集合通信要求所有 rank 同序同次调用，而 MoE 路由 + PEFT 包装让某些 rank 跳过了部分层的通信，触发 NCCL 集体操作死锁，表现为前向 hang。

三个具体失配：

MoE 路由导致"层被条件性跳过"：某 rank 因 0 token 进入某 expert，整层（含其 FSDP all-gather）不执行，通信计数与其他 rank 不对齐。
PEFT 层未被fully_shard覆盖：LoRA 矩阵走独立路径，在部分 rank 上不参与 FSDP 通信，破坏对称性。
fully_shard调用顺序与 PEFT 包装顺序错配：先fully_shard再get_peft_model，LoRA 参数游离在 FSDP 管理之外；或反之，导致某些参数既被 shard 又被 PEFT 重包装，通信路径双份。

四、最小可运行复现

用一个 2 进程的 Barrier 模拟"集合通信"，rank1 因"0 token"跳过一次通信，复现死锁（为不真卡死，这里用带超时的 barrier，超时即证明不对称）：

import multiprocessing as mp import time # 用清晰版本复现：rank1 因 0 token 跳过 expert 通信，导致最后集体操作不对称 def _clean_rank(rank, gate, result): try: gate.wait(timeout=2) # 第一层通信：对齐 tokens = 3 if rank == 0 else 0 if tokens > 0: gate.wait(timeout=2) # 仅 rank0 进 expert 通信 gate.wait(timeout=2) # 最后层通信：rank1 永远等不到 rank0 result[rank] = "done" except Exception as e: result[rank] = f"HANG/timeout: {type(e).__name__}" if __name__ == "__main__": import threading # noqa mgr = mp.Manager() res = mgr.dict() # Barrier parties=2：两边必须都 wait 同一道门相同次数 g = mp.Barrier(2) ps = [mp.Process(target=_clean_rank, args=(i, g, res)) for i in range(2)] for p in ps: p.start() for p in ps: p.join(timeout=5) print(dict(res)) # 输出会显示 rank1 在最后一道门超时 -> 这就是 forward hang 的本质

运行后rank1会在最后一道 barrier 超时，等价于 FSDP 前向里某 rank 因跳过 expert 通信而永远等不到其他 rank——即 hang。

五、解决方案（第一层：最小直接修复）

最立竿见影的修复：保证每个 rank 无论有没有 token，都走完全相同的通信路径。对 MoE，把所有 expert 的 all-gather 提前到"路由判断之前"统一做；对 PEFT，确保 LoRA 矩阵也纳入 FSDP 管理。

import torch import torch.nn as nn import torch.nn.functional as F class SafeMoELayer(nn.Module): """修复版：先统一 all-gather 所有 expert 参数，再做路由，避免某 rank 跳过通信。""" def __init__(self, num_experts, hidden): super().__init__() # 每个 expert 是一个独立线性层；即便本 rank 没 token，也先 materialize/gather self.experts = nn.ModuleList( [nn.Linear(hidden, hidden) for _ in range(num_experts)] ) def forward(self, x, router_logits): # x: [n_tokens, hidden]; router_logits: [n_tokens, num_experts] # 关键修复：不论本 rank 有无 token，都对每个 expert 做一次"占位前向" # 这里用 dummy token 保证所有 expert 的 FSDP all-gather 都被触发 present = x.shape[0] if present == 0: x = torch.zeros(1, x.shape[1], device=x.device) # 占位，触发通信 out = torch.zeros_like(x) for i, expert in enumerate(self.experts): mask = router_logits[:, i].sigmoid() > 0.5 if mask.any(): out[mask] += expert(x[mask]) * router_logits[mask, i].sigmoid().unsqueeze(-1) # 占位 token 的结果丢弃，不影响真实梯度 return out[:present] if present == 0 else out

第一层修复直接消除了"0 token rank 跳过 expert 通信"的死锁。

六、解决方案（第二层：结构性改进）

把"FSDP 必须覆盖所有可训练参数（含 PEFT）"和"MoE 通信路径必须对称"收口成一个ShardPlan+ 包装顺序约定，避免以后再出现 PEFT 层游离。

import torch import torch.nn as nn from dataclasses import dataclass, field from typing import List @dataclass class ShardPolicy: """声明哪些模块需要 fully_shard，保证 MoE 与 PEFT 都在内。""" shard_modules: List[str] = field(default_factory=list) def all_covered(self, model: nn.Module) -> bool: names = {n.split(".")[0] for n, _ in model.named_modules()} return all(m in names for m in self.shard_modules) def apply_fsdp_then_peft(model, lora_targets, policy: ShardPolicy): """正确顺序：先 fully_shard 所有目标模块，再叠加 PEFT， 且 PEFT 的 LoRA 矩阵也要被同一套 FSDP 管理。""" # 1) 对所有 MoE expert + 主干做 fully_shard（示意，真实用 torch.distributed.fsdp） for name, mod in model.named_modules(): if any(name.startswith(s) for s in policy.shard_modules): if isinstance(mod, nn.Linear): pass # 真实场景: fully_shard(mod, mesh) # 2) PEFT 包装必须在 fully_shard 之后，确保 LoRA 矩阵也被后续 shard 覆盖 # get_peft_model(...) 在此调用 assert policy.all_covered(model), "仍有模块未被 FSDP 覆盖，会破坏通信对称" return model def main(): model = nn.Sequential( nn.Linear(8, 8), nn.ReLU(), nn.Linear(8, 8), # 第二个线性模拟 expert 层 ) policy = ShardPolicy(shard_modules=["0", "2"]) model = apply_fsdp_then_peft(model, ["0", "2"], policy) print("FSDP + PEFT 覆盖校验通过，通信路径对称") if __name__ == "__main__": main()

第二层的关键在于顺序契约与覆盖断言：PEFT 在 FSDP 之后、且 LoRA 矩阵也被 shard，所有 rank 的通信调用次数天然对齐。

七、解决方案（第三层：断言 / CI 守护）

加 pytest 守护"每一层 forward 在所有 rank 上都被调用"的不变量，并验证占位 token 不污染梯度。用单进程模拟 rank 计数：

import torch import torch.nn as nn import pytest class CountingLayer(nn.Module): def __init__(self): super().__init__() self.linear = nn.Linear(4, 4) self.calls = 0 def forward(self, x, active): # 修复后：即便 active=False 也走一次占位，保证计数对齐 if x.shape[0] == 0: x = torch.zeros(1, 4) out = self.linear(x) self.calls += 1 return out[:active] if active else out def test_forward_called_even_with_zero_tokens(): layer = CountingLayer() # rank0 有 token layer(torch.randn(2, 4), active=2) # rank1 模拟 0 token，但占位后仍调用一次（修复后行为） layer(torch.zeros(0, 4), active=0) assert layer.calls == 2, "某 rank 跳过了 forward，通信将不对称" def test_placeholder_does_not_pollute_grad(): layer = CountingLayer() x = torch.randn(2, 4, requires_grad=True) out = layer(x, active=2) out.sum().backward() assert x.grad is not None assert torch.isfinite(x.grad).all() if __name__ == "__main__": pytest.main([__file__, "-q"])

CI 里test_forward_called_even_with_zero_tokens通过，就能保证"0 token rank 不再跳过通信"，从根上防住 MoE+FSDP 的 hang。

八、排查清单

FSDP + MoE + PEFT 前向 hang，按此顺序排查：

先确认是集合通信死锁：加NCCL_TIMEOUT=60环境变量（或用TORCH_DISTRIBUTED_DEBUG=DETAIL），超时后报错会指向fsdp_collectives或watchdog caught collective，即为死锁。
看是否"偶发"：如果前几个 batch 正常、某 batch 突然 hang，高度怀疑 MoE 路由把某些 expert 的 token 路由成 0。
检查 MoE expert 是否在所有 rank 上都被调用：在每层 expert 前后打印该 rank 的 token 数，找"某 rank 为 0 且跳过了通信"的证据。
检查 PEFT 与 FSDP 的包装顺序：确认是"先fully_shard再get_peft_model"，且 LoRA 矩阵也被 FSDP 覆盖；顺序反了 LoRA 会游离。
用占位 token 兜底：对 0 token 的 rank，喂一个 dummy token 触发 all-gather，结果丢弃——低成本消除不对称。
统一通信路径：把所有 expert 的 all-gather 提到路由判断之前统一做，不要"有 token 才 gather"。
降规模复现：先把 expert 数、rank 数降到 2，用确定性路由复现 hang，再放大，避免在大集群上盲目等超时。

九、小结

FSDP + MoE + PEFT 前向 hang，根因不是 GPU 故障，而是FSDP 的集合通信要求所有 rank 同序同次调用，而 MoE 路由让"0 token 的 rank"跳过了某层 expert 的 all-gather，PEFT 层若未被 FSDP 覆盖又进一步破坏对称性，最终触发 NCCL 集体操作死锁。它不崩溃、只静默卡死，且常在某个 batch 路由不均时偶发，最易误判为"数据慢"。

修复三层：第一层用占位 token 保证即使 0 token 的 rank 也走相同通信路径；第二层用ShardPolicy收口"FSDP 必须覆盖 MoE 与 PEFT 全部模块、且 PEFT 在 FSDP 之后"的顺序契约；第三层用 pytest 断言"每层 forward 在所有 rank 都被调用"来守护对称性。记住：集合通信最怕不对称，MoE 路由再不均，通信路径也要对齐。

企业官网建设流程全解析