fix(mtp): generalize context parallelism across MTP models

[HuiyingLi]

What changed

• Add common preparation of per-depth MTP token, position, label, and validity streams before context-parallel sharding.
• Use explicit model capabilities and model-owned position policy instead of assuming one model-specific MTP hook/layout.
• Wire MTP+CP through DeepSeek V4, Qwen3.5, MiniMax M3, and Step3.7, with cross-rank parity and recipe coverage.

Why

The existing recipe support handled MTP+CP around a specific input convention. Other MTP models can use different position semantics or embedding paths, so they need one common global-input contract while retaining model-owned preparation where their layout differs. The CP sharder then performs the actual local tensor placement.

This PR intentionally keeps MiniMax's missing-MTP-checkpoint compatibility out of scope.

Impact

Models explicitly advertising MTP+CP support can use the shared recipe path without duplicating CP index calculations. Models without that capability keep their existing behavior and fail early if MTP+CP is requested.

Validation

• Added unit and distributed parity coverage for common MTP rolling, recipe preparation, and supported model layouts.
• Static compilation and diff checks pass in the current checkout.
• Full pytest was not rerun in the current shell because PyTorch is unavailable there.

50-step MTP+CP parity

• DeepSeek V4: CP1 and CP2 both completed 50 steps on 16 nodes / EP128 / PP1 / GBS256 / LBS1 with sequence length 1024. Every step used identical num_label_tokens; the mean absolute loss difference was 0.000164 and the maximum was 0.0006. CP1 pipeline, CP2 pipeline. The parity probe used sequence length 1024 because the full 16K recipe exceeded H100 memory.
• Step3.7: CP1 and CP8 both completed 50 steps on 18 nodes / EP144 / PP1 / GBS144 / LBS1 with sequence length 2048. Every step used identical num_label_tokens; the mean absolute loss difference was 0.0114 and the maximum was 0.0469 (2.2% maximum relative difference). CP1 pipeline, CP8 pipeline.
• Qwen3.5 dense: CP1 and CP2 both completed 50 steps on Qwen/Qwen3.5-4B with one MTP depth, sequence length 1024, GBS2, LBS1, and LR0. Raw MTP loss was isolated with paired MTP-on (weight 1.0) and MTP-off runs; a direct calculate_mtp_loss probe confirmed that the subtraction gives the returned raw MTP loss. Every step used identical num_label_tokens and num_tokens_per_step. Mean raw MTP loss was 2.361566 for CP1 and 2.361717 for CP2 (0.0064% difference); the mean per-step relative difference was 0.42% and the maximum was 1.41%, within the base-loss CP variation. CP1 MTP-on pipeline, CP1 MTP-off pipeline, CP2 MTP-on pipeline, CP2 MTP-off pipeline.

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

[HuiyingLi]

/claude review

[claude]

Light review — context-parallel MTP support across DeepSeek V4, Step3.7, MiniMax M3, and Qwen3.5, plus the new MTPPositionPolicy enum and an fp32 LM-head dtype cast in calculate_mtp_loss.

Overall the distributed logic reads correctly and is well-tested:

• The CURRENT vs FUTURE position policy matches each model's non-CP default (DSV4/Step/MiniMax retain the backbone query position; Qwen advances with the future token), and the enum's else branch raises on unknown values.
• The changed hunks contain no run-to-run gradient-determinism hazards — only roll_tensor, masked_fill, and round-robin sharding, whose backwards are deterministic; no new accumulating gathers into trainable parameters.
• The fp32 cast (mtp_output.to(dtype=lm_weight/lm_head_weight.dtype)) is correct and matches DSV4's fp32 LM head; the new unit test asserts fp32 loss dtype and finite grads.
• test_mtp_cp_training.py gives real 2-rank CP1==CP2 parity on both losses and MTP gradients across all four families — the meaningful invariant for this change.

Findings (2, both non-blocking, same root cause):

• Tensor-contract docstrings missing on the MTP forward methods that gained new tensor parameters (DeepseekV4MTPModule.forward, MiniMaxM3MTP.forward). The per-depth CP tensors have non-obvious per-rank-local layouts that should be documented. Inline suggestions posted.

The common MTPModule.forward (qwen path) already documents input_ids_per_depth/position_ids_per_depth, so it's a good template for the two model-local copies.

[claude]

DeepseekV4MTPModule.forward gained new tensor parameters (input_ids_per_depth, position_ids_per_depth) and materially changed its tensor-handling body, but has no docstring documenting any tensor layouts. Per the tensor-contract rule, add a Google-style Args/Returns section documenting the semantic shapes, e.g.:

Args:
    hidden_states: Tensor of shape [batch, sequence, hidden].
    input_ids: Optional token IDs of shape [batch, sequence]; rolled per depth.
    input_ids_per_depth: Optional tuple of ``num_depths`` tensors, each [batch, sequence],
        holding the pre-shifted future-token IDs for the corresponding depth (CP path).
    position_ids: Optional [batch, sequence] positions applied at every depth.
    position_ids_per_depth: Optional tuple of ``num_depths`` tensors, each [batch, sequence],
        overriding per-depth positions (CP path).

Returns:
    List of ``num_depths`` prediction-hidden tensors, each [batch, sequence, hidden].

The input_ids_per_depth/position_ids_per_depth local-vs-global (per-rank CP shard) semantics are non-obvious and should be stated explicitly.

[HuiyingLi]

/ok to test f64478c