关联 Issue (Related Issue)

https://gitcode.com/mindspore/org-issues/issues/42712
https://gitcode.com/mindspore/org-issues/issues/42716

修改描述 (Description)

背景

在 TP×CP 并行下，MoE 路由的 token 视图被切分到多个 rank 上：

• TP 维度：router 输入在 TP rank 间被复制（路由前未做 token 切分），但其后的 dispatch 阶段会切分
• CP 维度：context parallel 把 sequence 切到不同 rank

旧实现中，TopKRouter 直接用本地 routing_map.sum(dim=0) 计算 tokens_per_expert，没有在 TP×CP 组内做 all-reduce 求和。这导致：

1. aux-loss 计算时每个 rank 只看到自己那份 token 的直方图，与单卡基线不一致
2. total_num_tokens 用的是 seq_length（local），没有还原成全局 token 数
3. 在 recompute 触发时，MoE 层的 tokens_per_expert 会被累加两次，aux-loss 随步数线性漂移
4. expert_bias 更新用的是本地直方图，bias delta 与单卡不一致
5. MTP 损失在 recompute 时也会被 tracker 累加两次
6. LossCallback 直接调用 tracker 函数，耦合度高，无法适配非 GPTModel 模型

修复

核心修改 — 在 TP×CP 组内做正确的 all-reduce：

• 新增 moe_utils.get_tokens_per_expert_and_token_count(routing_map, reduce_group, topk, with_padding_mask)：
  • 用 routing_map.sum(dim=0) 算 local per-expert 直方图
  • 沿 reduce_group（仅 TP×CP，不含 DP）做 in-place SUM all-reduce，还原 global 直方图
  • 推导 total_num_tokens = local_num_tokens * group.world_size()
  • DP 故意排除：DP 各 rank 本就该分别贡献自己的 aux-loss，由调用方用 /bsz 做平均；如果对 DP 做 all-reduce 会重复计数
• 新增 moe_utils.get_world_size_from_group(group)，兼容 hyper_parallel / 裸 group / get_world_size 三种 group size 获取方式，失败时降级为 1

TopKRouter 改造：

• 新增实例属性 self.aux_loss_group = None、self.aux_loss_group_size = 1，由外部在并行模式下注入
• seq_aux_loss 和 global_aux_loss 两条路径都改为调用 get_tokens_per_expert_and_token_count(...)，去掉手写的 tokens_per_expert = routing_map.sum(dim=0) 和手算的
  total_num_tokens
• 注释里说明 topk * bsz 是为了和 Megatron helper 签名对齐（当前未启用 padding 分支）

Recompute 双计防护：

• MoELayer.forward：在 tokens_per_expert.add_(...) 之前加 if not is_in_recompute(): 守卫，避免重计算阶段把 token 数累加两遍
• multi_token_prediction.save_to_mtp_losses_tracker：同样的 is_in_recompute() 守卫，避免 MTP 损失在 tracker 里被重复累加

Expert bias 更新修复：

• GPTModel._update_expert_bias(metric_group, metric_group_size)：
  • 在 metric_group（dp×cp 域）上对每层 tokens_per_expert 做 SUM all-reduce，恢复成 global-batch 直方图
  • DP 必须包含：每张 DP 卡看到的是 global batch 的不同 shard，本地直方图只是部分计数，不 reduce 就会和单卡基线发散
  • TP 故意排除：MoE plan 把 router 输入在 TP 间复制，TP 看到同样的 token，对 TP 归约会把计数乘以 tp_size
  • 用 _no_grad() 上下文执行 bias 更新和 tokens_per_expert 归零
• 引入 _get_expert_bias_modules() / _get_global_aux_loss_modules()，用实例属性缓存模块列表，VPP 下每个 chunk 一份独立 cache，互不干扰

模型接口泛化（LossCallback 解耦）：

• GPTModel 暴露 5 个方法取代 callback 直接调用 tracker：
  • _update_expert_bias(metric_group, metric_group_size)
  • get_load_balancing_loss(...) — 调用 track_moe_metrics（PP-group 组合 + dp/cp 域归约），每 stage 每 step 都必须调用以避免 deadlock
  • reset_model_temporary_tensors() — global_aux_loss 模式下重置 tracker
  • get_mtp_loss(metric_group, metric_group_size) — 调用 track_mtp_metrics
  • get_index_loss() — 调用 track_indexer_metrics
• PyNativeDeepseekV3ForCausalLM 透传这 5 个方法到内部 self.model
• LossCallback 改为通过 hasattr 探测方法是否存在并调用，缺失时只 warn 一次（用 logger_record 字典去重），错误信息附带具体模型类名

行为变化

• TopKRouter 现在要求外部在并行模式下注入 aux_loss_group；单卡/未启用并行的场景保持 None，行为与旧实现一致
• aux-loss 不再随训练步数漂移（TP×CP + DP 都正确还原）
• expert_bias 增量与单卡基线一致
• 触发 recompute 时 MoE token 计数和 MTP loss 不再被双重累加
• LossCallback 不再硬编码 GPTModel，可被其它 TrainModelMixin 模型复用

修改类型 (Type of Change)

• Bug 修复 (Bug fix)
• 新特性 (New feature)
• 文档更新 (Documentation update)
• 代码重构 (Refactoring)
• 性能优化 (Performance improvement)
• 其他 (Other)

测试结果 (Test Results)

检查清单 (Checklist)

• 我已自验过该功能/修复 (Self-checked)
• 我已通过了本地的单元测试 (Passed local UTs)
• (如适用) 我已更新了对应文档 (Updated documentation)
• (如适用) 我已添加了新的测试用例 (Added new tests)