已开启
动态图 DTensor in-place 算子 bypass 将梯度累加器降级为普通 Tensor,TP-4 梯度累加报 "Cannot found __fallback__" #190
xuxinglei创建于  6月6日
xuxinglei成员
6月6日 创建

该问题是怎么引起的?

hyper_parallel 的 DTensor dispatch 在 hyper_parallel/core/shard/_op_dispatch.py 里,对 bypass 白名单中的 in-place 算子(InplaceAddExt / InplaceSubExt / InplaceMul / InplaceDiv),会把参数从 DTensor 解包成本地 tensor、执行 op_call,然后返回那个本地结果

if self._should_bypass_dispatch(op_name):
    return op_call(*self._unwrap_args(args), **self._unwrap_kwargs(kwargs))

当第一个操作数是 DTensor 时,这等于把它降级成了普通 Tensor。MindFormers PynativeTrainer 的梯度累加 loss += micro_lossloss 初值为 0.0)因此:

  • 第 1 次 loss += micro_loss0.0 + DTensor 走 out-of-place,loss 变 DTensor;
  • 第 2 次:DTensor += DTensor 命中 in-place 白名单 bypass,返回本地结果,loss降级为普通 Tensor
  • 第 3 次:普通 Tensor += DTensor,去普通 Tensor 上找 __fallback__ 而失败。

因此只要梯度累加跑到 >= 3 个 micro-step 就复现:TP-4(dp_shard=2、每 dp rank 4 个 micro-step)会崩;TP-2(2 个 micro-step)到不了第 3 步所以不崩。

重现步骤

单机 8x910B,MindFormers PynativeTrainer(run_mindformer.py --mode 1)+ hyper_parallel:

  1. GLM-4.7 Flash 减层配置(4 层 / hidden 1024 / 8 experts,moe_router_force_expert_balance: false);
  2. 并行布局 FSDP-2 x TP-4 x EP-2tensor_parallel=4, expert_parallel=2, sequence_parallel=true);
  3. training.deterministic: True,训练 5 步。

结果:第 0 步在梯度累加 loss += micro_loss 处崩溃;改成 TP-2 则正常。

报错信息

File ".../mindformers/pynative/trainer/trainer.py", line 726, in _inner_train_loop
    loss += micro_loss
  ...
RuntimeError: Cannot found __fallback__ attr from object Tensor
likedislike