From 3807d11f30524cbcee243f35d44d4d34a13f911c Mon Sep 17 00:00:00 2001 From: Deepak Narayanan Date: Thu, 18 Jun 2026 12:41:57 -0700 Subject: [PATCH] Rename CP batch helpers to describe balancing granularity Rename get_sft_batch_on_this_cp_rank to _get_batch_on_this_cp_rank_per_document_balancing and get_pretrain_batch_on_this_cp_rank to _get_batch_on_this_cp_rank_per_sequence_balancing. Both helpers apply the same zigzag load-balanced chunking strategy, but at different granularities: per-document (via TE's thd_get_partitioned_indices) vs per-sequence (global). The old names implied the functions were tied to specific training modes (SFT vs pretraining). Prefixed with underscores since they are internal to get_batch_on_this_cp_rank. Signed-off-by: Deepak Narayanan Co-Authored-By: Claude Opus 4.6 Signed-off-by: Deepak Narayanan --- megatron/core/utils.py | 79 +++++++++---------- .../models/mimo/test_mimo_partition.py | 2 +- 2 files changed, 40 insertions(+), 41 deletions(-) diff --git a/megatron/core/utils.py b/megatron/core/utils.py index 2e916482433..169aebc27f9 100644 --- a/megatron/core/utils.py +++ b/megatron/core/utils.py @@ -2253,30 +2253,28 @@ def _broadcast_cu_seqlens(): ######################## -def get_sft_batch_on_this_cp_rank( +def _get_batch_on_this_cp_rank_per_document_balancing( batch: dict[str, torch.Tensor], cp_group: torch.distributed.ProcessGroup ): - """Partition an SFT packed-sequence batch across context-parallel ranks using THD indexing. + """Partition a batch across CP ranks with per-document zigzag load balancing. - For SFT workloads the batch contains multiple variable-length sub-sequences - packed contiguously (THD format). This function uses Transformer Engine's - ``thd_get_partitioned_indices`` to compute the token indices assigned to the - current CP rank and gathers only those tokens from every sequence-dimension - tensor in the batch. - - Metadata keys ('attention_mask', 'cu_seqlens', 'cu_seqlens_padded', - 'max_seqlen', 'local_cp_size', 'hybrid_cp_group') are left unchanged - because TE's attention kernels consume them directly. + Applies zigzag load-balanced chunking independently within each + sub-sequence (document) using Transformer Engine's + ``thd_get_partitioned_indices``. Each document length must be + divisible by ``2 * cp_size``. Sequence-dimension tensors (tokens, + labels, loss_mask, position_ids) are index-selected to this CP + rank's partition; metadata keys (cu_seqlens, cu_seqlens_padded, + max_seqlen, etc.) are left unchanged. Args: batch (dict[str, torch.Tensor]): Batch dict with tensors of shape ``[micro_batch_size, seq_length, ...]``. - cp_group (torch.distributed.ProcessGroup): The context-parallel process - group. + cp_group (torch.distributed.ProcessGroup): The context-parallel + process group. Returns: dict[str, torch.Tensor]: The batch with sequence-dimension tensors - index-selected to this CP rank's partition. + partitioned to this CP rank. """ cp_size = torch.distributed.get_world_size(cp_group) cp_rank = torch.distributed.get_rank(cp_group) @@ -2305,31 +2303,31 @@ def get_sft_batch_on_this_cp_rank( return batch -def get_pretrain_batch_on_this_cp_rank( +def _get_batch_on_this_cp_rank_per_sequence_balancing( batch: dict[str, torch.Tensor], cp_group: torch.distributed.ProcessGroup ): - """Partition a pretraining batch across context-parallel ranks with load-balanced chunking. - - With causal masking, each token only attends to its prior tokens. Simply splitting - the sequence into CP chunks can result in severe load imbalance, as chunks at the - end of the sequence have bigger workloads than earlier ones. To address this, the - sequence is split into ``2 * cp_size`` chunks and assigned in a zigzag pattern: - for CP=2 the 4 chunks are assigned as (chunk_0, chunk_3) -> GPU 0 and - (chunk_1, chunk_2) -> GPU 1, balancing the workload across the CP group. - - All tensor-valued entries in the batch are partitioned along their sequence - dimension (``seq_dim=1`` by default, ``seq_dim=2`` for 'attention_mask'). - None-valued entries are left unchanged. + """Partition a batch across CP ranks with per-sequence zigzag load balancing. + + Applies zigzag load-balanced chunking across the entire sequence. The + sequence is split into ``2 * cp_size`` equal chunks and assigned in a + zigzag pattern: for CP=2, the 4 chunks are assigned as + (chunk_0, chunk_3) -> GPU 0 and (chunk_1, chunk_2) -> GPU 1, balancing + compute for causal attention where later tokens attend to more + predecessors. The sequence length must be divisible by + ``2 * cp_size``. All tensor-valued entries in the batch are + partitioned along their sequence dimension; metadata keys + (cu_seqlens, cu_seqlens_padded, max_seqlen, etc.) and None-valued + entries are left unchanged. Args: batch (dict[str, torch.Tensor]): Batch dict with tensors of shape ``[micro_batch_size, seq_length, ...]``. - cp_group (torch.distributed.ProcessGroup): The context-parallel process - group. + cp_group (torch.distributed.ProcessGroup): The context-parallel + process group. Returns: dict[str, torch.Tensor]: The batch with sequence-dimension tensors - sliced to this CP rank's zigzag partition. + partitioned to this CP rank. """ cp_size = torch.distributed.get_world_size(cp_group) @@ -2376,15 +2374,14 @@ def get_batch_on_this_cp_rank( Routes to the appropriate CP partitioning strategy based on the batch contents and parallelism mode: - - **SFT (packed sequences)**: When ``cu_seqlens`` is present and - ``is_hybrid_cp`` is False, delegates to ``get_sft_batch_on_this_cp_rank`` - which uses THD index-based partitioning. + - **Per-document zigzag**: When ``cu_seqlens`` is present and + ``is_hybrid_cp`` is False, delegates to + ``_get_batch_on_this_cp_rank_per_document_balancing``. - **Hybrid CP**: When ``cu_seqlens`` is present and ``is_hybrid_cp`` is True, creates a local hybrid CP group (via ``hybrid_cp_group_func``) - and delegates to ``get_pretrain_batch_on_this_cp_rank`` with that group. - - **Pretraining**: When ``cu_seqlens`` is None, delegates to - ``get_pretrain_batch_on_this_cp_rank`` with zigzag load-balanced - chunking. + and delegates to ``_get_batch_on_this_cp_rank_per_sequence_balancing``. + - **Per-sequence zigzag**: When ``cu_seqlens`` is None, delegates to + ``_get_batch_on_this_cp_rank_per_sequence_balancing``. Args: batch (Dict[str, Any]): Input batch tensors. Must contain a @@ -2408,12 +2405,14 @@ def get_batch_on_this_cp_rank( ), "local_cp_size is required for hybrid context parallel" if batch['local_cp_size'].item() > 1: hybrid_cp_group = hybrid_cp_group_func(group_size=batch['local_cp_size'].item()) - batch = get_pretrain_batch_on_this_cp_rank(batch, cp_group=hybrid_cp_group) + batch = _get_batch_on_this_cp_rank_per_sequence_balancing( + batch, cp_group=hybrid_cp_group + ) batch["hybrid_cp_group"] = hybrid_cp_group else: - batch = get_sft_batch_on_this_cp_rank(batch, cp_group=cp_group) + batch = _get_batch_on_this_cp_rank_per_document_balancing(batch, cp_group=cp_group) else: # NOTE(asolergi-nv): Pretrain case - batch = get_pretrain_batch_on_this_cp_rank(batch, cp_group=cp_group) + batch = _get_batch_on_this_cp_rank_per_sequence_balancing(batch, cp_group=cp_group) return batch diff --git a/tests/unit_tests/models/mimo/test_mimo_partition.py b/tests/unit_tests/models/mimo/test_mimo_partition.py index da5c1eb440a..72071def5d7 100644 --- a/tests/unit_tests/models/mimo/test_mimo_partition.py +++ b/tests/unit_tests/models/mimo/test_mimo_partition.py @@ -330,7 +330,7 @@ def test_thd_path_raises_when_te_unavailable(self): def _expected_cp_zigzag_shard(tensor: torch.Tensor, cp_size: int, cp_rank: int) -> torch.Tensor: """Reconstruct the CP zigzag shard of ``tensor`` along the sequence dim (dim 1). - Mirrors ``get_pretrain_batch_on_this_cp_rank``: the sequence is split into + Mirrors ``_get_batch_on_this_cp_rank_per_sequence_balancing``: the sequence is split into ``2 * cp_size`` equal chunks and rank ``r`` keeps chunks ``r`` and ``2*cp_size - r - 1`` (concatenated in that order). Implemented independently here so the real-distributed assertions do not lean on the production helper.