Star139
196
代码
代码
Issues347
Pull Requests255
项目讨论
Wiki
分析
Star139
196
已开启
[Bug] FSDP CPU offload 开启后，梯度裁剪报 No backend type associated with device type cpu #379
fandengdong创建于  9月10日
fandengdong
fandengdong
9月10日 创建

Checklist

🐛 Describe the bug

开启 FSDP CPU offload（fsdp_config.enable_offload: true）后，训练在第一个 optimizer step 的梯度裁剪处崩溃：

RuntimeError: No backend type associated with device type cpu
  File ".../hyper_parallel/platform/torch/clip_grad.py", line 536, in _total_norm_fsdp2_aligned
    dist.all_reduce(local_p, op=dist.ReduceOp.SUM, group=group)

根因：offload 后参数与其分片梯度常驻 CPU，clip_grad_norm_ 因此把每个 rank 的 norm 标量也建在 CPU 上，再对该标量做跨设备进程组（hccl/nccl）all-reduce；这两种后端都没有注册 CPU 后端，于是报错。未开 offload 时标量本就在加速器上，因此不受影响。

复现步骤

  1. 任意训练配置加上 fsdp_config.enable_offload: true（FSDP / EP 组合均可复现）；
  2. 启动训练（torch 文本入口或 VLM 入口均可）；
  3. 第 1 步 forward/backward 正常通过，在 optimizer.step() 之前的梯度裁剪处崩溃。

最小复现规模：VLM 小模型 FSDP-8 + enable_offload: true 即可复现。实际触发场景为全量文本塔（61 层 × 384 路由专家，约 1T 参数）+ FSDP2-128 + EP128（16 节点 128 卡）。

已排查的相似 issue（非本问题）：#313 [RFC]: swap optimizer（optimizer state swap 机制设计，走 swap_optimizer() 路径）、#178（clip_grad 范数重复计数的数值问题）。

Expected behavior

enable_offload: true 时训练应正常进行——offload 只改变参数/梯度/优化器状态的驻留位置，不应导致梯度裁剪失败。

Additional context

  • 已完成定位与修复，修复 PR：#1385 fix(platform): reduce CPU norm scalars on accelerator
  • 相关背景：
    • #1217 feat: add fsdp mix_precision and cpu offload config（当初引入 enable_offload 的 PR）
    • #313 [RFC]: swap optimizer（另一个"训练状态放 CPU"的机制，走 swap_optimizer() 接口，与本问题不是同一路径）
  • 修复方式：把 norm 标量临时搬到本卡加速器完成归约、结果拷回原设备；非 offload 路径逐字不变
  • 验证结果：UT tests/ut/platform/torch/test_clip_grad.py 5/5；VLM FSDP-8 与 TP2×EP2 开 offload 后 loss 轨迹与非 offload 逐位一致；全量 1T + offload 在单超节点 128 卡上 3/3 步跑通（该配置不开 offload 会 OOM）

Environment info

  • 硬件：Ascend 910（64GB/卡），单超节点 16 节点 × 8 卡
  • torch 2.9.1、torch_npu 2.9.1、transformers 5.13.0、Python 3.11
  • CANN 9.0.0

Thanks for contributing 🎉!

likedislike
庄昭雄
庄昭雄
28 天前 评论：

补两点定位信息。先说路径：这条 issue 按 platform/torch/clip_grad.py:536 描述，但这个文件在 master 上不存在了，现在是 hyper_parallel/core/utils/clip_grad.py（933 行），对应的行是：

:491  def _total_norm_fsdp2_aligned(grad_groups, norm_type, mesh_cache, device, …)
:534  total_p = torch.tensor(0.0, device=device, dtype=torch.float32)
:538      norms = _per_tensor_norms(grads, norm_type, device)
:543      local_p = torch.tensor(0.0, device=device, dtype=torch.float32)   # 空桶
:546      dist.all_reduce(local_p, op=dist.ReduceOp.SUM, group=group)

关于根因，issue 的判断我核对后是成立的，可以再具体一点：这里 local_p 的设备完全由 device 决定，而 device 在 offload 下就是 CPU。all_reduce 要求输入与进程组的 backend 同设备，于是 CPU 张量碰上 NPU 的组。这个函数里只有一个 device 概念，没有区分「梯度所在设备」和「通信设备」—— 同文件 :226、:259、:281、:299、:303、:337、:344、:351、:366 也全是 .to(device=device)。

所以修法应当是把集合通信的输入搬到通信设备上、算完再搬回来。但我没能把「通信设备」确定下来：全仓搜下来 collectives/cc.py:98 的 get_backend(group) 只返回字符串；core/tensor_parallel/mc2.py:70 那种用法是 group._get_backend(torch.device("npu"))，得先知道设备才能问组，反推不出来。要落这个修法，得先定一个 backend→device 的映射（hccl→npu、nccl→cuda、gloo→cpu），那是新增的平台相关约定，我倾向于由你们决定放在哪里。

另外说明我为什么没有提交修复：这个错误需要「有 backend 的进程组 + CPU 张量」才会出现，而 gloo 本身就支持 CPU，所以纯 CPU 环境里设备与 backend 是匹配的，复现不出来。我手上没有 NPU，无法给出「修前失败／修后通过」的实测判据，就不把一个只靠推理的改动提上来了。

likedislike