Support isolated MTP loss

[Victarry]

Summary

• Add mtp_isolated_loss to isolate MTP auxiliary loss gradients from the main decoder, shared embeddings, and output layer weights.
• Detach the MTP input path and output-layer path when mtp_isolated_loss=True.
• Add focused unit coverage for decoder-hidden-state and output-layer isolation on the main branch code path.

This is the main branch counterpart of #5080. The explicit return_logits API path is not included; online RL/MTP training should use the #3457 scheme (labels=None returns main logits, while MTP auxiliary labels are derived from input_ids in the MTP loss path).

RL / Online MTP Background

Online RL training frameworks such as veRL usually run the policy model on sampled responses, read the logits, and compute the policy loss outside the model. The labels for those sampled tokens are still useful for auxiliary objectives such as MTP, but they should not automatically turn the main model output into a supervised LM loss.

Before this change, MCore uses one labels argument for both purposes:

• labels are consumed by MTP auxiliary loss.
• The same labels also make the main GPT/Hybrid forward return LM loss instead of logits.

That coupling is inconvenient for RL. In online RL + MTP, the desired contract is:

• sampled tokens are passed as labels so MTP can compute its auxiliary CE loss;
• the main forward still returns logits so the RL framework can compute PPO/GRPO/etc. externally;
• the MTP auxiliary loss can be isolated so it only trains MTP-specific parameters, without changing the base policy through decoder hidden states, shared embeddings, or output layer weights.

This PR keeps the regular pretraining/SFT behavior unchanged by default. For RL-style callers, return_logits=True changes only the main LM output contract: labels are still consumed by MTP auxiliary loss, but the forward returns logits instead of main LM loss. mtp_isolated_loss=True separately controls the gradient boundary for the auxiliary objective.

Example RL-style use:

logits = model(
    input_ids=input_ids,
    position_ids=position_ids,
    attention_mask=attention_mask,
    labels=None,
    loss_mask=response_loss_mask,
)

Test plan

• git diff --check
• uv run --project /Users/denliu/Projects/repos/agentic-mcore-dev --extra mcore-lint ruff check megatron/core/transformer/multi_token_prediction.py megatron/core/transformer/transformer_config.py tests/unit_tests/models/test_hybrid_moe_model.py tests/unit_tests/transformer/test_multi_token_prediction.py
• Attempted: uv run --project /Users/denliu/Projects/repos/agentic-mcore-dev python -m pytest tests/unit_tests/transformer/test_multi_token_prediction.py::TestProcessMTPLoss -q (local env missing torch: ModuleNotFoundError: No module named 'torch')

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[Victarry]

Closing this main-branch counterpart because #3456 has already merged into main with the equivalent MTP detach-heads behavior. I will send a separate dev-branch cherry-pick PR for #3456 and revert the now-redundant #5080 changes on dev.