feat(megatron): support MTP in split training

[yfw]

What does this PR do ?

Ports the Megatron multi-token prediction (MTP) training lifecycle from #2801 into the SingleController split training path.

• Derives the MTP loss mask from the existing token and sample masks before microbatch packing.
• Leaves mtp_grad_scale_func unset so MCore inherits the optimizer's loss scale, including fp16 dynamic scaling.
• Normalizes detached MTP auxiliary gradients by the global valid-token count independently of the main loss denominator.
• Rejects nonzero attached-head MTP with a sequence-level main loss because that mixed backbone gradient cannot be corrected after split-step accumulation.
• Collects the independently clipped MTP gradient norm plus per-layer loss and acceptance metrics after the optimizer step.
• Propagates worker-reduced MTP metrics through TQPolicy to the training logger.

No TQ schema or new user configuration fields are required.

Issues

Follow-up to #2801.

• Split path: support MTP + SEQUENCE_LEVEL loss with mtp_detach_heads=False #3889 tracks support for attached-head MTP with a sequence-level main loss.
• Split path: MoE auxiliary loss inherits the main loss denominator instead of 1/global_valid_toks #3888 tracks the analogous pre-existing MoE auxiliary-loss denominator issue.

Usage

Existing MTP-enabled Megatron recipes now use the same MTP training behavior when driven through the SingleController split API.

Before your PR is "Ready for review"

Pre checks:

• Read and followed the contributor guidelines.
• Added focused lifecycle and TQ propagation tests.
• Ran the available focused unit tests and static checks locally.
• No documentation update is needed because this restores parity for existing configuration.

Additional Information

Validation:

• Focused unit tests: 10 passed, 1 skipped.
• Ruff lint and formatting checks passed.
• Repository pre-commit hooks, including Pyrefly, passed.
• Python compilation and git diff checks passed.

The MCore-specific state-machine tests were skipped locally because Transformer Engine is unavailable on this host. Installing the full MCore extras is also blocked here by the missing infiniband/mlx5dv.h system header.

[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.

[yfw]

The MTP lifecycle port looks sound. scale_gradients does reach the MTP parameters (it iterates buffers + expert_parallel_buffers and scales each whole buffer, with no MTP or expert exclusion), so the raw-accumulate-per-chunk then normalize-once-at-finish contract holds. The loss mask is load-bearing too — without it process_mtp_loss gets loss_mask=None and trains MTP over prompt and padding tokens.

Three inline comments: a latent fp16 issue in the gradient-scale hook, a test assertion that cannot currently fail, and one intent question about auxiliary-branch normalization that is pre-existing on this path.

Generated by Claude Code

[yfw]

/ok to test 60c55c8