[megatron] feat: model engine support mtp

[ArronHZG]

What does this PR do?

model engine support mtp

#5323 break use mtp in mbridge, revert.

Unload the KV cache before parameter synchronization (SGLang supports this first).

Throughput increased from an initial 3900 token/s to 4800 token/s, representing a 23% improvement.
The speculative acceptance rate increased from 44% to 54%, representing a 22% improvement.

Checklist Before Starting

• Search for similar PRs. Paste at least one query link here: ...
• Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)
  • {modules} include fsdp, megatron, veomni, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data, cfg, reward, fully_async, one_step_off
  • If this PR involves multiple modules, separate them with , like [megatron, fsdp, doc]
  • {type} is in feat, fix, refactor, chore, test
  • If this PR breaks any API (CLI arguments, config, function signature, etc.), add [BREAKING] to the beginning of the title.
  • Example: [BREAKING][fsdp, megatron] feat: dynamic batching

Test

For changes that can not be tested by CI (e.g., algorithm implementation, new model support), validate by experiment(s) and show results like training curve plots, evaluation results, etc.

API and Usage Example

Demonstrate how the API changes if any, and provide usage example(s) if possible.

# Add code snippet or script demonstrating how to use this

Design & Code Changes

Demonstrate the high-level design if this PR is complex, and list the specific changes.

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

• Read the Contribute Guide.
• Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
• Add / Update the documentation.
• Add unit or end-to-end test(s) to the CI workflow to cover all the code. If not feasible, explain why: ...
• Once your PR is ready for CI, send a message in the ci-request channel in the verl Slack workspace. (If not accessible, please try the Feishu group (飞书群).)
• If your PR is related to the recipe submodule, please also update the reference to the submodule commit via git submodule update --remote or cd recipe && git pull origin main.

[gemini-code-assist]

Code Review

This pull request introduces support for Multi-Token Prediction (MTP) in the model engine, primarily affecting Megatron-based training and inference. Key changes include refactoring model forward passes to handle MTP-specific preprocessing and postprocessing, updating configuration files to enable MTP, and adding a new example script for MTP training. The changes also include improvements in handling nested tensors and position IDs. However, there are several areas that require attention to improve robustness, maintainability, and portability, particularly concerning hardcoded values, manual configuration steps, and potential behavioral changes in patched functions.

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

examples/mtp_trainer/test_dapo_mimo_7b_with_mtp_math_megatron_4_4.sh (34)

The comment indicates a manual step to modify max_position_embeddings in config.json. Manual intervention for configuration is prone to human error and reduces the reproducibility and automation of the setup. This step should ideally be automated within the script or handled programmatically by the model loading logic.

examples/mtp_trainer/test_dapo_mimo_7b_with_mtp_math_megatron_4_4.sh (30-31)

The default values for NNODES and NGPUS_PER_NODE are set to 16 and 8 respectively, which are very high for a test script. These are then overridden to 1 and 4 within the fully_async array (lines 77-80). This inconsistency can lead to confusion regarding the actual resource allocation and might cause unexpected resource consumption or errors if the script is run without careful inspection. It is best to define these variables once or clearly indicate which values take precedence.

NNODES=${NNODES:-1}
NGPUS_PER_NODE=${NGPUS_PER_NODE:-4}

examples/mtp_trainer/test_dapo_mimo_7b_with_mtp_math_megatron_4_4.sh (48-49)

The use of magic numbers 2 and 3 in the calculations for actor_ppo_max_token_len and infer_ppo_max_token_len reduces the readability and maintainability of the script. It would be clearer to define these values as named variables with descriptive names, explaining their purpose.

ACTOR_PPO_FACTOR=2
INFER_PPO_FACTOR=3
actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * ACTOR_PPO_FACTOR))
infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * INFER_PPO_FACTOR))

examples/mtp_trainer/test_dapo_mimo_7b_with_mtp_math_megatron_4_4.sh (76)

The calculation 512*100 for rollout.total_rollout_steps is hardcoded. While simple, it would be more explicit and easier to modify if 512 and 100 were defined as variables, or if the resulting value was directly assigned as a literal if it's a fixed constant.

ROLLOUT_BASE_STEPS=512
ROLLOUT_MULTIPLIER=100
rollout.total_rollout_steps=$(((ROLLOUT_BASE_STEPS*ROLLOUT_MULTIPLIER)))

verl/experimental/fully_async_policy/shell/runtime_env_4_4.yaml (14-15)

The paths for RAY_DATA_HOME and TENSORBOARD_DIR are hardcoded to user-specific directories (/home/hadoop-djst-algoplat). This significantly reduces the portability and reusability of this runtime environment configuration across different systems or users. These paths should be made configurable (e.g., via environment variables that can be set externally) or use more generic relative paths.

verl/models/mcore/model_forward.py (191-192)

In the _convert_to_nested_tensor function, if vi.shape[0] < target_len, the tensor is padded with torch.ones. Depending on the context and the data represented by vi (e.g., token IDs), padding with 1 might be semantically incorrect if 1 is a valid token ID. This could lead to unintended model behavior. It would be safer to pad with a specific pad_token_id from the tokenizer or a value that is guaranteed not to interfere with valid data.

            vi = torch.cat([vi, torch.full((target_len - vi.shape[0],), self.pad_token_id, dtype=vi.dtype, device=vi.device)])

verl/models/mcore/mtp_patch.py (78-81)

The refactoring of _megatron_gptmodel_postprocess removes the explicit delegation to self._postprocess_backup for inference paths (when labels is None). While the new logic might cover the training path, it's crucial to verify that the inference behavior remains unchanged. If _postprocess_backup contained specific logic or optimizations for inference that are not replicated in the new combined logic, this change could introduce regressions or alter the model's behavior during evaluation.