Add HybridModel MTP runtime controls

[Phlip79]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Adds a backward-compatible, per-forward run_mtp_forward control to HybridModel. Target user is VeRL for Nemotron 3 models that use MTP.

HybridModel currently executes its training/evaluation MTP forward whenever MTP layers are present and the model is outside InferenceMode. Consumers may need to load, save, synchronize, or export MTP parameters without executing the MTP forward or attaching its auxiliary loss on every model call.

This is required for native HybridModel integration in veRL. veRL distinguishes between loading MTP parameters, training them, and using them for speculative rollout. Its current workaround monkeypatches GPTModel, which does not apply to HybridModel and should not become the long-term integration path as GPTModel is deprecated.

Behavior

run_mtp_forward=True preserves the existing behavior and remains the default.

When run_mtp_forward=False, HybridModel.forward:

• skips the training/evaluation MTP forward;
• skips process_mtp_loss and therefore does not attach the MTP auxiliary loss;
• leaves the MTP module, parameters, and checkpoint/state-dict structure intact;
• returns the ordinary base-model output at both post-processing and intermediate pipeline stages;
• remains safe on pipeline ranks that do not own an MTP module; and
• does not disable serial MTP speculative decoding, which remains controlled by InferenceMode, the inference context, and the speculative decoding controller.

The control is call-time rather than construction-time because one model instance may alternate between forward-only operations and training updates while preserving the same loaded MTP parameters.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[Phlip79]

/ok to test b722c55

[Phlip79]

/ok to test 582f111

[Phlip79]

/claude strict-review

[claude]

Strict Review Summary

Findings — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1

Assessment: Low risk. Recommend merge after considering the one naming nit.

This is a small, well-scoped, backward-compatible change adding a keyword-only compute_mtp_loss: bool = True control to HybridModel.forward.

Correctness verified:

• The default path (compute_mtp_loss=True) is provably behavior-identical to the original:
  • mtp_forward_ran gains an and compute_mtp_loss term that is True by default → identical to the prior expression.
  • The changed elif not in_inference_mode: → elif mtp_forward_ran: at L579 is exactly equivalent in context: it sits under if ... self.mtp_process (L563) and in the elif of if is_spec_decode, so mtp_forward_ran there reduces to not in_inference_mode.
• The disabled path (compute_mtp_loss=False) cleanly skips both the MTP forward and process_mtp_loss while leaving the MTP module, parameters, and checkpoint/state-dict structure untouched.
• Speculative-decoding / inference paths are correctly unaffected — the spec-decode branch (L565-578) never depended on mtp_forward_ran, and compute_mtp_loss only gates the training/eval block. The dedicated test test_forward_mtp_loss_control_does_not_disable_speculative_decoding confirms this.
• Non-MTP pipeline ranks are safe (guarded by self.mtp_process), and the intermediate-stage return (not self.post_process) is covered by a test.

Backward compatibility: New arg is keyword-only (after *) with a safe default, so no existing positional callers break, and no checkpoint/config/default-hyperparameter changes are introduced.

Tests: Good coverage — parametrized enabled/disabled/default, non-MTP rank, pre-post-process boundary, and spec-decode. import types and InferenceMode are properly imported.

Process-group guidance: No new parallel_state.get_*_group() reads introduced.

The single SUGGESTION is a non-blocking naming observation on compute_mtp_loss (the flag also gates the MTP forward pass, not just the loss).

[Phlip79]

/ok to test cf3b1d9

[Phlip79]

/claude strict-review

[claude]

Strict review passed — no significant issues found. LGTM

Verified:

• Backward compatibility: run_mtp_forward is a keyword-only arg (after *) defaulting to True; no existing caller passes it, so existing behavior is preserved.
• Behavioral equivalence for the default: at the process_mtp_loss guard site self.mtp_process is already True and is_spec_decode is False, so with run_mtp_forward=True, mtp_forward_ran reduces to not in_inference_mode — exactly the prior elif not in_inference_mode condition. The refactor also correctly couples process_mtp_loss to whether the MTP forward actually ran.
• Spec decode unaffected: is_spec_decode handling and mtp_decoder_hidden_states population remain gated only by inference context/controller state, not the new flag.
• No new global process-group reads introduced in megatron/core.
• Test coverage exercises defaults, explicit enable/disable, non-MTP pipeline ranks, the pre-post-process return path, and non-interference with speculative decoding.