Repository navigation
Add an opt-in route_norm flag for the legacy MoE gate - #8769
Aravind-11 wants to merge 1 commit into
Conversation
top1gating leaves the chosen expert's softmax probability in the combine weight. top2gating and topkgating rescale the kept weights so they sum to 1. Making those agree changes the MoE output for existing checkpoints, so route_norm defaults to None and each gate keeps its historical rule. True or False overrides all three gates. Auxiliary loss is unchanged. Signed-off-by: ARAVINDHAN T <arvindhant01@gmail.com>
|
This looks good to me at a glance — thanks for putting it together. |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 79df454fae
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| enable_expert_tensor_parallelism: bool = False, | ||
| top2_2nd_expert_sampling: bool = True) -> None: | ||
| top2_2nd_expert_sampling: bool = True, | ||
| route_norm: Optional[bool] = None) -> None: |
There was a problem hiding this comment.
Because this SHA has one parent, it is a non-merge commit, but git log --format='%b' b11e0a1924c934183ecfbe01e225c81a04a123af shows that its message has no Signed-off-by trailer. Recreate the commit with --signoff using the configured identity so that it satisfies the repository's commit/CI policy.
AGENTS.md reference: AGENTS.md:L6-L9
Useful? React with 👍 / 👎.
This branch has not been deployed
2 participants
Draft. Please say whether this shape is welcome before review-for-merge. The measurement and the questions for maintainers are in #8768.
What
Legacy
MoE/TopKGategain an opt-inroute_norm: Optional[bool] = None.Nonekeeps the historical combine weights.top1gatingleaves the chosen expert's softmax probability.top2gatingandtopkgatingrescale the kept weights so they sum to 1.Truerescales for every k. A kept top-1 weight becomes 1.Falseleaves the raw softmax probabilities for every k.MOELayer) and the dense return (direct callers and inference).0 / epsis 0.deepspeed/ops/transformer/inference/moe_inference.pystill constructsTopKGatepositionally and stays on the historical rule.This is not
expert_parallel.route_norm. That key configures the AutoEP router, which already has its own flag.Why the default is not one bool
On
991ccf0f, the same logits give top-1 combine weights[0.710100, 0.870166, 0.378175, 0.799164]andtopkgating(k=1)weights of all ones. The numbers and the setup are in #8768. Either unifying choice changes the MoE output of existing checkpoints.Noneis the default that preserves both rules. Router parameter tensors are unchanged; the layer output and the expert gradients change when the flag is set.Tests
CPU, torch 2.5.1:
10 passed, 37 deselected. The six new tests check:
None, and passing the historical bool are bitwise equal, for k=1, k=2, and k=3, on the sparse and dense returnsTrue/Falsechanges only the combine weights; expert indices, capacity, the dispatch mask, andl_auxstay equalroute_norm=Truemakes a kept weight 1, including a dropped token staying 0topkgating(k=1)l_auxmatch, top-2 and top-kl_auxdo not, androute_normdoes not move either lossMoE()andMoE(route_norm=None)match after a state-dict copy;MoE(k=1, route_norm=True)changes the output and notl_aux;MoE(k=2, route_norm=True)matches the default androute_norm=Falsedoes notTopKGateconstructor leavesroute_normunsetThe distributed classes in that file were not run.
yapf(column limit 119) andflake8are clean on the edited Python files.Not in this diff
used_tokenfork != 1, and the top-1 no-drop tensor-parallel capacity clamp. That is question 5 in [REQUEST] Legacy MoE gate: opt-in route_norm, with historical combine weights left as the default #8768.@tohtana is the code owner for
deepspeed/moe/. The questions are in #8768. I will change the flag name or the default if you want a different shape, and I will close this if you would rather not take it.