[HaochenYuan]

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

What does this PR do?

#7000 added pipeline_p2p_fixed_shape to skip the per-transfer shape handshake for packed (THD) sequences padded to max_seqlen_per_dp_cp_rank, but rejected it under virtual pipeline parallelism: forward_backward_pipelining_with_interleaving sized its receive buffers from seq_length // cp_size × micro_batch_size instead of calling get_tensor_shapes(), so skipping the handshake would post mismatched irecv buffers.

This PR makes the interleaved schedule derive the same fixed packed shape as get_tensor_shapes() — [max_seqlen_per_dp_cp_rank, 1, hidden_dim], then the sequence-parallel split — when variable_seq_lengths and pipeline_p2p_fixed_shape, and removes the VPP rejection (and its layout-derived re-check in TransformerConfig). The two #7000 tests that asserted the rejection now assert acceptance.

Validation

• DSv4 hybrid (mHC, DSA, MoE) PP2×VPP2 THD proxy, eager and TE CUDA graphs: runs with the flag, losses in band with the flag-less runs; step time −6% (eager) / −3.5% (graphs).
• 128 GPUs, PP2×VPP4, CP64, 1M tokens: 50 iterations, no NaN, 1074 TFLOP/s/GPU vs 1059 without the flag.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

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
• If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
• 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.