[Mellonta]

Background

TE softmax attention requires a zigzag CP layout, while linear attention does not. With
chunkwise CP, using the zigzag layout for linear-attention layers doubles their communication. This PR wires the megatron/core/context_parallel module from #6764 into hybrid models.

Summary

• Add linear_cp_layout and attention_cp_layout to TransformerConfig, both defaulting to zigzag, so existing runs are unchanged.
• Shard the batch into contiguous CP slices when the linear layout is contiguous, keeping the dense attention mask in zigzag.

Tests

uv run python -m torch.distributed.run --nproc-per-node 4 -m pytest -q \
    tests/unit_tests/data/test_get_batch.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.

[Mellonta]

/claude strict-review

[Phlip79]

This file is already way too big (not your fault).

Can you move this function to megatron/core/context_parallel/utils.py or similar?

[Phlip79]

Please update PR description as well

[Mellonta]

/claude strict-review

[Mellonta]

/ok to test fb9538a

[Mellonta]

/ok to test 2acc3d1

[Mellonta]

/ok to test e368ac8

[Mellonta]

/claude review

[zhongbozhu]

LGTM

[claude]

Light review — a few things worth addressing before merge.

1. forward_step reads args but never binds it (pretrain_hybrid.py:333)

forward_step has no args = get_args(), and there is no module-level args in pretrain_hybrid.py (the only bindings are locals at lines 113, 238, 367, 431). args.seq_length on line 346 is pre-existing and already raises NameError whenever cu_seqlens is not None, but this PR adds a second read on the same path:

if args.linear_cp_layout == "contiguous" and args.context_parallel_size > 1:

Compare pretrain_gpt.py:303, which binds args = get_args() at the top of its forward_step.

Worth noting that the new test cannot catch this — it does:

monkeypatch.setattr(pretrain_hybrid, "args", SimpleNamespace(...), raising=False)

raising=False creates a module attribute that does not exist in production, so the test passes while the real code path fails. Suggested fix in forward_step:

def forward_step(data_iterator, model: HybridModel):
    ...
    args = get_args()
    timers = get_timers()

Then drop raising=False from the test (or monkeypatch pretrain_hybrid.get_args) so it exercises the real lookup.

2. _get_batch_on_this_cp_rank_contiguous uses an allow-list where the zigzag path uses a deny-list

megatron/core/context_parallel/utils.py shards only ('tokens', 'labels', 'loss_mask', 'position_ids'), while _get_batch_on_this_cp_rank_per_sequence_balancing (megatron/core/utils.py:2536) shards everything except an explicit METADATA_KEYS set. Any future sequence-dim batch key will be silently left unsharded on the contiguous path — a wrong-results failure mode rather than an error. Consider mirroring the deny-list so the two paths stay in sync, or at least raising on an unrecognized non-metadata tensor key.

3. Docstring bullet order does not match dispatch precedence

In get_batch_on_this_cp_rank the new Contiguous CP bullet is appended last, but the implementation checks use_contiguous_cp first, so it overrides use_per_sequence_balancing and is_hybrid_cp. Worth moving the bullet to the top and stating that it takes precedence — otherwise a caller passing both use_per_sequence_balancing=True and use_contiguous_cp=True will reasonably expect the former to win.

4. Test coverage for the new integration points

The layout unit tests from #6764 cover ContextParallelLayoutManager in isolation, and the new tests cover batch sharding plus packed_seq_params selection. Not covered by any test in this PR:

• HybridStack.forward / checkpointed_forward with cp_size > 1 and linear_cp_layout != attention_cp_layout — that prepare_layer/finalize_layer fire at the right layer indices and that the recompute and non-recompute paths agree.
• sequence_is_contiguous=True in MambaContextParallel / GDPContextParallel — the pre_conv_ssm / post_conv_ssm branches that now skip _undo/_redo_attention_load_balancing. A round-trip test asserting post_conv_ssm(pre_conv_ssm(x)) reconstructs x under both flag values would pin this down.
• The new _validate_cp_layouts rejections (invalid layout string, attention_cp_layout='contiguous' with CP, MTP, hybrid_context_parallel) — cheap pytest.raises cases, no GPU needed.

5. Minor: linear_cp_layout is only honored by HybridStack

_validate_cp_layouts accepts linear_cp_layout='contiguous' for any TransformerConfig, but the layout manager is built only in HybridStack, and sequence_is_contiguous is threaded only into MambaMixer / GatedDeltaProductMixer. A GPT model with linear-attention layers (linear_attention_freq) would accept the flag and silently ignore it. If contiguous is hybrid-only for now, consider rejecting it elsewhere or documenting the scope on the config field.

[nemo-automation-bot]

🔄 Merge queue validation started!

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