feat: add DSV4.1 sequence and modality load balancing

[Er1ckPang]

Paired: GitHub #1067 ↔ GitCode !1717

What type of PR is this?

/kind feature

---

What does this PR do / why do we need it:

为 DeepSeek V4.1 增加逐逻辑样本的 seq-auxloss，以及文本/图像独立统计、每个训练步独立更新 correction bias 的 aux-loss-free 负载均衡。避免不同样本或模态的偏载在合并统计后互相抵消。

• seq-auxloss：全专家原始 scores 归一化，单独使用未加 bias 的 topk(scores) 计算每个样本的目标，对非空样本平均。保留 [B,S,H] 的 batch 边界及 packed cu_seq_lens 的逻辑样本边界，支持有效 token mask 和空分片。
• 实际 dispatch 与模态负载使用 topk(scores + correction_bias)；专家输出权重仍从原始 scores gather。bias 改变实际路由时，不能把该集合复用于 seq 目标。
• 分模态 bias：text/image 各有一行 int64 实际专家计数，在 backward 中累计，避免 checkpoint forward 重算重复计数。每步所有 optimizer 完成后归约，分别执行 bias += rate * sign(mean(load) - load) 并清零；缺席模态不更新。
• 保留 bias/bias_vl checkpoint 名称，关闭其梯度。支持 HP DTensor 分片后的原地更新；非持久化计数在 checkpoint materialization 后重建为零。
• 统计域：seq 目标归约 CP 与 EP 场景的 TP token 分区；bias 另汇总 DP 各轴。EP dispatch 组不额外作为统计轴。模型 adapter 同步切分 sample ID、image mask 和 token mask。
• 归属：公式、计数与更新逻辑位于 models/deepseek_v41/aux_loss.py；公共层保留 Router/Indexer 共用的梯度注入与通用训练生命周期。
• text/VLM builder 与 recipe 使用 router_aux_loss_coeff=0.0001、router_bias_update_rate=0.001，可分别用0关闭；注入梯度跟随 foundation loss 的 microstep/token/并行倍率。

依据 V4.1 官方报告第8、22页及其继承的 V3 sequence-wise 定义。公开 V4.1 推理代码仅用于核对实际路由，没有训练 auxloss。参考 TorchTitan 固定版本的注入/更新机制；其 folded stream 统计及带 bias routing map 不作为本实现的 seq 公式 oracle。

基于 master 68fc3e33，保留已合入 PR1066 的 FP32 clamp/SwiGLU、shared-expert 绑定和 grouped-GEMM 路径；aux 接线兼容已有文本/视觉 EP forward 签名。已回退与 aux 无关的 parity 重排及 pylint 屏蔽，新增标注逐项核对后仅保留12处静态误报豁免。

---

Which issue(s) this PR fixes:

暂无关联 issue（未提供 issue/SR/AR 编号）。

---

Test Plan and Test result：What scenarios were tested, and what were the verification results（Function, performance, reliability, etc.）：

• CPU UT：python -m pytest tests/ut/auto_models/models/deepseek_v41 tests/ut/components/functional tests/ut/core/utils/test_moe_utils.py tests/ut/trainer -q，116 passed、86 subtests passed。
• 新增真实 router 回归：独立改变 text/image bias，实际路由和计数随之变化；seq loss、router weight/input 的辅助梯度保持与未加 bias 的独立公式一致。测试在修正前失败，修正后通过。
• 双进程 CPU/Gloo：两 rank 各 1 passed；非零模态 bias、真实跨 rank SUM、完整样本参考梯度、空 rank、模态归约和 HP DTensor Shard bias 更新。
• 继续覆盖不等长/packed 样本、mask、原生模型/EP 接线、两类 checkpoint、梯度累积、双 optimizer、冻结 router、缺席模态及 checkpoint materialization。
• 门禁 FSDP 重构去除platform抽象 #596 的5条 W0640 已通过两份测试内显式绑定每轮变量修复；没有新增 pylint 屏蔽或放宽断言。使用 Jenkins 同款远端 pylintrc 与 Pylint 3.3.7 检查两个修改文件，10.00/10；定向回归14 passed、5 subtests passed。
• 门禁 [Feature]: add mega_kernel profiling #600 的三项 UT 失败已修复：两类 CPU checkpoint 测试显式选择 CPU 并在清理时恢复原默认设备；recipe wrapper UT 改为随测试模块发布的 YAML fixture，保留 text/VLM、production/validate、EP 合并及梯度断言。本次仅改三份测试，不修改生产逻辑或新增 pylint 屏蔽。重新回归 116 passed、86 subtests passed；三份文件使用门禁同款远端 pylintrc 检查 10.00/10；模拟 NPU 默认设备下11项测试通过、硬件能力查询0次且每次恢复默认值；仅 tests、不含 examples 的打包目录中原失败的3项用例通过。fixture 与当前两份源码 recipe 完全一致。真实 NPU 完整训练仍待远端门禁验证。
• gate-doctor AutoModels routed EP 公共执行逻辑与路由兼容性 #619 修复：rebase master 68fc3e33（保留其 graph compiler 集成）；auxloss eager Trainer fixture 补齐 graph_compiler=None，tied-weight UT 接收 helper 的第四个返回值。全部数值、共享权重和梯度断言保留。另清理该 tracer 测试两项现有未使用 import/接收变量，满足同款 Pylint 10.00/10，无新增屏蔽。定向回归 137 passed、86 subtests passed，修复前可复现3项失败。扩展编译回归为291 passed、2 failed；两个 PP 弱引用生命周期用例在本地原始 master 同样失败，且在 AutoModels routed EP 公共执行逻辑与路由兼容性 #619 远端通过，故未扩改。记录与 patch 已归档到本地 _codex-docs/2026-10-12-pr1067-gate-doctor/。本轮未验收 V4.1 joint-graph auxloss 的完整训练梯度。
• 原 master f96a3f08 快照按仓库 .pylintrc 的差异检查为新增诊断0、既有58条；该历史检查不代表当前全仓 lint 通过。最新两份测试按 Jenkins 同款规则为10.00/10，[Bug]: 初始化阶段会有两次权重加载过程 #621 的14项实际 Check 均已通过。

已知限制：评论中的 FSDP16+EP16 meta embedding 初始化失败仍未在本 PR 修复。普通多维 Shard 的 dim0 全局整数索引不受当前 DTensor 路径支持；用户要求将初始化方案放在独立特性分支，提交 17db1014，不并入本 PR。已集成 master 的 clamp grouped-GEMM，保留 use_grouped_gemm=true 支持；CPU 替身覆盖 clamp、dtype 和 forward/backward 数值，未验证实际 NPU kernel。没有声称本次本地验证完成 NPU/GPU 的完整 CP×TP×EP×FSDP 训练或性能验收。

语义边界：microbatch 内按样本均权，跨 microbatch 沿用 Trainer 的 foundation-loss 权重，未新增全局 sample-count denominator。last_aux_loss 是未乘系数的 detached 观测值，不应再次加入主 loss。

---

Self-checklist:

• 设计：报告依据、公式、两种 top-k 职责、并行归约与生命周期已记录。
• 测试：新增 UT/ST 随 PR 提交，非零 bias 回归覆盖 value/gradient/实际负载。
• 验证：记录本地 CPU 结果与明确的硬件/初始化验证边界。
• 资料：介绍、设计、代码走读、评论快照与日志归档到本地仓根 _codex-docs/2026-10-12-auxloss-rebase-master/；既有 aux 文档已同步更正，资料不随代码 PR 上传。

[Er1ckPang]

PR1067 验证结果与建议修复（基于 db45bb6b）

我在单机 16 张 Ascend 910 NPU 上验证了 PR1067 的精确 head db45bb6b。模型采用 4 层缩减版
DeepSeek-V4.1，不加载权重。原 PR 的相关 CPU 测试为 100 passed，Gloo 分布式测试为 1 passed，
关闭 FSDP/EP 的单卡训练 step 也能正常完成。不过，当前实现仍存在两个需要修复的问题。

1. 序列级辅助损失使用了错误的 Top-K 集合

问题

路由器使用模态 correction bias 选择实际专家，同时使用原始 score 计算专家权重，这部分实现是正确的：

indices = torch.topk(scores + correction_bias, self.top_k, dim=-1, sorted=False).indices
weights = scores.gather(1, indices)

但随后又把同一组经过 bias 修正的 indices 传入了序列级辅助损失：

aux_loss = sequence_load_balancing_loss(
    scores, indices, sequence_ids.reshape(-1), num_sequences, token_mask, sequence_partition_groups,
)

DeepSeek-V4.1 报告说明保留了小权重的 sequence-level balance
loss，但没有重新展开公式。其沿用的 DeepSeek-V3 公式明确使用
原始 affinity score s_i,t 的 TopK(s_i,t) 计算 f_i。Correction bias 只应用于实际专家选择，
不应改变辅助损失的 Top-K 统计。

当 correction bias 改变实际路由结果时，两者会产生明显差异。构造两个 token，原始 score 分别为
[0.9, 0.1]、[0.8, 0.2]，并使用 bias 把两个 token 都路由到 expert 1：

• 原始 score Top-K 公式：loss 为 1.70000005；
• PR1067 当前 bias-corrected Top-K：loss 为 0.30000001；
• 绝对差为 1.40000004；
• 两个专家对应的 score 梯度符号均相反。

根因

代码复用了 dispatch 所需的实际路由 indices。启用 auxiliary-loss-free correction bias 后，下面两个
概念必须分开：

• 实际 dispatch 和模态负载计数应使用 scores + correction_bias 的 Top-K；
• sequence auxiliary loss 应使用未加 bias 的原始 scores 的 Top-K。

建议修改

diff --git a/hyper_parallel/models/deepseek_v41/aux_loss.py b/hyper_parallel/models/deepseek_v41/aux_loss.py
@@
-        selected_experts: Actual bias-corrected routing decisions, shaped ``[T, K]``.
+        selected_experts: Top-K experts selected from unbiased scores, shaped ``[T, K]``.

diff --git a/hyper_parallel/models/deepseek_v41/modeling_deepseek_v41.py b/hyper_parallel/models/deepseek_v41/modeling_deepseek_v41.py
@@
+            unbiased_indices = torch.topk(scores, self.top_k, dim=-1, sorted=False).indices
             aux_loss = sequence_load_balancing_loss(
-                scores, indices, sequence_ids.reshape(-1), num_sequences, token_mask, sequence_partition_groups,
+                scores, unbiased_indices, sequence_ids.reshape(-1), num_sequences, token_mask,
+                sequence_partition_groups,
             )

后面的 accumulate_router_load(..., indices, ...) 应保持不变，因为模态 bias 更新需要统计真正接收了
token 的专家。

建议增加以下回归测试，确保 bias 可以改变实际路由，但不会改变辅助损失的 Top-K 定义：

def test_aux_loss_uses_unbiased_topk_while_routing_uses_bias(self) -> None:
    router = make_router()
    with torch.no_grad():
        router.weight.copy_(torch.tensor([
            [1.0, 1.0, 1.0, 1.0],
            [0.0, 0.0, 0.0, 0.0],
            [-1.0, -1.0, -1.0, -1.0],
        ]))
        router.bias.copy_(torch.tensor([0.0, 10.0, 0.0]))
    module = "hyper_parallel.models.deepseek_v41.modeling_deepseek_v41"
    with patch(f"{module}.sequence_load_balancing_loss", return_value=torch.tensor(0.0)) as aux_loss:
        _, _, routed_indices = router(torch.ones(1, 2, 4))
    torch.testing.assert_close(routed_indices, torch.ones(2, 1, dtype=torch.long))
    torch.testing.assert_close(aux_loss.call_args.args[1], torch.zeros(2, 1, dtype=torch.long))

2. PR 当前 head 在 FSDP16+EP16 下初始化失败

问题与复现结果

使用 16 卡 FSDP16+EP16 运行时，所有 rank 都会在开始训练前失败：

ValueError: For __getitem__, indexing with int on non-replicate dim 0 is not supported,
but got sharding fsdp_shard on dim 0.

失败调用路径为：

DeepseekV41ForCausalLM._init_weights
  -> _initialize_embedding_shard_safe
  -> module.weight[module.padding_idx].zero_()

根因

执行这段初始化逻辑时，embedding 的 dim0 已经被 FSDP 分片。此时直接通过全局 padding_idx 对
DTensor 做整数索引是不支持的。应根据 DTensor layout 找到 padding row 所属的 local shard，并且只在
对应 rank 上清零本地行。

这也是 PR1066 与 PR1067 两个兄弟分支的集成遗漏：PR1066 已包含 shard-safe helper，但 PR1067 直接
基于 master 开发，因此没有带入该修复。

建议修改

diff --git a/hyper_parallel/models/deepseek_v41/modeling_deepseek_v41.py b/hyper_parallel/models/deepseek_v41/modeling_deepseek_v41.py
@@
+from hyper_parallel.core.dtensor.layout import infer_slice_area_by_layout
@@
 def _initialize_embedding_shard_safe(module: nn.Embedding, std: float) -> None:
-    """Initialize an embedding and clear its global padding row through DTensor."""
+    """Initialize an embedding and clear its global padding row shard-safely."""
     nn.init.normal_(module.weight, mean=0.0, std=std)
     if module.padding_idx is None:
         return
-    module.weight[module.padding_idx].zero_()
+    weight = module.weight
+    layout = getattr(weight, "layout", None)
+    to_local = getattr(weight, "to_local", None)
+    if layout is None or not callable(to_local):
+        weight[module.padding_idx].zero_()
+        return
+    inner_rank = layout.rank_list.index(layout.mesh.rank)
+    slice_area = infer_slice_area_by_layout(layout, inner_rank, weight.shape)
+    row_start, row_end = slice_area[0]
+    if row_start <= module.padding_idx < row_end:
+        to_local()[module.padding_idx - row_start].zero_()

3. 与 PR1066/grouped-gemm 的集成问题

PR1067 当前 head 会主动拒绝 use_grouped_gemm=true：

ValueError: DeepSeek-V4.1 clamp semantics currently require use_grouped_gemm=false

考虑到 PR1066 和 PR1067 是两个兄弟分支，这一限制可以理解；但这也意味着仅基于 PR1067 无法完成
FSDP+EP × grouped-gemm on/off 的组合验收。合入两个 PR 时，需要处理 MoE adapter 和模型文件上的重叠
修改，并移除旧的 grouped-gemm 拒绝逻辑，而不是在 PR1067 中复制另一份 clamp 实现。

PR1066 的 grouped-gemm 实现已单独对比公开的 V4.1 推理公式：forward、输入梯度、gate/up 参数梯度、
down 参数梯度均逐位一致。

应用上述两个修复后的验证结果

• 序列/模态负载均衡定向测试：15 passed；
• FSDP16+EP16、grouped-gemm off、1 个训练 step：通过；
• loss：11.9285；
• grad norm：19.1775；
• 处理 token：8192；
• 吞吐：1879.57 tokens/s；
• 关闭 FSDP/EP 的单卡 smoke 仍然通过。

PR1067 的模态 bias 主流程本身与 V4.1 报告一致，包括：文本/图像独立 bias、bias 只影响选择、原始
score 用于输出权重、每个 optimizer step 后独立更新、更新速度 0.001、序列辅助损失系数
0.0001。建议修改范围可以限制在辅助损失 Top-K 定义和 FSDP-safe 初始化/兄弟 PR 集成路径。

[Er1ckPang]

/check-pr

[Er1ckPang]

/retest

[mindspore-ai-bot]

🔵 The pipeline #581 is running. Please wait a moment... (Link)

[mindspore-ai-bot]

Project Name	Build_Stage	Build Result	Details
Hyper-parallel_Github_Gate	-	❌ FAILURE	581
	Check_Pylint	FAILURE	581
	Check_AI_Keyword	FAILURE	581
	Smoke_Ascend910b_Torch	FAILURE	6176

[mindspore-ai-bot]

@Er1ckPang, your gate pipeline failed at the code check stage. You can configure some code check tools locally following the (pre-push quick guide), to enable automatic checking when you push your code every time. Tools included are: Check_CMakelint, Check_Codespell, Check_Cpplint, Check_Lizard, Check_Pylint, Check_ClangFormat, Check_ShellCheck.

[Er1ckPang]

/retest

[mindspore-ai-bot]

🔵 The pipeline #590 is running. Please wait a moment... (Link)

[mindspore-ai-bot]

Project Name	Build_Stage	Build Result	Details
Hyper-parallel_Github_Gate	-	❌ FAILURE	590
	Check_Pylint	FAILURE	581
	Check_AI_Keyword	FAILURE	581
	Smoke_Ascend910b_Torch	FAILURE	6176

[mindspore-ai-bot]

@Er1ckPang, your gate pipeline failed at the code check stage. You can configure some code check tools locally following the (pre-push quick guide), to enable automatic checking when you push your code every time. Tools included are: Check_CMakelint, Check_Codespell, Check_Cpplint, Check_Lizard, Check_Pylint, Check_ClangFormat, Check_ShellCheck.

[Er1ckPang]

/retest

[mindspore-ai-bot]

🔵 The pipeline #606 is running. Please wait a moment... (Link)

[mindspore-ai-bot]

Project Name	Build_Stage	Build Result	Details(Full/Incremental)	Detail
Coverage	Line Coverage	❌ FAILURE	66%/94%	Link

[mindspore-ai-bot]

Platform Type	Module	Excution Time	Time Limit
_torch	-	-	-
-	clip_grad	186.782s	120.0s
-	pipeline_parallel	622.917s	600.0s
-	optimizer	225.272s	300.0s
-	fully_shard	277.695s	600.0s
-	shard	150.153s	300.0s
-	swap_optimizer	72.051s	120.0s
-	tensor_parallel	191.316s	300.0s
-	accuracy	72.263s	120.0s
-	activation_memory	113.02862s	120.0s
-	distributed_checkpoint	63.667s	120.0s
-	context_parallel	72.652s	120.0s
-	process_group	61.121s	90.0s
-	dtensor	70.112s	120.0s
-	init_weights	66.343s	118.75s
-	expert_parallel	62.101s	120.0s
-	llama3_examples	60.813s	120.0s
-	multicore	70.56425s	120.0s

[mindspore-ai-bot]

Project Name	Build_Stage	Build Result	Details
Hyper-parallel_Github_Gate	-	❌ FAILURE	606
	Check_Pylint	FAILURE	581
	Check_AI_Keyword	FAILURE	581
	Check_Pylint	FAILURE	596
	UT_TEST	FAILURE	600
	Smoke_Ascend910b_Torch	FAILURE	6176
	coverage	FAILURE	600

[mindspore-ai-bot]

@Er1ckPang, your gate pipeline failed at the code check stage. You can configure some code check tools locally following the (pre-push quick guide), to enable automatic checking when you push your code every time. Tools included are: Check_CMakelint, Check_Codespell, Check_Cpplint, Check_Lizard, Check_Pylint, Check_ClangFormat, Check_ShellCheck.

[Er1ckPang]

/check-pr

[Er1ckPang]

/retest

[mindspore-ai-bot]

🔵 The pipeline #619 is running. Please wait a moment... (Link)

[Er1ckPang]

/check-pr

[Er1ckPang]

/retest

[mindspore-ai-bot]

🔵 The pipeline #621 is running. Please wait a moment... (Link)

[Er1ckPang]

/check-pr

[mindspore-ai-bot]

Project Name	Build_Stage	Build Result	Details(Full/Incremental)	Detail
Coverage	Line Coverage	✅ SUCCESS	66%/94%	Link

[mindspore-ai-bot]

Platform Type	Module	Excution Time	Time Limit
_torch	-	-	-
-	activation_memory	127.85225s	120.0s
-	shard	170.089s	300.0s
-	optimizer	252.448s	300.0s
-	pipeline_parallel	564.435s	600.0s
-	swap_optimizer	83.014s	120.0s
-	fully_shard	293.604s	600.0s
-	accuracy	76.233s	120.0s
-	clip_grad	83.559s	120.0s
-	expert_parallel	80.33s	120.0s
-	multicore	74.92775s	120.0s
-	context_parallel	80.15s	120.0s
-	dtensor	77.443s	120.0s
-	distributed_checkpoint	70.868s	120.0s
-	tensor_parallel	231.353s	300.0s
-	init_weights	74.225s	118.75s
-	process_group	64.313s	90.0s
-	llama3_examples	63.388s	120.0s

[mindspore-ai-bot]

Project Name	Build_Stage	Build Result	Details
Hyper-parallel_Github_Gate	-	✅ SUCCESS	621