已开启
[Bug] MindSpore DSA Indexer stateless dispatch错误封装变长序列参数导致grad_norm发散 #334
lzy0920232创建于  8月15日
lzy0920232
8月15日 创建

[Bug] MindSpore DSA Indexer stateless dispatch错误封装变长序列参数导致grad_norm发散

问题描述

Hyper-Parallel !1107 将 MindSpore custom op 从 AclnnOpRunner 迁移到 v2 stateless dispatch 后,Dense/Sparse Lightning Indexer 的 actual_seq_qlen 和 actual_seq_klen 仍以 optional<vector<int64_t>> 直接封装为 IntArray 参数。变长 TND/压缩 EOD 场景中,Indexer 前后向从第 3 个训练 step 开始偏离,后续 indexer_loss 与 grad_norm 指数级放大。

复现配置

  • MindFormers: 53429d9f21070ea7ed6c2987196661f2be82bd6f
  • Hyper-Parallel bad: f25e48d2ed78610c92418327138ed31239a17d2f 及最新 2fc3bfdc2438617147531da77cd28da0bca467d9
  • Hyper-Parallel good: fed8671b2d99405de064deb35be88d4145f9fe0d
  • 8 卡:TP1 + CP4(colossal) + DP2/full-FSDP,seq_length=2560,DSA1 dense indexer loss,LR=1e-3,从同一 MLA step100 checkpoint 训练到 step150。

A/B 结果

  • good fed8671b:step150 indexer_loss=0.020570,grad_norm=0.251021。
  • bad f25e48d2:step142 indexer_loss=111.299591;step150 grad_norm=2232283.0。
  • 最新 HP + 修复:step150 indexer_loss=0.020571,grad_norm=0.250388。

主 loss 在三组中一致,异常局限于 Indexer 辅助 loss/backward 链。

根因与方案

MS_DISPATCH_ACLNN 的 IntArray 输入需要 pair<vector<int64_t>, true>,不能传 pair<optional<vector<int64_t>>, true>。在 dense softmax LSE、dense grad KL loss 和 sparse grad KL loss wrapper 中先通过 value_or(empty vector) 解包,再传入 stateless dispatch。保留 v2 stateless dispatch,不回退整体接入方案。

likedislike
Llzy0920232
8月15日 关联了pull request:fix(mindspore): 修复DSA Indexer变长序列参数封装导致的梯度发散
Llzy0920232
8月15日 issue类型由 Documentation 改变为 Bug
Mmindspore-ci-bot成员
9月18日 关联了pull request:[mirror] mindspore-ai/hyper-parallel#82: fix(mindspore): 修复DSA Indexer变长序列参数封装导致的梯度发散