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

[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，不回退整体接入方案。

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

Locate the Dense/Sparse Lightning Indexer wrappers for dense softmax LSE, dense grad KL loss, and sparse grad KL loss, then trace their MindSpore stateless dispatch arguments. Verify how optional sequence-length vectors are passed and compare the result with the reported good and bad revisions. Done means variable-length TND/compressed EOD training no longer causes Indexer loss or grad_norm to diverge.