已开启
【bug】MindSpore fully_shard 参数 view 被非重入重计算捕获导致 unpack 报错 #164
DavidFFFan创建于  5月26日
DavidFFFan
DavidFFFan成员
5月26日 创建

该问题是怎么引起的?

在 MindFormers pynative 训练中,同时满足以下条件时会触发问题:

  • transformer layer 配置非重入重计算(use_reentrant=False)。
  • transformer layer 内部的 norm 子模块配置 fully_shard。
  • backward 阶段由 recompute 重跑 transformer layer forward。

重算 forward 会再次执行 norm 的 fully_shard forward pre hook,并进入 MindSporeHSDPParamV2.reset_sharded_param()。其中用于刷新 sharded local view 的 local_tensor.narrow(...) 会进入自动微分,被非重入 recompute 的 saved tensor hooks 记录为 placeholder。

该 tensor 属于 fully_shard 内部参数 view 维护,并非模型 forward activation。它进入 placeholder 序列后,原始 forward 与重算 forward 的 saved tensor 序列可能不一致,导致 backward unpack 时无法找到匹配的重算 tensor。

重现步骤

  1. 在 MindFormers pynative 配置中,为 transformer layer 启用 full recompute,并对 layer 内部 norm 启用 fully_shard。
  2. 使用双卡运行训练:
ASCEND_RT_VISIBLE_DEVICES=1,2 \
PYTHONPATH=/home/fdw/code/hyper-parallel:$PYTHONPATH \
/home/fdw/miniforge3/bin/conda run -n py310 \
bash scripts/msrun_launcher.sh "run_mindformer.py --config pynative.yaml --mode 1" 2
  1. 训练进入首个 step 的 backward 后,非重入重计算恢复 activation 时触发异常。

定位结果:pack 的调用链为:

_forward_pre_hook
-> _hsdp_forward_pre_hook
-> hsdp_state.unshard()
-> hsdp_param.unshard()
-> _get_unsharded_param_data()
-> reset_sharded_param()
-> local_tensor.narrow(...)

报错信息

RuntimeError: Unpack is being triggered for a tensor, make sure to do this only once!

报错位置:

mindspore/common/recompute.py::_CreatePlaceholderHook.unpack

期望行为:fully_shard 内部用于维护 sharded parameter view 的操作不应污染 recompute 的 activation placeholder 序列,同时应保持 local tensor 的 requires_grad 状态不变。

likedislike