已合并
fix: use deterministic placeholder initialization for DeepSeek4 hash layers #4898
丁子叉创建于 17 天前
fix: use deterministic placeholder initialization for DeepSeek4 hash layers #4898
已合并
Pull Request已成功合入, 合并人@ascend-robot
（感谢 丁子叉 的贡献）
atomgit-bot
17 天前 评论：
17 天前 评论：
变更摘要
此 PR 主要对 mindspeed_llm/core/transformer/moe/router.py 进行代码风格规范化与少量功能修复。核心功能变更是修改了 Hash 层 tid2eid 的初始化方式，从 torch.randperm 随机排列改为轮询（round‑robin）占位初始化；同时移除了未使用的变量 topk_group_idx_for_aux_loss 和无用导入 partial，并为多个函数添加了 pylint 抑制注释。其余大部分改动为代码格式化，包括长行拆分、缩进调整、末尾逗号补齐和空白行规范化。
主要改动
- Hash 层
tid2eid初始化方式变更: 将topk_router_build_hash_module中self.tid2eid的初始化从torch.randperm随机生成改为基于(token_ids + offset) % num_experts的轮询式初始化，并注明该实现仅为使 Hash 层可从零运行而设的占位方案，不适用于真实训练。 - 移除未使用的变量
topk_group_idx_for_aux_loss: 在group_limited_greedy_topKgating函数中删除了计算后未被引用的topk_group_idx_for_aux_loss，消除无效计算。 - 清理无用导入: 移除了
from functools import partial，该导入在文件中未被使用。 - 新增 pylint 抑制注释: 为
topk_router_forward_patch、topk_router_gating_func、global_aux_loss_topk_router_forward添加了# pylint: disable=redefined-builtin，为global_load_balancing_loss_func添加了# pylint: disable=possibly-used-before-assignment，以抑制静态分析误报。 - 大规模代码格式化: 对整个文件进行了统一的风格调整，包括长表达式拆分换行、函数参数缩进对齐、末尾逗号补全、多余空行清理以及注释格式规范化（如
# TODO改为完整句子描述）。
ascend-robot
17 天前 评论：
17 天前 评论：
atomgit-bot
17 天前 评论：
17 天前 评论：
17 天前 添加了label：ascend-cla/yes
此处折叠了123条消息 查看更多
6 天前 添加了label：approvedlgtm
6 天前 合入了pull request
ascend-robot
6 天前 评论：
6 天前 评论：
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
ascend-robot
6 天前 评论：
6 天前 评论：
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
ascend-robot
6 天前 评论：
6 天前 评论：
The MR is merging by another one
If you want to solve this problem, you can click here to do it in the FAQs.
What this PR does / why we need it?
This PR updates the initialization logic of the DeepSeek4 hash-layer
tid2eidtable.The previous implementation generated one
torch.randperm(moe_router_topk)for each token. This only permuted expert IDs within[0, moe_router_topk)and consumed a large amount of RNG state during model construction.Following NVIDIA/Megatron-LM#4481, this PR:
tid2eidas a non-trainable parameter;int64dtype required by the existingtorch.gatherandtorch.scatterrouting path;DeepSeek4 checkpoints provide a pre-trained
tid2eidtable. No public recipe for initializing this table for real training is currently available, so the placeholder initialization in this PR must not be treated as a real training initialization strategy.Reference:
https://github.com/NVIDIA/Megatron-LM/pull/4481
The remaining changes in
router.pyare formatting and static-check fixes produced while making the file pass the repository pre-commit checks. They do not change the existing routing behavior.Does this PR introduce any user-facing change?
No CLI or configuration interface is changed.
For models loaded from a valid DeepSeek4 checkpoint, the checkpoint-provided
tid2eidtable overwrites the placeholder value, so checkpoint loading behavior is unchanged.For hash-layer models constructed without a checkpoint, the initial routing table changes from random permutations limited to the first
moe_router_topkexperts to deterministic round-robin assignments across all experts.The placeholder initialization is not suitable for real-world training.
How was this patch tested?
mindspeed_llm/core/transformer/moe/router.pytid2eidshape and deterministic round-robin values;tid2eidremains anint64non-trainable parameter;torch.gatherandtorch.scatterrouting path;tid2eid.