What does this PR do ?

Unifies batch-processing utilities (get_batch, get_batch_on_this_tp_rank, get_batch_on_this_cp_rank) behind a single framework that handles pretraining, SFT (THD packed sequences), and Hybrid Context Parallel, and enables pipeline parallelism for SFT with THD layout.

Changes

Core refactor (megatron/core/utils.py, megatron/training/utils.py)

• Move get_batch_on_this_tp_rank / get_batch_on_this_cp_rank from megatron/training/utils.py to megatron/core/utils.py.
• Rewrite get_batch_on_this_tp_rank around an explicit, length-prefixed broadcast protocol for THD metadata (cu_seqlens, cu_seqlens_padded, max_seqlen, local_cp_size), with dedicated branches for PP-first / PP-last / PP-intermediate / MTP stages.
• Split CP partitioning into three strategies dispatched by get_batch_on_this_cp_rank:
  • get_sft_batch_on_this_cp_rank — THD index-based partitioning via tex.thd_get_partitioned_indices (SFT).
  • get_pretrain_batch_on_this_cp_rank — zigzag load-balanced chunking (pretrain / Hybrid CP).
  • A dispatcher that selects between them based on cu_seqlens presence and is_hybrid_cp.

PP for SFT (THD)

• get_batch in both entrypoints now broadcasts only the required THD metadata (cu_seqlens, cu_seqlens_padded, max_seqlen) to intermediate PP stages; tokens/labels/loss_mask/position_ids are intentionally left None.
• get_sft_batch_on_this_cp_rank CP-partitions whichever of tokens / labels / loss_mask / position_ids are present on the current PP stage

MTP API cleanup (megatron/core/transformer/multi_token_prediction.py)

• mtp_on_this_rank now takes layout: PipelineParallelLayerLayout and mtp_num_layers: int directly instead of a full TransformerConfig.

Entry scripts (pretrain_gpt.py, pretrain_mamba.py)

• Unified, sorted BATCH_KEYS tuple; single call into the new core utilities.
• Build PackedSeqParams from the returned tensors, passing Python ints (int(max_seqlen.item()), int(local_cp_size.item())) to match PackedSeqParams' declared types.
• pretrain_mamba.py: also pass total_tokens=int(cu_seqlens_for_params[-1].item()) so PackedSeqParams.post_init populates seq_idx for the Mamba SSM kernel (prevents state bleeding across packed sequence boundaries). Works at every PP stage and for any CP size.

Tests (tests/unit_tests/data/test_get_batch.py)

• New suite covering:
  • test_sft_batch — SFT + TP/PP/CP product over {1,2,4}, seq_length ∈ {1024, 4096}, with per-PP-stage shape / dtype / mask / metadata assertions (including the new intermediate-PP branch).
  • test_pretrain_batch — pretrain path across TP/PP/CP × micro_batch_size ∈ {1,4} × create_attention_mask ∈ {True, False}.
  • test_hybrid_cp_batch — Hybrid CP broadcasts and zigzag partitioning.

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 the @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.