feat: correct routed_experts from sglang(trick of routing_replay)#884
feat: correct routed_experts from sglang(trick of routing_replay)#884rchardx merged 31 commits intoinclusionAI:mainfrom
Conversation
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as resolved.
This comment was marked as resolved.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
ZiyiTsang
force-pushed
the
r3replay
branch
2 times, most recently
from
24fc9c1 to
7adde36
Compare
ZiyiTsang
changed the title
ZiyiTsang
changed the title
|
(Moved to inline comment on |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces support for retrieving routed_experts information from the sglang backend, which is useful for analyzing Mixture of Experts (MoE) models. The changes span across configuration, data structures, and engine logic to propagate the option and handle the returned data. The implementation looks good, with proper error handling for unsupported backends and documentation updates. I have a couple of suggestions to improve maintainability by adding a clarifying comment and removing a redundant validation check.
| num_sgl_token = ( | ||
| meta_info["prompt_tokens"] + meta_info["completion_tokens"] - 1 | ||
| ) |
There was a problem hiding this comment.
For better maintainability, could you add a brief comment explaining why 1 is subtracted in the num_sgl_token calculation? This would clarify the logic for future developers who might not be familiar with the specifics of the sglang API's token counting.
| if self.config.rollout.return_routed_experts: | ||
| raise ValueError( | ||
| "return_routed_experts is not supported with vLLM backend. Please disable return_routed_experts or switch to SGLang backend." | ||
| ) |
There was a problem hiding this comment.
This validation check is redundant because the _validate_cfg method, which is called earlier in __init__, already performs the same check. To avoid code duplication, this if block can be removed.
…nclusionAI#884) The routing replay becomes a trick for stabilizing the RL training in MOE. The sglang 0.5.7 has support to return routed_experts in API. Therefore this P.R support to get the routed_experts from sglang, which will be put into Megatron later. --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Wentai Zhang <zhangwentai.zwt@antgroup.com>
…nclusionAI#884) The routing replay becomes a trick for stabilizing the RL training in MOE. The sglang 0.5.7 has support to return routed_experts in API. Therefore this P.R support to get the routed_experts from sglang, which will be put into Megatron later. --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Wentai Zhang <zhangwentai.zwt@antgroup.com>
2 participants
Description
The routing replay proposed in this Paper becomes a trick for stabilizing the RL training in MOE.
The sglang 0.5.7 has support to return
routed_expertsin API.therefore this P.R support to get the routed_experts from sglang, which will be put into Megatron later.
To test with:
python3 -m areal.launcher.local examples/math/gsm8k_rl.py --config examples/math/gsm8k_grpo.yaml gconfig.return_routed_experts=TrueRelated Issue
N/A
Type of Change
work as expected)
Checklist
jb build docs/gemini review)Breaking Change Details (if applicable):
Additional Context
Need help? Check the Contributing Guide or ask in
GitHub Discussions!