Wire CP into MIMO

[Mellonta]

Summary

Integrates the shared CP batch/layout handling into MIMO and extends heterogeneous MIMO grids with GTP sequence partitioning from #7748.

• Partition language embeddings, token IDs, positions, labels, loss masks, and MTP masks together through ContextParallelBatch, keeping dense and packed inputs aligned across contiguous and zigzag CP layouts.
• Build separate weight and batch/sequence views of each language grid. Weight groups control GTP weight sharding; sequence groups combine --mimo-llm-cp with sequence partitions inside GTP. Remaining GTP partitions process independent microbatches.

For example, with TP=1, weight shards=64, sequence shards=64, and --mimo-llm-cp 2, the language grid partitions each sequence across 128 ranks while retaining 64-rank weight-sharding groups.

Test

uv run python -m torch.distributed.run --nproc-per-node 8 -m pytest -q \
  tests/unit_tests/models/mimo \
  tests/unit_tests/pipeline_parallel/test_bridge_communicator.py \
  tests/unit_tests/data/test_get_batch.py \
  tests/unit_tests/transformer/test_multi_token_prediction.py

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

[fanshiqing]

/review

[svcnvidia-nemo-ci]

Light review completed across the changed files and directly affected CP/GTP paths. One logic error and one configuration-documentation issue remain. Tests were not executed; suggested changes below require author verification.

1. Dual-layout batches crash GPT-backed MIMO

Failure. With CP=2, linear_cp_layout="contiguous", attention_cp_layout="zigzag", and a dense [1, 8] input, the adapter creates two layout views. The new language-forward helper consequently passes cp_batch to the language model. However, GPTModel.forward has no such parameter, so GPT-backed MIMO raises TypeError: ... unexpected keyword argument 'cp_batch'. GPT remains a MIMO receiver, including in the mock provider.

Correction. Distinguish the receiver contracts when configuring the adapter. Legacy GPT's TransformerBlock execution does not perform CP layout transitions, so GPT-backed MIMO should prepare one batch in the attention layout. Hybrid-backed MIMO must retain its linear boundary layout, alternate attention view, and cp_batch. Merely suppressing the keyword would leave GPT consuming contiguous activations as zigzag attention inputs.

A small suggested patch for MimoModel initialization: add imports for dataclasses.replace and GPTModel, then insert this immediately after self._initialize_language_model():

if self.partition_adapter is not None and isinstance(
    unwrap_model(self.language_model), GPTModel
):
    self.partition_adapter = PartitionAdapter(
        replace(
            self.partition_adapter.cfg,
            linear_cp_layout=self.partition_adapter.cfg.attention_cp_layout,
        )
    )

This restores the invariant that each receiver gets activations in its supported boundary ordering and only supported keywords.

Preservation. Ordinary GPT zigzag CP must remain unchanged; CP=1 must remain a passthrough; SP must still scatter only decoder embeddings. Hybrid/Mamba receivers must continue receiving both layouts, including the converted MTP validity mask. Encoder-only ranks must still have no partition adapter.

Acceptance — proposed regression checks, not executed. Add a GPT-backed MIMO case with CP=2, TP=1, the differing layouts above, sequence length 8, and identifiable embedding rows. Keep GPT's real forward signature: if the environment lacks a CP attention backend, replace only its decoder as the existing MIMO tests do. Forward must not raise, rank 0 must receive embedding rows [0, 1, 6, 7], rank 1 rows [2, 3, 4, 5], and labels/loss masks must follow those same rows. Also retain the CP/SP shape-and-content assertions in the distributed partition tests, and the assertion that Hybrid MTP receives the zigzag mask in the MTP regression. Verify statically that the normalization is limited to GPT receivers.

2. Document the added PartitionConfig fields

Failure. The new public fields linear_cp_layout, attention_cp_layout, and tp_cp_group in PartitionConfig lack adjacent MyST docstrings, contrary to the base configuration-documentation rule. This is a documentation issue, not a demonstrated runtime failure.

Correction. Add descriptions immediately after those declarations, for example:

linear_cp_layout: CPLayout = "zigzag"
"""Boundary CP ordering used when preparing language-model inputs."""
attention_cp_layout: CPLayout = "zigzag"
"""Attention/MTP ordering; a differing layout requests an additional batch view."""
cp_group: Optional[ProcessGroup] = None
tp_group: Optional[ProcessGroup] = None
tp_cp_group: Optional[ProcessGroup] = None
"""Explicit TP x CP group used for sequence-parallel layout conversion."""

Preservation and acceptance — static checks. Keep the existing field types, defaults, and group forwarding unchanged. Inspect that each of the three added public fields has its own immediately following docstring and that the descriptions agree with the adapter's layout selection and the conversion helper's explicit-group requirements. No runtime test is needed for this documentation-only correction.

Note: formal approval is not available for this request, so this review is posted as a comment only.

Review generated by codex

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38688830700