已合并
fix: use deterministic placeholder initialization for DeepSeek4 hash layers #4898
丁子叉创建于 16 天前
fix: use deterministic placeholder initialization for DeepSeek4 hash layers #4898
已合并
丁
丁子叉成员
16 天前
16 天前
## What this PR does / why we need it?
This PR updates the initialization logic of the DeepSeek4 hash-layer `tid2eid` table.
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:
- replaces the random initialization with deterministic round-robin initialization across all experts;
- keeps `tid2eid` as a non-trainable parameter;
- preserves the `int64` dtype required by the existing `torch.gather` and `torch.scatter` routing path;
- documents that this initialization is only a placeholder for constructing a runnable model from scratch.
DeepSeek4 checkpoints provide a pre-trained `tid2eid` table. 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.py` are 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 `tid2eid` table 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_topk` experts to deterministic round-robin assignments across all experts.
The placeholder initialization is not suitable for real-world training.
## How was this patch tested?
- Ran the complete pre-commit checks for:
- `mindspeed_llm/core/transformer/moe/router.py`
- Verified:
- `tid2eid` shape and deterministic round-robin values;
- expert indices remain within the configured expert range;
- `tid2eid` remains an `int64` non-trainable parameter;
- compatibility with the existing `torch.gather` and `torch.scatter` routing path;
- non-hash layers do not create `tid2eid`.
- Ran the DeepSeek4 checkpoint-loading ST case for 15 iterations. Training completed successfully and the checkpoint loaded without state-dict errors.
- The existing DeepSeek4 ST baseline JSON is intentionally not modified by this PR.
Pull Request已成功合入, 合并人@ascend-robot
（感谢 丁子叉 的贡献）
atomgit-bot
16 天前 评论：
16 天前 评论：
变更摘要
此 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
16 天前 评论：
16 天前 评论：
atomgit-bot
16 天前 评论：
16 天前 评论：
16 天前 添加了label：ascend-cla/yes
此处折叠了123条消息 查看更多
4 天前 添加了label：approvedlgtm
4 天前 合入了pull request
ascend-robot
4 天前 评论：
4 天前 评论：
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
ascend-robot
4 天前 评论：
4 天前 评论：
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
ascend-robot
4 天前 评论：
4 天前 评论：
The MR is merging by another one
If you want to solve this problem, you can click here to do it in the FAQs.