mindspore-ai / mindspore-ai/hyper-parallel

动态图 DTensor in-place 算子 bypass 将梯度累加器降级为普通 Tensor,TP-4 梯度累加报 "Cannot found __fallback__"

Open
#269 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
53
Forks
63
Avg merge
23h 45m
Merged PRs (30d)
63

Description

该问题是怎么引起的?

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

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

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in hyper_parallel/core/shard/_op_dispatch.py at the _should_bypass_dispatch branch for the listed in-place operators, then trace how its returned value is used by MindFormers' trainer.py loss += micro_loss path. Reproduce with the TP-4 configuration described in the issue; done means gradient accumulation reaches the third micro-step without downgrading the DTensor or raising the fallback error.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
58/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.