[dev] moe(perf): Refactor CP layout - auto layout conversion

[yuzhongw-nvidia]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Depends on: #6387
PR for main: #6233

Summary

This PR makes context-parallel token layout an explicit per-batch / per-layer contract instead of a single global TransformerConfig.cp_partition_mode. The model can now keep the input batch layout, convert hidden states only at layout-sensitive layer boundaries, and restore model outputs to the expected input layout at the postprocess boundary.

The main target is long-context hybrid training where different modules prefer different CP layouts:

• Standard attention / MLP paths continue to work with the existing zigzag layout.
• Chunkwise GatedDeltaNet uses contiguous CP layout.
• DSv4 / compressed sparse attention uses contiguous THD layout.
• Packed and non-packed batches carry their current CP layout through PackedSeqParams.

Key Changes

• Added megatron.core.context_parallel_layout as the shared CP-layout package:
  • SBHD and THD zigzag/contiguous conversion primitives.
  • THD route prebuild helpers stored directly on PackedSeqParams.
  • Per-layer required-layout declarations.
  • Stage-entry layout policy helpers for GPT / Hybrid models.
• Updated GPT and Hybrid model/block forward paths to:
  • infer entry layout from batch metadata,
  • convert hidden states only when a layer requires a different layout,
  • convert batch-side tensors when MTP or postprocess needs them aligned,
  • restore postprocess outputs to the input layout.
• Updated GDN, DSv4, MLA/RoPE, MTP, and sequence-packing paths to consume explicit CP layout metadata.
• Deprecated TransformerConfig.cp_partition_mode for model-level planning. It is still validated if set, but model-level layout is inferred from module requirements and batch metadata.
• Added and refined unit coverage for:
  • CP layout primitive correctness,
  • THD route prebuild / CUDA graph metadata behavior,
  • GPT and Hybrid output-layout restoration,
  • GDN and DSv4 layout requirements,
  • GDN + MoE CP loss parity for SBHD and THD.

Compatibility / Limits

• PackedSeqParams.cp_partition_mode is expected for CP > 1 model forwards. For now, GPT / Hybrid model forward paths can synthesize temporary metadata with a warning, but callers should pre-populate it. This warning path is intentionally not exercised by the updated unit tests.
• THD routes should be prebuilt in data-preparation code with prebuild_thd_cp_partition_routes. Runtime fallback still exists with a warning for compatibility, but tests and pretrain batch paths prebuild the routes explicitly.
• Non-packed / SBHD MTP currently supports zigzag layout. If contiguous SBHD MTP is requested, it raises rather than silently rolling with the wrong layout.
• Review note for discussion_r3691308680: _get_thd_freqs_on_this_cp_rank is not currently on an in-repo runtime call path after _apply_rotary_pos_emb_thd moved to inline freq_pos computation. It is kept with cp_partition_mode handling as a compatibility helper for external / bridge-like code that may still call it directly.

Timeline

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

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

[yuzhongw-nvidia]

/claude review

[claude]

LGTM

[yuzhongw-nvidia]

/claude strict-review

[claude]

Strict Review Summary

Reviewed this CP-layout refactor with focus on correctness, training performance, and backward compatibility.

Findings by severity: CRITICAL: 1, IMPORTANT: 0, SUGGESTION: 0

Most impactful finding

• [CRITICAL Correctness] Removed unknown-mode guard in get_thd_context_parallel_rank_indices (megatron/core/context_parallel_layout.py). The refactor deleted the top-of-function if layout not in ("zigzag", "contiguous"): raise ValueError(...) check. Any mode that is not "contiguous" now silently falls through to the zigzag path. This (a) breaks the still-present test test_thd_context_parallel_rank_indices_reject_unknown_layout, which expects a ValueError for "interleaved", and (b) diverges from every sibling helper in this file that still raises on unknown modes — turning a typo'd/future mode into a silently-wrong partition rather than a loud failure. Restore an explicit final guard.

Overall assessment

This is a well-structured, well-documented refactor that centralizes CP partition-mode ownership at the TransformerBlock/GPTModel level and threads cp_partition_mode cleanly through the RoPE, batch-partitioning, and attention-validation paths. The default "zigzag" everywhere preserves backward compatibility for existing standard-attention configs, and the new copy-on-write PackedSeqParams handling plus the attention-side validation asserts give good defense-in-depth. Duck-typed layer→mode mapping fails loudly on unknown layer types, and deferred args (tp_cp_group) are explicitly documented.

The zigzag↔contiguous conversion math (THD contiguous positions, freq offsets, SP gather/scatter path) traced consistently across context_parallel_layout.py, rope_utils.py, attention.py, and gpt_model.py.

Risk level: Low-to-moderate, contingent on fixing the one CRITICAL. It is a genuine correctness/loud-failure regression and will fail the existing unit test in CI. Everything else looked correct.

[yuzhongw-nvidia]

/claude strict-review

[claude]

Strict Review Summary

Reviewed the CP partition-mode refactor (zigzag to contiguous layout ownership moved from module-local conversions in GatedDeltaNet up to block/model-level scheduling). The core design is sound and the RoPE/THD index math checks out: I verified the contiguous THD freq_pos derivation, the exact-packed-freqs shift, and the fused-RoPE zigzag-only fallback. New public callers correctly use keyword args, so the inserted cp_partition_mode parameters do not break positional call sites.

Findings by severity: CRITICAL: 1, IMPORTANT: 1, SUGGESTION: 2

Most impactful:

1. [CRITICAL Compatibility] HybridModel defaults cp_stage_entry_partition_mode=None and forwards it unconditionally into HybridStack._build_cp_partition_mode_plan(), which raises ValueError on None at construction time. Every HybridModel(...) call that omits the new argument now crashes on init, including megatron/post_training/model_builder.py and roughly twenty unit tests. Peer models (GPTModel, HybridStack, TransformerBlock) all default to zigzag. The docstring also claims a pattern-based inference that does not exist. Fix by defaulting to zigzag or implementing the promised inference.

2. [IMPORTANT Correctness] get_hybrid_stage_input_cp_partition_mode_for_stage accepts tp_group/dp_cp_group then deletes them; all four call sites fetch and pass these groups. Either wire them in or drop them from the signature and callers (per the repo unused-new-argument check).

Risk assessment: High risk as-is. The CRITICAL finding is a hard break of the public HybridModel constructor that surfaces immediately in CI (unit tests omit the new arg). Once the default is corrected, the remaining items are low-risk cleanups. The standard non-hybrid/non-experimental GPT path is unaffected (source==target partition mode short-circuits with no extra work). Recommend re-running the hybrid + MTP unit suites after fixing the default.

[yuzhongw-nvidia]

/claude strict-review

[claude]

Strict Review Summary

Reviewed the CP partition-mode refactor across RoPE utils, transformer/hybrid blocks, GPT/Hybrid models, and batch partitioning. The core design — a per-stage cp_stage_entry_partition_mode plan driving layer-boundary layout conversions, with attention modules validating (not mutating) the mode they receive — is sound, and the removal of module-local zigzag↔contiguous shuffles from gated_delta_net.py in favor of block-level scheduling is a clean improvement. Backward compatibility is preserved: all new parameters default to "zigzag"/None, the CP-size-1 fast path short-circuits, and packed_seq_params.cp_partition_mode defaults to zigzag. Unit-test coverage for the new planning helpers is reasonable.

Findings by severity:

• CRITICAL: 0
• IMPORTANT: 1
• SUGGESTION: 2

Most impactful:

1. [IMPORTANT Correctness] _call_rope_helper in examples/multimodal_dev/models/qwen35_vl/specs.py silently drops kwargs (including cp_partition_mode) not present in the target helper's signature. Against an older MCore where _apply_rotary_pos_emb_thd lacks cp_partition_mode, a contiguous request silently degrades to zigzag RoPE — wrong positional embeddings with no error. Suggest raising when a non-default mode can't be passed through.
2. [SUGGESTION Simplification] The new cp_partition_mode/contiguous branch in _get_thd_freqs_on_this_cp_rank (rope_utils.py) is dead code — the function has no in-repo callers.
3. [SUGGESTION Simplification] The MTP layout-conversion block (nested _convert_sequence_tensor/_convert_rotary closures) is duplicated near-verbatim across gpt_model.py and hybrid_model.py, and _convert_rotary_cp_partition_mode is duplicated across transformer_block.py/hybrid_block.py. Worth extracting a shared helper.

Overall risk: Low–Moderate. No correctness defects found in the core CP conversion logic (SP gather/scatter symmetry, THD index computation, and RoPE freq mapping check out). The one IMPORTANT item is confined to example VLM code and is a graceful-degradation hazard rather than a core-path bug. Recommend addressing the silent-kwarg-drop before this pattern is copied to other bridges. Functional tests exercising contiguous CP with GDN + MTP + PP would strengthen confidence.

[claude]

Review summary — CP layout refactor

Reviewed the full diff for implementation correctness, distributed correctness (TP/PP/SP/CP/EP), training-performance, and backward compatibility, then read the full source of the highest-risk changed files (rope_utils.py, conversion.py, gpt_model._postprocess, gated_delta_net.py, policy/routes/metadata, transformer_block.py, multi_token_prediction.py, the get_batch entrypoints, and the data/layout helpers).

Overall this is a well-structured refactor: the split context_parallel_layout/ package has clear ownership boundaries, the route-tensor precompute + single-all-to-all conversion is a sound design, None-mode fast-paths are consistently guarded by cp_size > 1, and the layered validation (get_stage_entry_partition_mode -> per-layer get_required_cp_partition_mode_for_layer -> attention-level _validate_packed_seq_params_cp_partition_mode) fails loud on layout mismatches rather than silently mis-partitioning. I did not find a correctness bug in the layout conversion, the contiguous-RoPE branch, or the _postprocess loss round-trip (postprocess_to_input_converter is correctly initialized to None; the loss is converted back to the input layout at the boundary as documented).

Findings

Severity	Area	Location
IMPORTANT	Numerics/init: GDN conv1d init changed (kaiming weight + bias + RNG fork) shifts golden values	gated_delta_net.py:308
IMPORTANT	Backward-compat: TransformerConfig.cp_partition_mode default zigzag->None, deprecated, and dsv4_hybrid/zigzag validation block removed	transformer_config.py:312
Suggestion	Perf: inspect.signature called per-RoPE-application on the hot path	qwen35_vl/specs.py:27
Suggestion	Maintainability: tp_cp_group accepted but del-ed with no active use path (TODO P2.1)	conversion.py:213

Counts: 0 Critical, 2 Important, 2 Suggestions.

Key action items before merge:

1. Attach the "Run functional tests" label and regenerate/verify golden values for GDN & hybrid-with-GDN recipes — the conv1d init change alters fresh-init numerics (checkpoint resume is unaffected).
2. Confirm the removed dsv4_hybrid->contiguous config invariant is still enforced on the new per-layer inference path so a stale user config cannot silently run the wrong layout, and note the cp_partition_mode deprecation in the release notes.

No blocking correctness issues found. The two Important items are about numeric-baseline management and public-config compatibility rather than incorrect logic.

Process-group usage: no new direct parallel_state.get_*_group() reads were introduced in megatron/core production code (the one get_context_parallel_group() in llava_model.py is pre-existing example-path code).

[hxbai]

Seems that it conflicts with overlap_moe_expert_parallel_comm and full recompute when the layout changes between layers. We should add an assertion if we do not want to support it for now.

[yuzhongw-nvidia]

The auto mode is not in plan. Main works are transferred into #6387.