[fix] bypass r3 for mtp layer.

[guapisolo]

Related megatron commit: radixark/Megatron-LM#10

R3 cannot be applied to mtp layer, since mtp will not forward verify bonus tokens during rollouts. So bypass r3 replay for mtp head.

This impl is a little bit hard code, but maybe the best impl in current design. Mainly because we cannot get the module name before the layer is fully registered, so the Replay class cannot record anything about the layer name. Also, the original MoE routing function does not included module info, so we cannot directly skip MoE mtp replay during runtime... So add the element is_mtp as instead.

[gemini-code-assist]

Summary of Changes

Hello @guapisolo, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a bug concerning the interaction between Multi-Token Prediction (MTP) layers and Rollout Routing Replay (R3) by ensuring that MTP layers do not utilize the R3 mechanism. The changes involve introducing a specific flag to identify MTP layers within the Mixture of Experts (MoE) routing logic, allowing for conditional bypassing of R3. Additionally, a new script for the GLM-4.7-Flash model demonstrates the application of these features, including advanced speculative decoding techniques.

Highlights

• MTP and Rollout Routing Replay Compatibility: Introduced an is_mtp flag across MoE and router components to explicitly mark Multi-Token Prediction (MTP) layers. This flag prevents MTP layers from engaging with the Rollout Routing Replay (R3) mechanism, addressing a reported bug.
• Conditional Rollout Routing Replay: Modified the topk_routing_with_score_function to conditionally apply the routing_replay_manager based on the is_mtp flag, ensuring R3 is skipped for MTP layers.
• Model Provider Integration: Implemented logic in the model provider to temporarily disable the routing_replay_manager when constructing MTP block specifications, further reinforcing the separation of MTP from R3.
• New GLM-4.7-Flash Script: Added a new script, scripts/run_glm47_flash.py, which configures and runs the GLM-4.7-Flash model, including settings for SGLang speculative decoding (EAGLE) and explicit enablement of the miles router with rollout routing replay.

Changelog

• docker/patch/dev/megatron.patch
  • Added a set_is_mtp method to BaseMoELayer in megatron/core/transformer/moe/moe_layer.py.
  • Modified topk_routing_with_score_function in megatron/core/transformer/moe/moe_utils.py to include an is_mtp parameter and conditionally apply routing_replay_manager.get_topk_fn.
  • Introduced an is_mtp attribute to the Router class in megatron/core/transformer/moe/router.py and passed it to the routing function.
  • Integrated a call to set_is_mtp on the MLP layer within MultiTokenPredictionLayer in megatron/core/transformer/multi_token_prediction.py.
• miles/backends/megatron_utils/model_provider.py
  • Imported logging and routing_replay_manager.
  • Implemented logic to temporarily disable routing_replay_manager when building MTP block specifications if use_rollout_routing_replay is enabled, and re-enable it afterwards.
• scripts/run_glm47_flash.py
  • Added a new script for running the GLM-4.7-Flash model.
  • Configured various arguments for checkpointing, rollout, evaluation, performance, GRPO, optimizer, and SGLang (EAGLE speculative decoding, MTP).
  • Explicitly enabled --use-miles-router and --use-rollout-routing-replay in the script's arguments.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request aims to fix a bug where Multi-Token Prediction (MTP) layers were incorrectly interacting with the rollout routing replay (r3) feature, by introducing an is_mtp flag in MoE layers. However, a critical security vulnerability has been identified: the new script scripts/run_glm47_flash.py introduces several command injection vulnerabilities due to unsafe interpolation of command-line arguments. Additionally, there are architectural issues in megatron-core with a direct dependency on the miles project, violating dependency inversion principles, and a robustness issue in miles where a global flag is modified without a try...finally block, which could leave the system in an inconsistent state.

I am having trouble creating individual review comments. Click here to see my feedback.

scripts/run_glm47_flash.py (26)

The script scripts/run_glm47_flash.py is vulnerable to command injection. Several command-line arguments (args.model_dir, args.data_dir, args.model_org, args.model_name) are directly interpolated into shell command strings that are then executed via U.exec_command. Since these arguments are not sanitized or escaped, an attacker who can control these arguments can execute arbitrary shell commands on the system.

For example, if args.model_dir is set to ; rm -rf / ;, the command executed will be mkdir -p ; rm -rf / ; /root/datasets.

To remediate this, use shlex.quote() to escape command-line arguments before interpolating them into shell command strings, or better yet, use subprocess.run() with a list of arguments instead of a shell string.

scripts/run_glm47_flash.py (31-34)

Similar to the issue on line 26, the args.model_org, args.model_name, and args.model_dir arguments are interpolated into a shell command string without sanitization, leading to a command injection vulnerability.

scripts/run_glm47_flash.py (165-171)

The args.extra_args argument is directly interpolated into train_args, which is then executed as a shell command in U.execute_train. This allows for arbitrary command injection if the user provides malicious input in extra_args.

docker/patch/dev/megatron.patch (545-549)

Importing from miles within megatron-core creates a problematic dependency inversion. The core library (megatron-core) should not depend on an application-level project (miles) that uses it. This makes the core library less reusable and harder to maintain.

Consider refactoring this to use a plugin or hook system, where miles can register its custom logic with megatron-core at runtime without megatron-core having a hard dependency on miles. For example, megatron-core could expose a registry for top-k functions that miles can add to.

miles/backends/megatron_utils/model_provider.py (196-204)

The temporary disabling of routing_replay_manager.enabled is not robust. If get_gpt_mtp_block_spec or subsequent code raises an exception, routing_replay_manager.enabled will not be reset to True, leaving it in an inconsistent state for the rest of the process. This should be wrapped in a try...finally block to ensure the flag is always restored.

            try:
                if getattr(args, "use_rollout_routing_replay", False):
                    routing_replay_manager.enabled = False
                    logger.warning(
                        "Rollout routing replay is not applicable for MTP modules, so skipped replay registration"
                    )
                mtp_block_spec = get_gpt_mtp_block_spec(config, transformer_layer_spec, **mtp_kwargs)
                kwargs["mtp_block_spec"] = mtp_block_spec
            finally:
                if getattr(args, "use_rollout_routing_replay", False):
                    routing_replay_manager.enabled = True

[yueming-yuan]

could you also create a patch-equivalent PR in megatron fork miles-main branch? thanks! https://github.com/radixark/Megatron-LM/tree/miles-main

[guapisolo]

could you also create a patch-equivalent PR in megatron fork miles-main branch? thanks! https://github.com/radixark/Megatron-LM/tree/miles-main

done