From ce55b79fce1f7bb32ac00e312e8f059d0059b7b8 Mon Sep 17 00:00:00 2001 From: "kayeon.song" Date: Tue, 29 Sep 2026 03:38:15 +0000 Subject: [PATCH 01/13] feat(data): intra-microbatch reordering for MegatronMIMO scalable reads Rebalance each canonical-grid micro-batch across a module's DP ranks by a per-sample cost (encoder patches, optionally language tokens) so paired vision/language replicas keep identical sample order. Adds ReorderingBuffer (window exchange, optional prefetch overlap), dataset config fields megatron_mimo_intra_microbatch_reorder / *_reorder_encoder_cost_weight / *_reorder_language_cost_weight / *_reorder_overlap / *_reorder_window_size, and the --intra-microbatch-reorder wiring in the Qwen3.5-VL example. A dp=1 module reorders locally into the same canonical order. The example collate now always ships the cheap [B, S] fields (input_ids, position_ids, attention_mask) on every stage so data-path consumers (reorder cost, packing lengths, MRoPE) need no per-feature gating; forward_step alone decides what the model receives and always drops the tokenizer attention_mask. Signed-off-by: kayeon.song Signed-off-by: Yoonsik Kim Signed-off-by: Chanwoo Park Co-authored-by: Yoonsik Kim Co-authored-by: Chanwoo Park <45866990+sailor1493@users.noreply.github.com> --- .../qwen35_vl/finetune_qwen35_vl.py | 119 +- .../bridge/data/builders/direct_hf_sft.py | 28 + .../bridge/data/megatron_mimo/dp_utils.py | 31 + .../data/megatron_mimo/reorder_buffer.py | 1539 +++++++++++++++++ .../bridge/training/megatron_mimo_step.py | 3 + .../bridge/training/train_megatron_mimo.py | 5 + .../megatron_mimo/test_reorder_exchange.py | 175 ++ .../data/builders/test_direct_hf_sft.py | 50 + .../data/megatron_mimo/test_reorder_buffer.py | 1186 +++++++++++++ .../examples/test_qwen35_vl_mimo_finetune.py | 46 + .../megatron_mimo/test_megatron_mimo_step.py | 37 + 11 files changed, 3200 insertions(+), 19 deletions(-) create mode 100644 src/megatron/bridge/data/megatron_mimo/reorder_buffer.py create mode 100644 tests/functional_tests/test_groups/megatron_mimo/test_reorder_exchange.py create mode 100644 tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py diff --git a/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py b/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py index 338ca008225..7f60c25607d 100644 --- a/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py +++ b/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py @@ -11,8 +11,9 @@ Conversation examples are built with the standard HF VLM provider, then the resulting Qwen batch is adapted into the MIMO forward shape: - - language inputs: ``input_ids``, MRoPE ``position_ids``, labels, loss mask, and - (when packing) the tokenizer's ``attention_mask`` + - language inputs: ``input_ids``, MRoPE ``position_ids``, labels, loss mask, and the + tokenizer's ``attention_mask`` (a data-path length source; the forward step drops it + before the model) - image inputs: ``modality_inputs["images"]["qwen_visual"]`` Example 2-GPU smoke: @@ -29,6 +30,7 @@ from __future__ import annotations import argparse +import functools import logging import math import os @@ -58,8 +60,8 @@ chat_template_kwargs_from_example, ) from megatron.bridge.data.datasets.utils import IGNORE_INDEX -from megatron.bridge.data.megatron_mimo.canonical_sampler import build_canonical_mimo_data_loader -from megatron.bridge.data.megatron_mimo.dp_utils import get_megatron_mimo_sampling_info +from megatron.bridge.data.megatron_mimo.canonical_sampler import build_canonical_mimo_data_loader, canonical_grid_size +from megatron.bridge.data.megatron_mimo.dp_utils import _find_rank_module, get_megatron_mimo_sampling_info from megatron.bridge.data.samplers import build_pretraining_data_loader from megatron.bridge.data.sources.hf import hf_dataset_supports_split from megatron.bridge.data.token_utils import extract_skipped_token_ids @@ -132,7 +134,6 @@ class MIMOBatchSpec: labels: bool = True loss_mask: bool = True modality_inputs: bool = True - attention_mask: bool = False def describe(self) -> str: enabled = [ @@ -143,7 +144,6 @@ def describe(self) -> str: ("labels", self.labels), ("loss_mask", self.loss_mask), ("modality_inputs", self.modality_inputs), - ("attention_mask", self.attention_mask), ) if value ] @@ -279,31 +279,27 @@ def _batch_spec_for_rank(cfg: Any) -> MIMOBatchSpec: pp_rank = _grid_dim_rank(grid, "pp", 0) pp_size = _grid_dim_size(grid, "pp", 1) - is_first_pp = pp_rank == 0 is_last_pp = pp_rank == pp_size - 1 - dataset_cfg = getattr(cfg, "dataset", None) - packing_active = bool(getattr(dataset_cfg, "enable_in_batch_packing", False)) - + # Cheap [B, S] fields (input_ids, position_ids, attention_mask) are always shipped: the data path + # (reorder cost, packing lengths, MRoPE) reads them on every stage, and forward_step alone decides + # what the model receives. Only the expensive or last-stage-only fields are gated here. if module_name == MIMO_LANGUAGE_MODULE_KEY: return MIMOBatchSpec( - input_ids=is_first_pp, - # Qwen3.5-VL mRoPE needs position_ids on every language PP stage. + input_ids=True, position_ids=True, labels=is_last_pp, loss_mask=is_last_pp, modality_inputs=False, - # Packing length source; the packer nulls it again before the model. - attention_mask=packing_active, ) + # Only an encoder's first PP stage loads data, so nothing is stage-gated here. return MIMOBatchSpec( - # Encoder first stages need input_ids to attach per-sample split metadata. - input_ids=is_first_pp, + input_ids=True, position_ids=False, labels=False, loss_mask=False, - modality_inputs=is_first_pp, + modality_inputs=True, ) @@ -321,8 +317,6 @@ def _project_adapted_batch( adapted["loss_mask"] = None if not batch_spec.modality_inputs: adapted["modality_inputs"] = None - if not batch_spec.attention_mask: - adapted["attention_mask"] = None return adapted @@ -418,6 +412,11 @@ def _build_dataset_config(args: argparse.Namespace) -> DirectHFSFTDatasetConfig: enable_in_batch_packing=args.pack_sequences_in_batch, defer_in_batch_packing_to_step=True, megatron_mimo_scalable_dp=args.scalable_dp, + megatron_mimo_intra_microbatch_reorder=args.intra_microbatch_reorder, + megatron_mimo_reorder_encoder_cost_weight=args.reorder_encoder_cost_weight, + megatron_mimo_reorder_language_cost_weight=args.reorder_language_cost_weight, + megatron_mimo_reorder_overlap=not args.no_overlap_intra_microbatch_reorder, + megatron_mimo_reorder_window_size=args.reorder_window_size, do_validation=do_validation, do_test=False, trust_remote_code=args.trust_remote_code, @@ -953,6 +952,62 @@ def _build_data_iterators(cfg, _megatron_mimo_infra, *, train_state=None): # `pretrain_megatron_mimo` calls `next(data_iterator)` per microbatch, so # return an iterator (DataLoader is iterable but not itself an iterator). loader_iter: Iterator[dict[str, Any]] = iter(train_loader) + + # Gate on the canonical grid (same on every rank), not this rank's DP size: dp=1 modules must + # still join the world-collective PG creation below and reorder locally to stay paired. + reorder_on = ( + scalable_dp + and bool(getattr(cfg.dataset, "megatron_mimo_intra_microbatch_reorder", False)) + and canonical_grid_size(module_dps) > 1 + ) + if reorder_on: + from megatron.bridge.data.megatron_mimo.reorder_buffer import ( + ReorderingBuffer, + build_module_dp_process_groups, + sample_cost, + ) + + for _name, _par in cfg.model.megatron_mimo_parallelism_config.module_parallelisms.items(): + if _name != MIMO_LANGUAGE_MODULE_KEY and _par.pipeline_model_parallel_size > 1: + raise NotImplementedError( + "MegatronMIMO intra-microbatch reorder does not support encoder pipeline parallelism " + f"(module {_name!r} has PP={_par.pipeline_model_parallel_size}): non-first encoder stages " + "skip data loading, so the collective process-group creation below would hang." + ) + # n_groups is the canonical grid size (LCM of the module DP sizes), the geometry the canonical + # loader reads with, so vision and language ranks derive one identical per-sample assignment. + balance_n_groups = canonical_grid_size(module_dps) + _grid, _ = _find_rank_module(cfg.model._grids) + dp_rank, dp_size, dp_group_gloo, dp_group_nccl = build_module_dp_process_groups( + _grid.get_pg(["dp"]), overlap=cfg.dataset.megatron_mimo_reorder_overlap + ) + # Cost = encoder_cost_weight * patches + language_cost_weight * real_tokens. The patch count comes from the + # image-placeholder token count in input_ids (present on every module and PP stage), not + # from grid_thw, which the metadata collate omits on language shards. + cost_of = functools.partial( + sample_cost, + encoder_cost_weight=cfg.dataset.megatron_mimo_reorder_encoder_cost_weight, + language_cost_weight=cfg.dataset.megatron_mimo_reorder_language_cost_weight, + image_token_id=spec.image_token_id, + square_merge_size=spec.square_merge_size, + ) + + def image_count_of(b: dict[str, Any]) -> torch.Tensor: + # One vision-start token precedes each image, so this counts images, not patches. + return (b["input_ids"] == spec.vision_start_token_id).sum(dim=1).to(torch.long) + + loader_iter = ReorderingBuffer( + loader_iter, + dp_rank=dp_rank, + dp_size=dp_size, + n_groups=balance_n_groups, + cost_of=cost_of, + dp_group_gloo=dp_group_gloo, + dp_group_nccl=dp_group_nccl, + overlap=cfg.dataset.megatron_mimo_reorder_overlap, + image_count_of=image_count_of, + window_size=cfg.dataset.megatron_mimo_reorder_window_size, + ) if args.log_batches: loader_iter = _wrap_iter_logging(loader_iter, spec) return loader_iter, None @@ -1264,6 +1319,32 @@ def _parse_args() -> argparse.Namespace: "DP). Each rank processes its natural, unbalanced shard. Uses the same DP loss reduction as " "non-scalable runs.", ) + parser.add_argument( + "--intra-microbatch-reorder", + action="store_true", + help="Intra-microbatch reordering (requires --scalable-dp): rebalance each micro-batch's per-sample " + "vision load across the module DP group by a per-sample all-to-all so no rank is the straggler.", + ) + parser.add_argument( + "--reorder-encoder-cost-weight", + type=float, + default=1.0, + help="Reorder cost weight of the encoder term (per image patch).", + ) + parser.add_argument( + "--reorder-language-cost-weight", + type=float, + default=0.0, + help="Reorder cost weight of the language term (per sequence token, image tokens included); 0 = encoder-only.", + ) + parser.add_argument( + "--no-overlap-intra-microbatch-reorder", + action="store_true", + help="Run the reorder exchange synchronously instead of prefetching the next window on a side thread.", + ) + parser.add_argument( + "--reorder-window-size", type=int, default=1, help="Micro-batches exchanged together as one window." + ) parser.add_argument("--profile", choices=("none", "nsys", "pytorch"), default="none") parser.add_argument("--profile-step-start", type=int, default=1) parser.add_argument("--profile-step-end", type=int, default=2) diff --git a/src/megatron/bridge/data/builders/direct_hf_sft.py b/src/megatron/bridge/data/builders/direct_hf_sft.py index 3b4a1af3471..2a81b09cf5b 100644 --- a/src/megatron/bridge/data/builders/direct_hf_sft.py +++ b/src/megatron/bridge/data/builders/direct_hf_sft.py @@ -94,6 +94,13 @@ class DirectHFSFTDatasetConfig(DataloaderConfig): pad_to_multiple_of: int = 128 in_batch_packing_pad_to_multiple_of: int = 1 megatron_mimo_scalable_dp: bool = False + megatron_mimo_intra_microbatch_reorder: bool = False + # Reorder cost = encoder_weight * encoder work (image patches) + language_weight * sequence tokens + # (image placeholder tokens included); only the ratio matters. + megatron_mimo_reorder_encoder_cost_weight: float = 1.0 + megatron_mimo_reorder_language_cost_weight: float = 0.0 + megatron_mimo_reorder_overlap: bool = True + megatron_mimo_reorder_window_size: int = 1 def validate(self) -> None: """Validate declarative source and dataset settings.""" @@ -125,6 +132,27 @@ def validate(self) -> None: "megatron_mimo_scalable_dp requires drop_last=True; a partial final micro-batch " "gives modules unequal shares and misaligns the modality routing." ) + if self.megatron_mimo_intra_microbatch_reorder: + if not self.megatron_mimo_scalable_dp: + raise ValueError( + "megatron_mimo_intra_microbatch_reorder requires megatron_mimo_scalable_dp=True; the " + "per-sample exchange rebalances canonical-grid shards." + ) + if ( + self.megatron_mimo_reorder_encoder_cost_weight < 0 + or self.megatron_mimo_reorder_language_cost_weight < 0 + ): + raise ValueError("megatron_mimo_reorder_*_cost_weight must be >= 0.") + if ( + self.megatron_mimo_reorder_encoder_cost_weight == 0 + and self.megatron_mimo_reorder_language_cost_weight == 0 + ): + raise ValueError( + "megatron_mimo_intra_microbatch_reorder requires megatron_mimo_reorder_encoder_cost_weight > 0 " + "or megatron_mimo_reorder_language_cost_weight > 0; a zero cost has nothing to balance." + ) + if self.megatron_mimo_reorder_window_size < 1: + raise ValueError("megatron_mimo_reorder_window_size must be >= 1.") validate_declarative_mapping(self.hf_processor_kwargs, field_name="hf_processor_kwargs") if self.hf_processor_kwargs is not None and "trust_remote_code" in self.hf_processor_kwargs: raise ValueError( diff --git a/src/megatron/bridge/data/megatron_mimo/dp_utils.py b/src/megatron/bridge/data/megatron_mimo/dp_utils.py index 1e687f02e17..a38c61da965 100644 --- a/src/megatron/bridge/data/megatron_mimo/dp_utils.py +++ b/src/megatron/bridge/data/megatron_mimo/dp_utils.py @@ -302,3 +302,34 @@ def slice_batch_for_megatron_mimo( sliced[key] = value return sliced + + +def real_token_lengths(input_ids: torch.Tensor, *, attention_mask: "torch.Tensor | None") -> torch.Tensor: + """Per-sample real (non-pad) token length for a ``[B, S]`` batch. + + The length comes only from ``attention_mask`` (1 = real, 0 = pad), the padding mask the + tokenizer ships with the batch. Comparing ``input_ids`` against ``pad_token_id`` is not used: + tokenizers without a dedicated pad token alias ``pad_token = eos_token``, so a value-based + check would also drop genuine EOS tokens. ``loss_mask`` is not used either: it is a + supervision mask (prompt tokens are zero) and would under-count the real length. + + Args: + input_ids: Padded token ids ``[B, S]``. + attention_mask: ``[B, S]`` padding mask matching ``input_ids``. + + Returns: + An ``int64`` tensor of shape ``[B]`` with each sample's real length. + + Raises: + ValueError: If ``attention_mask`` is missing or does not match ``input_ids``. + """ + if not ( + isinstance(attention_mask, torch.Tensor) + and attention_mask.dim() == 2 + and attention_mask.shape == input_ids.shape + ): + raise ValueError( + "real_token_lengths requires a [B, S] attention_mask matching input_ids; " + f"got {None if attention_mask is None else tuple(attention_mask.shape)} for input_ids {tuple(input_ids.shape)}." + ) + return attention_mask.to(torch.bool).sum(dim=1).to(torch.long) diff --git a/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py b/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py new file mode 100644 index 00000000000..816c337dbce --- /dev/null +++ b/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py @@ -0,0 +1,1539 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-sample reorder buffer for MegatronMIMO intra-microbatch reordering (post-batch). + +Each MegatronMIMO data rank reads only its **disjoint** scalable-data-parallel shard of the global +micro-batch, and the per-microbatch cost imbalance between shards is corrected by a +**communication** step (cost all-gather + ragged sample all-to-all over the module DP +group) that can be overlapped with compute — instead of every rank reading the full +global micro-batch and slicing it locally. + +Contents, bottom-up: + +1. Pure, communication-free building blocks (all CPU-unit-testable without a process + group): + + - :func:`balanced_assignment` — the per-sample plan: for each global sample in a + micro-batch, the module-local DP rank that should own it and its canonical position. + Built on :func:`megatron.bridge.data.datasets.packing_utils.balanced_index_order`, + with het-DP handled via canonical ``n_groups = canonical_grid_size(module dp_sizes)`` (their LCM) (so the vision + and language modules, which may have different DP sizes, derive an identical + assignment and the ``BridgeCommunicator`` keeps vision replica *r* paired with + language replica *r*). + - Ragged (de)serialization of a sample's tensors to an 8-byte-aligned ``uint8`` buffer + (:func:`serialize_sample` / :func:`deserialize_sample`), so any dtype/shape mix can + ride a single ``all_to_all_single`` with per-rank byte splits; the receiver + reconstructs from metadata alone. + - :func:`prepare_sample_exchange` — from this rank's local samples plus the all-gathered + plan and metadata, partition kept-vs-exchanged samples, pack the contiguous send buffer, + and build the ``send_splits``/``recv_splits`` + ordered recv schedule for one + ``all_to_all_single``. + +2. The MIMO data plane: :func:`split_microbatch` / :func:`merge_samples` convert a nested + micro-batch to and from ``B`` per-sample flat (dotted-key) dicts, and :func:`sample_cost` + is the (collation-independent) per-sample cost — ``encoder_cost_weight·patches + language_cost_weight·real_tokens`` — + that drives the balanced assignment. + +3. The distributed glue: :func:`exchange_window` ties the above into one cost-balancing round + over a window of micro-batches on a module DP group, :class:`ReorderingBuffer` wraps a data + iterator to run that round per window (optionally on a background prefetch thread with a dedicated CUDA + stream + side NCCL PG, so step ``t+1``'s exchange overlaps step ``t``'s compute), and + :func:`build_module_dp_process_groups` creates the Gloo (metadata) and side NCCL + (exchange) process groups the buffer needs. + +The unit of exchange is a **single sample** (finer than a whole-microbatch swap), and the +assignment reuses the existing MIMO cost-balanced ordering rather than a separate +group-then-LPT pass. +""" + +import logging +import math +import os +import queue +import threading +from typing import Any, Dict, List, Tuple + +import torch + +from megatron.bridge.data.megatron_mimo.dp_utils import ( + _batch_dim_for_tensor, + _is_patch_packed_visual_dict, + real_token_lengths, +) + + +logger = logging.getLogger(__name__) + +# All sample tensors are serialized into uint8 byte buffers padded to this alignment, so +# the receiver can view-cast any dtype back from the received bytes without alignment +# errors and compute exact offsets from metadata alone. +_ALIGN = 8 + + +def _pad_to_align(nbytes: int) -> int: + """Zero-padding bytes to round ``nbytes`` up to the next :data:`_ALIGN` multiple.""" + return (-nbytes) % _ALIGN + + +# String -> (torch.dtype, byte width), for reconstructing tensors from metadata shared via +# all_gather_object. The itemsize is precomputed once so the (de)serialize byte-size loops do not +# allocate a throwaway tensor per field just to read ``element_size()``. +_DTYPE_INFO = { + name: (dt, torch.empty((), dtype=dt).element_size()) + for name, dt in { + "torch.float32": torch.float32, + "torch.float16": torch.float16, + "torch.bfloat16": torch.bfloat16, + "torch.int64": torch.int64, + "torch.int32": torch.int32, + "torch.bool": torch.bool, + "torch.uint8": torch.uint8, + }.items() +} + + +def _dtype_info(dtype_name: str) -> Tuple[torch.dtype, int]: + """Return ``(dtype, byte width)`` for a dtype encoded in sample metadata, rejecting unknown values.""" + try: + return _DTYPE_INFO[dtype_name] + except KeyError as exc: + raise ValueError(f"Unsupported tensor dtype in MegatronMIMO sample metadata: {dtype_name!r}.") from exc + + +# --------------------------------------------------------------------------- +# Per-sample balanced assignment +# --------------------------------------------------------------------------- + + +def balanced_index_order(costs: List[float], n_groups: int) -> List[int]: + """Permute indices so contiguous equal-size groups have balanced total cost. + + For intra-microbatch reordering when the downstream distributor slices the global + batch into ``n_groups`` **contiguous, equal-size** shards (e.g. + :func:`megatron.bridge.data.megatron_mimo.dp_utils.slice_batch_for_megatron_mimo`). + Returns a permutation of ``range(len(costs))`` laid out group-by-group, so that + taking contiguous block ``r`` (size ``len(costs)//n_groups``) yields a + cost-balanced shard for DP rank ``r``. + + Equal-cardinality LPT: items in decreasing cost are each placed into the + lowest-cost group that still has capacity (``len(costs)//n_groups``). Deterministic + (stable tie-break by group index) so paired modules reorder identically. + + Args: + costs: Per-item cost. + n_groups: Number of contiguous shards (e.g. DP size). Must divide ``len(costs)``. + + Returns: + A permutation of ``range(len(costs))``; the concatenation of the ``n_groups`` + balanced groups in order. + + Raises: + ValueError: If ``n_groups <= 0`` or does not divide ``len(costs)``. + """ + n = len(costs) + if n_groups <= 0: + raise ValueError(f"n_groups must be positive, got {n_groups}.") + if n % n_groups != 0: + raise ValueError(f"len(costs) ({n}) must be divisible by n_groups ({n_groups}).") + cap = n // n_groups + groups: List[List[int]] = [[] for _ in range(n_groups)] + group_cost = [0.0] * n_groups + for i in sorted(range(n), key=lambda j: (-costs[j], j)): + # lowest-cost group that still has capacity (tie-break by group index) + best = min( + (g for g in range(n_groups) if len(groups[g]) < cap), + key=lambda g: (group_cost[g], g), + ) + groups[best].append(i) + group_cost[best] += costs[i] + return [i for g in groups for i in g] + + +def balanced_assignment(costs: List[float], n_groups: int, dp_size: int) -> List[Tuple[int, int]]: + """Per-sample exchange plan for one micro-batch. + + Runs the canonical cost-balanced ordering over ``n_groups`` equal-cardinality groups, + then maps each global sample to the **module-local DP rank** that should own it and to + its **canonical position** in the globally balanced order. + + The owner mapping uses canonical groups (``n_groups = canonical_grid_size(module dp_sizes)`` (their LCM)): a module + of size ``dp_size`` covers ``n_groups // dp_size`` contiguous canonical groups, so + ``owner = canonical_group // (n_groups // dp_size)``. Because the vision and language + modules feed the same ``costs`` and ``n_groups``, they produce an identical plan and + paired replicas end up with the same samples in the same (canonical) order — required + by the ``BridgeCommunicator`` batch-dim fan-out. + + Args: + costs: Per-sample cost, indexed by the sample's global position in the micro-batch + (the patch-only cost from :func:`sample_cost`). + n_groups: Canonical group count (LCM of the module DP sizes). Must divide ``len(costs)``. + dp_size: This module's DP size. Must divide ``n_groups``. + + Returns: + For each global sample index ``g``: ``(owner_rank, canonical_pos)`` where + ``owner_rank`` is in ``[0, dp_size)`` and ``canonical_pos`` is in ``[0, len(costs))``. + Sorting a rank's owned samples by ``canonical_pos`` yields the order paired modules + agree on. + + Raises: + ValueError: If ``n_groups`` does not divide ``len(costs)`` or ``dp_size`` does not + divide ``n_groups``. + """ + n = len(costs) + if n_groups <= 0 or n % n_groups != 0: + raise ValueError(f"n_groups ({n_groups}) must be positive and divide len(costs) ({n}).") + if dp_size <= 0 or n_groups % dp_size != 0: + raise ValueError(f"dp_size ({dp_size}) must be positive and divide n_groups ({n_groups}).") + + perm = balanced_index_order(costs, n_groups) # global indices laid out group-by-group + cap = n // n_groups + groups_per_rank = n_groups // dp_size + + assignment: List[Tuple[int, int]] = [(0, 0)] * n + for canonical_pos, global_idx in enumerate(perm): + canonical_group = canonical_pos // cap + owner_rank = canonical_group // groups_per_rank + assignment[global_idx] = (owner_rank, canonical_pos) + return assignment + + +def intra_route(assignment: List[Tuple[int, int]], *, src_slot: int = 0) -> List[Tuple[int, int, int]]: + """Lift a single-slot 2D balanced assignment to the vehicle's 3D route. + + The inter-ready vehicle routes every sample by **three** coordinates + ``route[g] = (dst_slot, owner_rank, canonical_pos)`` (the destination micro-batch slot, the + module-local DP rank that owns it, and its canonical intra-slot position). The **intra** + balancer never moves a sample across micro-batch slots, so ``dst_slot == src_slot`` for every + sample; a future inter (cross-slot) balancer fills ``dst_slot`` freely behind this same seam. + + Args: + assignment: Per-slot output of :func:`balanced_assignment` — ``(owner_rank, canonical_pos)`` + per global sample index within the slot. + src_slot: The slot these samples were read from (and, for intra, stay in). Defaults to ``0`` + (the single-slot ``window_size == 1`` case). + + Returns: + For each global sample index ``g``: ``(dst_slot, owner_rank, canonical_pos)`` with + ``dst_slot == src_slot``. + """ + return [(src_slot, owner_rank, canonical_pos) for owner_rank, canonical_pos in assignment] + + +def build_window_route(costs_per_slot: List[List[float]], n_groups: int, dp_size: int) -> List[Tuple[int, int, int]]: + """Concatenate the per-slot intra assignments of a window into one 3D route. + + The window holds ``W = len(costs_per_slot)`` micro-batch slots, each with ``B_global`` samples + contiguous-sharded across ``dp_size`` ranks. The **intra** balancer runs **independently per + slot** (each slot cost-balanced across ranks); slots never mix. The window global index of slot + ``s``'s slot-local sample ``j`` is ``g = s · B_global + j`` (slots laid out back-to-back), so + concatenating each slot's :func:`intra_route` (stamped with ``src_slot = s``) in slot order + yields ``route[g]`` directly. ``dst_slot == src_slot == s`` for every sample (intra invariant). + + A future inter (cross-slot) balancer replaces this builder with one that fills ``dst_slot`` + across slots; the transport (:func:`exchange_by_route`) is unchanged. + + Args: + costs_per_slot: ``costs_per_slot[s]`` = slot ``s``'s per-sample costs indexed by slot-local + global position ``j`` (length ``B_global``, identical on paired modules). + n_groups: Canonical group count (LCM of the module DP sizes); must divide each ``B_global``. + dp_size: This module's DP size. + + Returns: + ``route`` of length ``W · B_global`` — ``(dst_slot, owner_rank, canonical_pos)`` per window + global index, slot-major. + """ + route: List[Tuple[int, int, int]] = [] + for slot, costs in enumerate(costs_per_slot): + route.extend(intra_route(balanced_assignment(costs, n_groups, dp_size), src_slot=slot)) + # §10.0 safety net: the intra balancer must never move a sample across slots. Cheap, always-on. + assert_intra_no_cross_slot(route, len(costs_per_slot[0]) if costs_per_slot else 0) + return route + + +def assert_intra_no_cross_slot(route: List[Tuple[int, int, int]], b_global: int) -> None: + """§10.0 "no cross-slot leak" guard: every sample's ``dst_slot`` equals its source slot. + + For the slot-major window layout (``g = src_slot · b_global + j``) the source slot of global + index ``g`` is ``g // b_global``. The **intra** balancer keeps every sample in its own slot, so + a mismatch means slot plumbing (route build / binning) corrupted the slot dimension — the §12 + "primary safety net while W>1 plumbing is new". A future inter (cross-slot) balancer would route + across slots by design and must **not** run this guard. + + Args: + route: The window route (``(dst_slot, owner_rank, canonical_pos)`` per global index). + b_global: Samples per micro-batch slot (``local · dp_size``); ``0`` skips the check. + + Raises: + RuntimeError: If any sample's ``dst_slot`` differs from its source slot. + """ + if b_global <= 0: + return + for g, (dst_slot, _owner_rank, _pos) in enumerate(route): + src_slot = g // b_global + if dst_slot != src_slot: + raise RuntimeError( + f"MegatronMIMO intra route cross-slot leak at global index {g}: " + f"dst_slot={dst_slot} != src_slot={src_slot} (the intra balancer must keep dst_slot == src_slot)." + ) + + +def window_cost_spread( + costs_per_slot: List[List[float]], route: List[Tuple[int, int, int]], dp_size: int +) -> List[Dict[str, float]]: + """Per-slot pre/post-balance per-rank cost spread + remote-sample count (the §10.5 balance probe). + + For each slot, totals the per-rank cost **before** balancing (the natural contiguous shard: rank + ``r`` holds slot-local positions ``[r·local, (r+1)·local)``) and **after** balancing (each sample + credited to its routed ``owner_rank``), plus how many samples change rank. Pure and cheap — used + only to log how much the reorder tightens the per-rank load (``after`` max/min should narrow + toward 1.0), the quantitative balance evidence the design's §10.3 validation looks for. + + Args: + costs_per_slot: ``costs_per_slot[s]`` = slot ``s``'s per-sample costs by slot-local position. + route: The window route (:func:`build_window_route`), indexed by ``g = s·B_global + j``. + dp_size: This module's DP size. + + Returns: + One dict per slot: ``{before_max, before_min, after_max, after_min, remote}``. + """ + spreads: List[Dict[str, float]] = [] + for slot, costs in enumerate(costs_per_slot): + n = len(costs) + local = n // dp_size + before = [0.0] * dp_size + after = [0.0] * dp_size + remote = 0 + for j, c in enumerate(costs): + src_rank = j // local + before[src_rank] += c + _dst_slot, owner_rank, _pos = route[slot * n + j] + after[owner_rank] += c + if owner_rank != src_rank: + remote += 1 + spreads.append( + { + "before_max": max(before), + "before_min": min(before), + "after_max": max(after), + "after_min": min(after), + "remote": float(remote), + } + ) + return spreads + + +# --------------------------------------------------------------------------- +# Ragged sample (de)serialization +# --------------------------------------------------------------------------- + + +def sample_keys(meta: Dict[str, Any]) -> List[str]: + """Sorted non-None tensor keys for a sample, computed identically on both sides.""" + return sorted(k for k, v in meta.items() if isinstance(v, dict) and "dtype" in v) + + +def tensor_metadata(flat: Dict[str, Any]) -> Dict[str, Any]: + """Per-field metadata for a flattened sample, keyed identically on both sides. + + Each key maps to ``{shape, dtype}`` for a tensor, ``{non_tensor: value}`` for a non-None + scalar/global field, or ``None``. Shared across ranks via ``all_gather_object`` so a + receiver can pre-allocate buffers and reconstruct the sample without seeing the data. + + A tensor whose dtype is not in :data:`_DTYPE_INFO` is rejected **here** — at metadata creation, + before the all-gather — via the same :func:`_dtype_info` check the deserializer uses, + so an unsupported dtype fails loudly and uniformly on the producing rank rather than only + surfacing later as a byte-size/deserialize ``ValueError`` on the receiver. + + Raises: + ValueError: If a tensor field has a dtype not supported by the ragged (de)serializer. + """ + meta: Dict[str, Any] = {} + for key, t in flat.items(): + if isinstance(t, torch.Tensor): + dtype_name = str(t.dtype) + _dtype_info(dtype_name) # reject unsupported dtypes early (same check as decode) + meta[key] = {"shape": list(t.shape), "dtype": dtype_name} + elif t is not None: + meta[key] = {"non_tensor": t} + else: + meta[key] = None + return meta + + +def serialize_sample(flat: Dict[str, Any], keys: List[str]) -> torch.Tensor: + """Serialize one sample's tensors into a 1D uint8 buffer (8-byte aligned, sorted keys). + + Args: + flat: Flattened tensor dict; ``None`` values are skipped. + keys: Sorted non-None tensor keys (from :func:`sample_keys`), same on both sides. + + Returns: + 1D uint8 CPU tensor with each tensor's raw bytes concatenated in ``keys`` order, + each padded to :data:`_ALIGN`. + """ + parts = [] + for key in keys: + t = flat.get(key) + if t is None: + continue + raw = t.contiguous().view(torch.uint8).reshape(-1) + pad = _pad_to_align(raw.numel()) + if pad: + raw = torch.cat([raw, torch.zeros(pad, dtype=torch.uint8, device=raw.device)]) + parts.append(raw) + return torch.cat(parts) if parts else torch.empty(0, dtype=torch.uint8) + + +def sample_byte_size(meta: Dict[str, Any], keys: List[str]) -> int: + """Byte size of a serialized sample from metadata alone (matches :func:`serialize_sample`).""" + total = 0 + for key in keys: + info = meta.get(key) + if info is None: + continue + _, elem = _dtype_info(info["dtype"]) + nbytes = math.prod(info["shape"]) * elem + nbytes += _pad_to_align(nbytes) + total += nbytes + return total + + +def deserialize_sample( + buf: torch.Tensor, offset: int, meta: Dict[str, Any], keys: List[str] +) -> Tuple[Dict[str, Any], int]: + """Inverse of :func:`serialize_sample`: reconstruct tensors from a uint8 buffer. + + Args: + buf: Received 1D uint8 buffer. + offset: Starting byte offset of this sample in ``buf``. + meta: This sample's field metadata (from :func:`tensor_metadata`: + ``{key: {shape, dtype} | {non_tensor: value} | None}``). + keys: Sorted non-None tensor keys (from :func:`sample_keys`). + + Returns: + ``(flat, new_offset)`` — the reconstructed tensor dict (tensors cloned off ``buf``, + with ``None`` and non-tensor keys restored from ``meta``) and the byte offset just + past this sample. + """ + flat: Dict[str, Any] = {} + cursor = offset + for key in keys: + info = meta.get(key) + if info is None: + flat[key] = None + continue + dtype, elem = _dtype_info(info["dtype"]) + nbytes = math.prod(info["shape"]) * elem + flat[key] = buf[cursor : cursor + nbytes].view(dtype).reshape(info["shape"]).clone() + cursor += nbytes + _pad_to_align(nbytes) + for key, info in meta.items(): + if key not in flat: + flat[key] = info.get("non_tensor") if isinstance(info, dict) and "non_tensor" in info else None + return flat, cursor + + +# --------------------------------------------------------------------------- +# All-to-all plan builder (one micro-batch, one module DP group) +# --------------------------------------------------------------------------- + + +def prepare_sample_exchange( + local_flats: List[Dict[str, Any]], + local_global_indices: List[int], + route: List[Tuple[int, int, int]], + all_global_indices: List[List[int]], + all_tensor_meta: List[List[Dict[str, Any]]], + dp_rank: int, + dp_size: int, + *, + window_size: int = 1, +) -> Dict[str, Any]: + """Prepare everything one per-sample ``all_to_all_single`` needs, in three parts. + + Pure and communication-free: all cross-rank information (which rank currently holds which + global sample, and each sample's tensor metadata) is passed in, having been shared by a + prior ``all_gather_object``. This single pass produces three distinct things: + + 1. **Partition** — split this rank's samples into those it already owns (kept locally) and + those it must send/receive (``local_samples`` + the send/recv schedules). + 2. **Pack** — serialize the outgoing samples into one contiguous ``send_buf`` ordered by + destination rank. + 3. **Plan** — the per-rank byte ``send_splits``/``recv_splits`` and the ordered + ``recv_schedule`` that drive (and decode) the ``all_to_all_single``. + + Each sample carries its destination micro-batch slot (``dst_slot``) end-to-end so the receiver + can bin it into the right slot bucket on reassembly (:func:`reassemble_window`); this is the + only genuinely new plumbing vs. the single-slot path (§4a of the inter-ready design). Sender + and receiver order the moves identically by ``(dst_slot, canonical_pos, src_local_idx)`` so the + packed bytes line up. + + Args: + local_flats: This rank's flattened samples (tensor dicts), in local order. + local_global_indices: Global index of each entry in ``local_flats`` (same length). + route: Output of :func:`intra_route` (or a future inter balancer) — + ``(dst_slot, owner_rank, canonical_pos)`` per global sample index. + all_global_indices: ``all_global_indices[r]`` = the global indices currently held by + rank ``r`` (the all-gathered ``local_global_indices``). + all_tensor_meta: ``all_tensor_meta[r][i]`` = tensor metadata of rank ``r``'s ``i``-th + local sample (the all-gathered :func:`tensor_metadata`). + dp_rank: This rank's module-local DP rank. + dp_size: This module's DP size. + window_size: Number of micro-batch slots ``W`` in the window; ``dst_slot`` must be in + ``[0, W)``. Defaults to ``1`` (the single-slot, byte-for-byte-with-intra case). + + Returns: + Dict with: + - (partition) ``local_samples``: ``[(dst_slot, canonical_pos, local_idx)]`` kept on this + rank. + - (pack) ``send_buf``: 1D uint8 tensor on the same device as the sample tensors (GPU + when CUDA is available), outgoing samples serialized contiguously by dst. + - (plan) ``send_splits`` / ``recv_splits``: per-rank byte counts for + ``all_to_all_single`` (``0`` for self). + - (plan) ``recv_schedule``: ``{src_rank: [(dst_slot, canonical_pos, src_local_idx)]}`` in + recv order, so the receiver deserializes ``recv_buf`` in the same order it was packed. + """ + # Map global index -> (current owner rank, current local index) from the all-gather. + location: Dict[int, Tuple[int, int]] = {} + for r, gidxs in enumerate(all_global_indices): + for i, g in enumerate(gidxs): + location[g] = (r, i) + + # Samples this rank must end up with: every global index whose owner is dp_rank. + local_samples: List[Tuple[int, int, int]] = [] + send_schedule: Dict[int, List[Tuple[int, int, int]]] = {r: [] for r in range(dp_size)} + recv_schedule: Dict[int, List[Tuple[int, int, int]]] = {r: [] for r in range(dp_size)} + + local_pos = {g: i for i, g in enumerate(local_global_indices)} + + for global_idx, (dst_slot, owner_rank, canonical_pos) in enumerate(route): + if not 0 <= dst_slot < window_size: + raise ValueError(f"dst_slot ({dst_slot}) for global sample {global_idx} out of range [0, {window_size}).") + src_rank, src_local_idx = location[global_idx] + if owner_rank == src_rank: + # Stays put; only this rank records it (when it is the owner/holder). + if owner_rank == dp_rank: + local_samples.append((dst_slot, canonical_pos, local_pos[global_idx])) + continue + if src_rank == dp_rank: + send_schedule[owner_rank].append((dst_slot, canonical_pos, src_local_idx)) + if owner_rank == dp_rank: + recv_schedule[src_rank].append((dst_slot, canonical_pos, src_local_idx)) + + # Deterministic order (by slot then canonical position) so sender and receiver agree. + for r in range(dp_size): + send_schedule[r].sort() + recv_schedule[r].sort() + + # Device of the (possibly GPU-resident) sample tensors, so the send buffer and any empty + # placeholders live on the same device for torch.cat / all_to_all_single. + dev = torch.device("cpu") + for f in local_flats: + t = next((v for v in f.values() if isinstance(v, torch.Tensor)), None) + if t is not None: + dev = t.device + break + + # Build the contiguous send buffer ordered by destination rank. + send_parts: List[torch.Tensor] = [] + send_splits: List[int] = [] + for dst in range(dp_size): + if dst == dp_rank or not send_schedule[dst]: + send_splits.append(0) + continue + chunk_parts = [] + for _dst_slot, _canonical_pos, src_local_idx in send_schedule[dst]: + keys = sample_keys(all_tensor_meta[dp_rank][src_local_idx]) + chunk_parts.append(serialize_sample(local_flats[src_local_idx], keys)) + chunk = torch.cat(chunk_parts) if chunk_parts else torch.empty(0, dtype=torch.uint8, device=dev) + send_parts.append(chunk) + send_splits.append(int(chunk.numel())) + + # Recv byte counts from metadata alone. + recv_splits: List[int] = [] + for src in range(dp_size): + if src == dp_rank or not recv_schedule[src]: + recv_splits.append(0) + continue + total = 0 + for _dst_slot, _canonical_pos, src_local_idx in recv_schedule[src]: + keys = sample_keys(all_tensor_meta[src][src_local_idx]) + total += sample_byte_size(all_tensor_meta[src][src_local_idx], keys) + recv_splits.append(total) + + send_buf = torch.cat(send_parts) if send_parts else torch.empty(0, dtype=torch.uint8, device=dev) + return { + "local_samples": local_samples, + "send_buf": send_buf, + "send_splits": send_splits, + "recv_splits": recv_splits, + "recv_schedule": {r: v for r, v in recv_schedule.items() if v}, + } + + +# --------------------------------------------------------------------------- +# MIMO data plane: micro-batch <-> per-sample flat dicts (flatten_fn / rewrap_fn) +# --------------------------------------------------------------------------- +# +# The exchange unit is a single sample. A MegatronMIMO micro-batch is nested: +# +# {input_ids: [B,S], position_ids: [3,B,S] (MRoPE), labels/loss_mask: [B,S] | None, +# attention_mask: None, modality_inputs: {modality: {encoder: {hidden_states: [ΣP,d], +# grid_thw: [n_img,3]}}}} +# +# ``split_microbatch`` turns it into B per-sample **flat** dicts (nesting flattened to +# dotted keys) so the Phase-0 serializer/all-to-all can move individual samples; after the +# exchange ``merge_samples`` rebuilds the (rebalanced) micro-batch. Sample slicing/vision +# attribution uses :func:`_apply_sample_dispatch` (LM by batch dim, vision by image via +# ``image_counts``). + + +def _flatten_nested(d: Dict[str, Any], prefix: str = "") -> Dict[str, Any]: + """Flatten a (possibly nested) batch dict to dotted keys; tensors and ``None`` preserved.""" + flat: Dict[str, Any] = {} + for key, val in d.items(): + path = f"{prefix}{key}" + if isinstance(val, dict): + flat.update(_flatten_nested(val, prefix=f"{path}.")) + else: + # tensors and None pass through; the buffer's batch has no per-sample list fields. + flat[path] = val + return flat + + +def _unflatten_nested(flat: Dict[str, Any]) -> Dict[str, Any]: + """Inverse of :func:`_flatten_nested`: rebuild the nested dict from dotted keys.""" + out: Dict[str, Any] = {} + for path, val in flat.items(): + parts = path.split(".") + node = out + for p in parts[:-1]: + node = node.setdefault(p, {}) + node[parts[-1]] = val + return out + + +def _cu_img_from_counts(image_counts: "Any") -> List[int]: + """Cumulative per-sample image-offset table from per-sample image counts. + + ``image_counts`` is a length-``B`` int tensor/sequence; returns ``[0, c0, c0+c1, ...]`` (length + ``B + 1``) so sample ``s`` owns ``grid_thw`` rows ``[cu_img[s], cu_img[s + 1])``. + """ + counts = image_counts.tolist() if isinstance(image_counts, torch.Tensor) else list(image_counts) + cu = [0] + for c in counts: + cu.append(cu[-1] + int(c)) + return cu + + +def empty_like_vision( + ref_hidden_states: torch.Tensor, ref_grid_thw: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + """Empty vision tensors matching a reference sample's dtype/device. + + Returns ``([0, d] hidden_states, [0, 3] grid_thw)`` — the placeholder a **text-only** sample + (0 images) contributes so the per-sample reorder/merge keeps uniform vision keys across samples + and ranks (the merge concat and the cross-rank byte-split metadata both stay symmetric). Used by + :func:`merge_samples`. + """ + empty_hidden_states = ref_hidden_states.new_empty((0, ref_hidden_states.shape[1])) + empty_grid_thw = ref_grid_thw.new_empty((0, ref_grid_thw.shape[1])) + return empty_hidden_states, empty_grid_thw + + +def patches_per_image(grid_thw: torch.Tensor) -> torch.Tensor: + """Per-image vision patch count ``prod(t, h, w)`` for a ``[n_images, 3]`` grid. + + Centralizes the Qwen-VL "patches per image = ``prod(grid_thw[i])``" convention shared by + the vision reorder and cost paths so a future grid layout change is a single edit. + """ + return torch.prod(grid_thw, dim=1) + + +def _reorder_vision_by_images( + hidden_states: torch.Tensor, + grid_thw: torch.Tensor, + image_perm: list[int], +) -> Tuple[torch.Tensor, torch.Tensor]: + """Permute vision ``hidden_states`` (patch-dim) and ``grid_thw`` (image-dim) by image. + + The patch block of image ``i`` (rows ``cu[i]:cu[i+1]`` where ``cu`` is the cumulative + patch count from ``prod(grid_thw[i])``) is gathered in ``image_perm`` order so the two + stay aligned — the patch-dim equivalent of ``grid_thw[image_perm]``. Reorders whole images + (keeping each image's patch block intact) rather than slicing across image boundaries. + + Args: + hidden_states: Vision patch features ``[Σ patches, dim]``. + grid_thw: Per-image grid ``[n_images, 3]``. + image_perm: Permutation of ``range(n_images)``. + + Returns: + ``(hidden_states_perm, grid_thw_perm)``. + """ + if not image_perm: + # No images selected (e.g. a text-only sample): return empty ``[0, d]`` / ``[0, 3]`` views + # (avoids ``torch.cat([])`` and preserves dtype/device/requires_grad). + return hidden_states[:0], grid_thw[:0] + patches = patches_per_image(grid_thw) # [n_images] + # Single host copy of the cumulative offsets (F14): materialize the cumsum once with one + # .tolist(), then index with plain Python ints — no per-image .item() D2H sync in the loop. + cu = torch.cat([patches.new_zeros(1), patches.cumsum(0)]).tolist() + blocks = [hidden_states[cu[i] : cu[i + 1]] for i in image_perm] + perm_idx = torch.as_tensor(image_perm, dtype=torch.long, device=grid_thw.device) + return torch.cat(blocks, dim=0), grid_thw.index_select(0, perm_idx) + + +def _gather_vision_subdict( + value: Dict[str, Any], + sample_indices: list[int], + n_samples: int, + *, + cu_img: "list[int] | None" = None, +) -> Dict[str, Any]: + """Reorder/gather a vision ``{hidden_states, grid_thw}`` sub-dict by sample. + + Supports a **variable number of images per sample** (0 = text-only, 1, or N). ``cu_img`` is the + cumulative per-sample image-offset table (length ``n_samples + 1``), so sample ``s`` owns + ``grid_thw`` rows ``[cu_img[s], cu_img[s + 1])``. The image permutation for ``sample_indices`` is + the concatenation of each selected sample's image range, gathered by + :func:`_reorder_vision_by_images` (which keeps ``hidden_states`` patch-dim and ``grid_thw`` + image-dim aligned). A text-only sample contributes an empty ``[0, d]`` / ``[0, 3]`` block. + + When ``cu_img`` is ``None`` the legacy one-image-per-sample mapping (image ``i`` ↔ sample ``i``) + is used and requires ``n_images == n_samples``. + + Args: + value: A sub-dict carrying ``hidden_states`` and ``grid_thw``. + sample_indices: Per-sample permutation/selection (global sample ids). + n_samples: Global batch sample count ``B``. + cu_img: Cumulative per-sample image offsets (length ``B + 1``), or ``None`` for the legacy + one-image-per-sample path. + + Returns: + A new sub-dict with permuted ``hidden_states`` / ``grid_thw`` (other keys carried). + + Raises: + ValueError: If ``cu_img`` is ``None`` and the batch is not one-image-per-sample, or if + ``cu_img`` does not sum to ``n_images`` (mis-sourced image counts). + """ + hidden_states, grid_thw = value["hidden_states"], value["grid_thw"] + n_images = grid_thw.shape[0] + if cu_img is None: + if n_images != n_samples: + raise ValueError( + f"Vision reorder without per-sample image counts expects one image per sample " + f"(n_images={n_images}, n_samples={n_samples}); pass cu_img for variable counts." + ) + image_perm = list(sample_indices) + else: + if cu_img[-1] != n_images: + raise ValueError( + f"Per-sample image counts sum to {cu_img[-1]} but grid_thw has {n_images} images; " + "image_count_of is likely mis-sourced (wrong token id / off-by-one)." + ) + image_perm = [img for s in sample_indices for img in range(cu_img[s], cu_img[s + 1])] + + hs, g = _reorder_vision_by_images(hidden_states, grid_thw, image_perm) + new_sub = dict(value) + new_sub["hidden_states"] = hs + new_sub["grid_thw"] = g + return new_sub + + +def _apply_sample_dispatch( + batch: Dict[str, Any], + sample_indices: list[int], + *, + n_samples: int, + cu_img: "list[int] | None" = None, +) -> Dict[str, Any]: + """Shared per-sample gather/permute visitor for tensors, vision sub-dicts, and lists. + + Used by :func:`split_microbatch` (per-sample local-shard selection); ``n_samples`` is the + canonical global batch size ``B`` so list/vision per-sample detection is consistent. ``cu_img`` + (cumulative per-sample image offsets) is forwarded to :func:`_gather_vision_subdict` for the + variable-images-per-sample reorder; ``None`` keeps the legacy one-image-per-sample path. + """ + idx = torch.as_tensor(sample_indices, dtype=torch.long) + out: Dict[str, Any] = {} + for key, value in batch.items(): + if isinstance(value, torch.Tensor): + batch_dim = _batch_dim_for_tensor(key, value) + out[key] = value.index_select(batch_dim, idx.to(value.device)) + elif isinstance(value, dict): + if _is_patch_packed_visual_dict(value): + out[key] = _gather_vision_subdict(value, sample_indices, n_samples, cu_img=cu_img) + else: + out[key] = _apply_sample_dispatch(value, sample_indices, n_samples=n_samples, cu_img=cu_img) + elif isinstance(value, list): + # Per-sample list -> permuted gather; global-metadata / empty list -> passthrough. + if len(value) == n_samples: + out[key] = [value[i] for i in sample_indices] + else: + out[key] = value + else: + out[key] = value + return out + + +def split_microbatch(batch: Dict[str, Any], *, cu_img: "List[int] | None" = None) -> List[Dict[str, Any]]: + """Split a micro-batch into ``B`` per-sample flat (dotted-key) dicts. + + Each sample is gathered with :func:`_apply_sample_dispatch` (so its image travels with it) + and then flattened to dotted keys for serialization. Leading batch dims are kept (size 1), + so :func:`merge_samples` is a plain concat. + + ``cu_img`` (cumulative per-sample image offsets, length ``B + 1``) lets the vision gather support + a variable number of images per sample (0 = text-only, 1, or N). When ``None``, the legacy + one-image-per-sample path is used. + """ + b = _batch_size(batch) + return [_flatten_nested(_apply_sample_dispatch(batch, [i], n_samples=b, cu_img=cu_img)) for i in range(b)] + + +def merge_samples(flats: List[Dict[str, Any]]) -> Dict[str, Any]: + """Rebuild a micro-batch from per-sample flat dicts (inverse of :func:`split_microbatch`). + + LM tensors concat on their batch dim (:func:`dp_utils._batch_dim_for_tensor`); vision + ``hidden_states`` concat on the patch dim and ``grid_thw`` on the image dim; ``None`` and + scalar/global fields are taken from the first sample. + """ + if not flats: + raise ValueError("merge_samples requires at least one sample.") + nested = [_unflatten_nested(f) for f in flats] + return _merge_nested(nested) + + +def _merge_nested(samples: List[Dict[str, Any]]) -> Dict[str, Any]: + # Iterate the union of keys (not just ``samples[0]``) and classify a vision sub-dict by scanning + # *all* samples: a text-only sample (0 images) carries the sub-dict as empty/None, so keying off + # ``samples[0]`` alone would misclassify it as a plain nested dict and silently drop the other + # samples' vision. With the empty-tensor convention every sample presents the sub-dict as a + # tensor, so the empty substitution below is a defensive backstop. + keys: List[str] = [] + seen = set() + for s in samples: + for k in s: + if k not in seen: + seen.add(k) + keys.append(k) + + out: Dict[str, Any] = {} + for key in keys: + vals = [s.get(key) for s in samples] + vis_ref = next((v for v in vals if _is_patch_packed_visual_dict(v)), None) + if vis_ref is not None: + # Vision sub-dict: concat hidden_states (patch dim) + grid_thw (image dim), substituting + # an empty [0,d]/[0,3] for any text-only sample (None/absent) so the concat stays aligned. + ref_hidden_states, ref_grid_thw = vis_ref["hidden_states"], vis_ref["grid_thw"] + hidden_parts: List[torch.Tensor] = [] + grid_parts: List[torch.Tensor] = [] + for v in vals: + if _is_patch_packed_visual_dict(v): + hidden_parts.append(v["hidden_states"]) + grid_parts.append(v["grid_thw"]) + else: + empty_hidden, empty_grid = empty_like_vision(ref_hidden_states, ref_grid_thw) + hidden_parts.append(empty_hidden) + grid_parts.append(empty_grid) + sub = dict(vis_ref) + sub["hidden_states"] = torch.cat(hidden_parts, dim=0) + sub["grid_thw"] = torch.cat(grid_parts, dim=0) + out[key] = sub + continue + + rep = next((v for v in vals if v is not None), None) + if isinstance(rep, dict): + out[key] = _merge_nested([v for v in vals if isinstance(v, dict)]) + elif isinstance(rep, torch.Tensor): + out[key] = torch.cat(vals, dim=_batch_dim_for_tensor(key, rep)) + else: + out[key] = rep # None / scalar / global metadata + return out + + +def _batch_to_cuda(batch: Dict[str, Any]) -> Dict[str, Any]: + """Recursively move a (possibly nested) batch dict's tensors to CUDA (non-blocking).""" + out: Dict[str, Any] = {} + for key, val in batch.items(): + if isinstance(val, torch.Tensor): + out[key] = val.cuda(non_blocking=True) + elif isinstance(val, dict): + out[key] = _batch_to_cuda(val) + else: + out[key] = val + return out + + +def _batch_size(batch: Dict[str, Any]) -> int: + """Sample count of a micro-batch from a known LM tensor (handles MRoPE position_ids).""" + for key in ("input_ids", "labels", "loss_mask", "position_ids"): + t = batch.get(key) + if isinstance(t, torch.Tensor): + return t.size(_batch_dim_for_tensor(key, t)) + raise ValueError("Cannot determine batch size: no input_ids/labels/loss_mask/position_ids tensor.") + + +# --------------------------------------------------------------------------- +# Per-sample joint cost (from a flat sample) +# --------------------------------------------------------------------------- + + +def sample_cost( + flat: Dict[str, Any], + *, + encoder_cost_weight: float, + language_cost_weight: float = 0.0, + image_token_id: "int | None" = None, + square_merge_size: int = 1, +) -> float: + """Joint per-sample cost: ``encoder_cost_weight·p + language_cost_weight·t``. + + Only the ratio of the two weights matters (the balancer is scale-invariant). Each term is + linear in an intrinsic per-sample workload count: + + - ``p`` is the vision patch count (mirrors vision-encoder FLOPs). + - ``t`` is the **real (non-pad) token count** of ``input_ids`` (mirrors LM FLOPs), counted + via :func:`real_token_lengths` (from the batch's ``attention_mask``). Image placeholder + tokens are part of the language sequence and are included. + + Both terms must be **collation-independent** so vision and language derive an identical + assignment with no cross-module communication: the patch count is intrinsic to the image, + and the real token count is intrinsic to the sample (invariant to how each module's collator + pads). The *padded* sequence length is collation-dependent and is deliberately **not** used — + it would differ between the two shards and mispair the ``BridgeCommunicator`` fan-out. + + Patch-count source (must be **module-independent**): + + - When ``image_token_id`` is given, ``p`` is the image-placeholder token count in ``input_ids`` + scaled by ``square_merge_size`` (each placeholder token stands for ``spatial_merge_size²`` + patches in Qwen-VL). ``input_ids`` is present and identical on **every** module/PP stage when + reorder is active, so vision and language derive the same cost. This is the correct source: + the rank-aware metadata collate (#4442) nulls ``modality_inputs``/``grid_thw`` on language + shards, so reading ``grid_thw`` there yields ``p = 0`` and mispairs the vision↔language fan-out. + - Otherwise ``p = Σ prod(grid_thw)`` (the legacy patch-only path, used by mock/tests where no + ``image_token_id`` is wired and ``grid_thw`` is present on the sample). + + With the default ``language_cost_weight=0.0`` the cost is patch-only (the historical behavior). + + Args: + flat: One per-sample flat (dotted-key) dict. + encoder_cost_weight: Weight of the encoder term (per patch here; patch-grid encoders only, + other encoders need their own ``cost_of``). + language_cost_weight: Weight of the language term (per sequence token). ``0.0`` disables it. + image_token_id: Placeholder token id whose count in ``input_ids`` is the module-independent + patch proxy. ``None`` falls back to the ``grid_thw`` patch sum. + square_merge_size: ``spatial_merge_size²`` — patches represented by one placeholder token. + Recovers the true patch count from the placeholder-token count (only used with + ``image_token_id``). + """ + patches = 0 + if image_token_id is not None: + input_ids = flat.get("input_ids") + if isinstance(input_ids, torch.Tensor) and input_ids.numel(): + patches = int((input_ids == image_token_id).sum().item()) * square_merge_size + else: + for key, val in flat.items(): + if key.endswith(".grid_thw") and isinstance(val, torch.Tensor) and val.numel(): + patches += int(torch.prod(val, dim=1).sum().item()) + + tokens = 0 + if language_cost_weight: + input_ids = flat.get("input_ids") + if isinstance(input_ids, torch.Tensor) and input_ids.numel(): + attention_mask = flat.get("attention_mask") + lengths = real_token_lengths( + input_ids, + attention_mask=attention_mask if isinstance(attention_mask, torch.Tensor) else None, + ) + tokens = int(lengths.sum().item()) + + return encoder_cost_weight * patches + language_cost_weight * tokens + + +# --------------------------------------------------------------------------- +# Per-slot reassembly + transport-only route exchange (W micro-batch slots) +# --------------------------------------------------------------------------- + + +def reassemble_window( + local_flats: List[Dict[str, Any]], + plan: Dict[str, Any], + recv_buf: torch.Tensor, + all_tensor_meta: List[List[Dict[str, Any]]], + *, + dp_size: int, + window_size: int, + local: int, +) -> List[Dict[str, Any]]: + """Rebuild this rank's ``W`` micro-batches from kept-local + received samples (§4a). + + Generalizes the single-micro-batch reassembler to ``W`` slot-buckets: kept-local samples and + deserialized received samples are binned by ``dst_slot``, sorted **within each slot** by + ``canonical_pos`` (the paired-module-identical order), and merged into one micro-batch per slot + via :func:`merge_samples`. ``deserialize_sample`` / ``merge_samples`` are reused verbatim — the + only new step is the per-slot binning. + + The §10.0 invariant asserts are baked in (cheap, always on): they turn a silent slot mispair or + a dropped/duplicated sample into a loud failure at the exact rank. + + Args: + local_flats: This rank's input samples (the kept-local ones are indexed by ``local_idx``). + plan: Output of :func:`prepare_sample_exchange` (carries ``local_samples`` / ``recv_schedule`` + with ``dst_slot``). + recv_buf: The ``all_to_all_single`` output buffer (samples this rank received). + all_tensor_meta: All-gathered per-rank tensor metadata (to decode ``recv_buf``). + dp_size: This module's DP size. + window_size: Number of micro-batch slots ``W`` to rebuild. + local: Expected samples per (slot) bucket for this rank (``B/dp`` per slot). + + Returns: + ``W`` rebuilt micro-batches, slot ``s`` at index ``s``. + """ + buckets: List[List[Tuple[int, Dict[str, Any]]]] = [[] for _ in range(window_size)] + for dst_slot, canonical_pos, local_idx in plan["local_samples"]: + buckets[dst_slot].append((canonical_pos, local_flats[local_idx])) + + offset = 0 + for src in range(dp_size): + for dst_slot, canonical_pos, src_local_idx in plan["recv_schedule"].get(src, []): + meta = all_tensor_meta[src][src_local_idx] + flat, offset = deserialize_sample(recv_buf, offset, meta, sample_keys(meta)) + buckets[dst_slot].append((canonical_pos, flat)) + + # Byte accounting: the deserialize cursor must consume exactly the received bytes. + if offset != recv_buf.numel(): + raise RuntimeError(f"MegatronMIMO reassembly consumed {offset} bytes but received {recv_buf.numel()}.") + + out: List[Dict[str, Any]] = [] + for slot, bucket in enumerate(buckets): + # Slot shape: each slot bucket must hold exactly `local` samples. + if len(bucket) != local: + raise RuntimeError(f"MegatronMIMO slot {slot} reassembled {len(bucket)} samples, expected {local}.") + bucket.sort(key=lambda x: x[0]) + # Order: canonical positions within a slot must be unique (no dropped/duplicated sample). + positions = [p for p, _ in bucket] + if len(set(positions)) != len(positions): + raise RuntimeError(f"MegatronMIMO slot {slot} has duplicate canonical positions: {positions}.") + out.append(merge_samples([f for _pos, f in bucket])) + return out + + +def exchange_by_route( + local_flats: List[Dict[str, Any]], + route: List[Tuple[int, int, int]], + all_global_indices: List[List[int]], + all_tensor_meta: List[List[Dict[str, Any]]], + *, + dp_rank: int, + dp_size: int, + window_size: int, + local: int, + dp_group_nccl: "Any", + nccl_stream: "Any" = None, +) -> List[Dict[str, Any]]: + """Transport-only per-sample exchange: route -> one ``all_to_all_single`` -> ``W`` micro-batches. + + The inter-ready seam (§3.2): given a 3D ``route`` and the all-gathered layout/metadata, build the + plan (:func:`prepare_sample_exchange`), run **one** ragged ``all_to_all_single``, and reassemble + ``W`` micro-batches per slot (:func:`reassemble_window`). Knows nothing about the balancer — an + intra route (``dst_slot == src_slot``, see :func:`intra_route`) and a future inter route flow + through unchanged. + + Args: + local_flats: This rank's input samples (GPU-resident when CUDA is available). + route: ``(dst_slot, owner_rank, canonical_pos)`` per global sample index. + all_global_indices: ``all_global_indices[r]`` = global indices currently held by rank ``r``. + all_tensor_meta: All-gathered per-rank tensor metadata. + dp_rank, dp_size: Module-local DP coordinates. + window_size: Number of micro-batch slots ``W`` in the window. + local: Samples per (slot) bucket for this rank. + dp_group_nccl: NCCL PG for the sample all-to-all. + nccl_stream: Optional CUDA stream to run the all-to-all on (cross-step overlap). + + Returns: + ``W`` rebuilt micro-batches for this rank, slot ``s`` at index ``s``. + """ + import torch.distributed as dist + + plan = prepare_sample_exchange( + local_flats, + all_global_indices[dp_rank], + route, + all_global_indices, + all_tensor_meta, + dp_rank, + dp_size, + window_size=window_size, + ) + + # send_buf is already on GPU (built from GPU-resident sample tensors); no H2D copy. + send_buf = plan["send_buf"] + recv_buf = torch.empty(sum(plan["recv_splits"]), dtype=torch.uint8, device=send_buf.device) + if nccl_stream is not None: + ev = torch.cuda.Event() + ev.record(torch.cuda.current_stream()) + nccl_stream.wait_event(ev) + with torch.cuda.stream(nccl_stream): + dist.all_to_all_single(recv_buf, send_buf, plan["recv_splits"], plan["send_splits"], group=dp_group_nccl) + done = torch.cuda.Event() + done.record(nccl_stream) + torch.cuda.current_stream().wait_event(done) + else: + dist.all_to_all_single(recv_buf, send_buf, plan["recv_splits"], plan["send_splits"], group=dp_group_nccl) + + return reassemble_window( + local_flats, plan, recv_buf, all_tensor_meta, dp_size=dp_size, window_size=window_size, local=local + ) + + +# --------------------------------------------------------------------------- +# Distributed per-sample exchange (a window of W micro-batches, one module DP group) +# --------------------------------------------------------------------------- + + +def exchange_window( + batches: List[Dict[str, Any]], + *, + dp_rank: int, + dp_size: int, + n_groups: int, + cost_of: "Any", + dp_group_gloo: "Any", + dp_group_nccl: "Any", + nccl_stream: "Any" = None, + image_count_of: "Any" = None, + probe: bool = False, +) -> List[Dict[str, Any]]: + """Cost-balance a **window** of ``W`` micro-batches across a module DP group in one exchange. + + Collects this rank's ``W`` disjoint shards, all-gathers their per-sample costs + tensor metadata + over Gloo **once**, builds the slot-major window route (:func:`build_window_route` — the intra + balancer run independently per slot), and runs **one** ragged ``all_to_all_single`` over the + whole window, reassembling ``W`` balanced micro-batches per slot (:func:`reassemble_window`). + The single collective lands once per window (≈ once per optimizer step when ``W == GA``). + + The window global index of slot ``s``'s slot-local sample ``j`` is ``g = s · B_global + j``; this + rank holds, in ``local_flats`` order, its slot-0 shard then slot-1 shard … (``local = B/dp`` per + slot). Both modules feed identical ``n_groups`` and per-slot costs, so paired replicas end with + the same samples in canonical order in every slot. + + Args: + batches: This rank's ``W`` disjoint micro-batch shards (each ``local = B/dp`` samples, same + ``local`` across slots). + dp_rank, dp_size: Module-local DP coordinates. + n_groups: Canonical group count (LCM of the module DP sizes). + cost_of: ``callable(flat) -> float`` per-sample patch cost (see :func:`sample_cost`). + dp_group_gloo: Gloo PG for the metadata all-gather. + dp_group_nccl: NCCL PG for the sample all-to-all. + nccl_stream: Optional CUDA stream to run the all-to-all on (cross-step overlap). + image_count_of: Optional ``callable(batch) -> LongTensor[local]`` giving a shard's per-sample + image count, enabling a variable number of images per sample (0/1/N). When ``None``, the + legacy one-image-per-sample path is used. + probe: When ``True``, log a one-line per-slot cost-spread digest (:func:`window_cost_spread`) + on ``dp_rank == 0`` — the §10.5 balance probe. Off by default (the caller throttles it). + + Returns: + ``W`` rebalanced micro-batches for this rank, slot ``s`` at index ``s``. + """ + import torch.distributed as dist + + window_size = len(batches) + if window_size == 0: + raise ValueError("exchange_window requires at least one micro-batch.") + + # Cost and metadata are computed on the CPU batches before the optional CUDA move, avoiding + # per-sample .item() synchronizations on the training stream. Per-slot image-count offset tables + # are derived once and reused for both the CPU-meta and GPU-payload splits. + cu_imgs = [_cu_img_from_counts(image_count_of(b)) if image_count_of is not None else None for b in batches] + cpu_flats_per_slot = [split_microbatch(b, cu_img=cu) for b, cu in zip(batches, cu_imgs)] + local = len(cpu_flats_per_slot[0]) + if any(len(slot) != local for slot in cpu_flats_per_slot): + raise ValueError( + f"All micro-batches in a window must have the same local size; got {[len(s) for s in cpu_flats_per_slot]}." + ) + + # Window-local order is slot-major: this rank's slot-0 shard, then slot-1, … (matches the + # g = s·B_global + j layout build_window_route assumes). + cpu_local_flats = [f for slot in cpu_flats_per_slot for f in slot] + local_meta = [tensor_metadata(f) for f in cpu_local_flats] + local_costs_per_slot = [[cost_of(f) for f in slot] for slot in cpu_flats_per_slot] + + # Keep the exchange GPU-resident after cost calculation: move this rank's (small) shards to the + # device once, then split / serialize / all_to_all / deserialize / merge all run on GPU and the + # returned batches are already on GPU for the forward pass. + if torch.cuda.is_available(): + gpu_flats_per_slot = [split_microbatch(_batch_to_cuda(b), cu_img=cu) for b, cu in zip(batches, cu_imgs)] + local_flats = [f for slot in gpu_flats_per_slot for f in slot] + else: + local_flats = cpu_local_flats + + payload = {"costs_per_slot": local_costs_per_slot, "meta": local_meta} + gathered: List[Any] = [None] * dp_size + dist.all_gather_object(gathered, payload, group=dp_group_gloo) + + b_global = local * dp_size + # Window global index layout: g = slot·b_global + (src_rank·local + i); rank r's local_flats are + # ordered slot-major so its window samples are at all_global_indices[r] (same order). + all_global_indices = [ + [slot * b_global + r * local + i for slot in range(window_size) for i in range(local)] for r in range(dp_size) + ] + all_tensor_meta = [g["meta"] for g in gathered] + # Per-slot global costs indexed by slot-local position j = r·local + i (contiguous sharding). + costs_per_slot: List[List[float]] = [[0.0] * b_global for _ in range(window_size)] + for r, g in enumerate(gathered): + for slot in range(window_size): + for i, c in enumerate(g["costs_per_slot"][slot]): + costs_per_slot[slot][r * local + i] = c + + # Balancing is driven by the patch-only cost (see sample_cost): identical across modules, so both + # compute the same per-slot canonical order with no cross-module communication. + route = build_window_route(costs_per_slot, n_groups, dp_size) + + if probe and dp_rank == 0: + digest = " ".join( + f"slot{s}:{sp['before_max']:.0f}/{sp['before_min']:.0f}->{sp['after_max']:.0f}/{sp['after_min']:.0f}" + f"(remote={int(sp['remote'])})" + for s, sp in enumerate(window_cost_spread(costs_per_slot, route, dp_size)) + ) + logger.info("MIMO reorder balance probe (W=%d) per-slot cost max/min before->after: %s", window_size, digest) + + return exchange_by_route( + local_flats, + route, + all_global_indices, + all_tensor_meta, + dp_rank=dp_rank, + dp_size=dp_size, + window_size=window_size, + local=local, + dp_group_nccl=dp_group_nccl, + nccl_stream=nccl_stream, + ) + + +def reorder_window_local( + batches: List[Dict[str, Any]], + *, + n_groups: int, + cost_of: "Any", + image_count_of: "Any" = None, +) -> List[Dict[str, Any]]: + """Reorder a window into canonical order on a ``dp_size == 1`` module (no exchange needed). + + Same costs + ``n_groups`` as the paired dp>1 module, so the result equals that module's + post-exchange order (rank 0's samples, then rank 1's, ...), keeping the batch-dim pairing. + + Args: + batches: This rank's ``W`` whole micro-batches. + n_groups: Canonical group count (LCM of the module DP sizes). + cost_of: ``callable(flat) -> float`` per-sample cost. + image_count_of: Optional ``callable(batch) -> LongTensor[B]`` per-sample image counts. + + Returns: + ``W`` micro-batches permuted into canonical order. + """ + out: List[Dict[str, Any]] = [] + for batch in batches: + cu_img = _cu_img_from_counts(image_count_of(batch)) if image_count_of is not None else None + flats = split_microbatch(batch, cu_img=cu_img) + costs = [cost_of(f) for f in flats] + assignment = balanced_assignment(costs, n_groups, dp_size=1) + order = sorted(range(len(flats)), key=lambda g: assignment[g][1]) + out.append(_apply_sample_dispatch(batch, order, n_samples=len(flats), cu_img=cu_img)) + return out + + +# --------------------------------------------------------------------------- +# Reorder buffer (data-iterator wrapper) +# --------------------------------------------------------------------------- + + +_SENTINEL = object() + + +class _WorkerException: + """Container used to re-raise overlap-worker failures on the consumer thread.""" + + def __init__(self, exc: BaseException) -> None: + self.exc = exc + + +def _record_stream_for_batch(value: Any, stream: torch.cuda.Stream) -> None: + """Recursively tie CUDA tensors in ``value`` to ``stream`` before cross-thread handoff.""" + if isinstance(value, torch.Tensor) and value.is_cuda: + value.record_stream(stream) + elif isinstance(value, dict): + for child in value.values(): + _record_stream_for_batch(child, stream) + elif isinstance(value, (list, tuple)): + for child in value: + _record_stream_for_batch(child, stream) + + +class ReorderingBuffer: + """Wrap a MegatronMIMO data iterator so each yielded micro-batch is cost-balanced. + + Each item this rank consumes is its disjoint scalable-data-parallel shard rebalanced by + :func:`exchange_window` (per-sample cost all-gather + ragged all-to-all over the module + DP group); at ``dp_size == 1`` it is permuted locally instead (:func:`reorder_window_local`). + + **Window buffering** (``window_size = W``): collect ``W`` micro-batches, cost-balance them in one + exchange over the window (:func:`exchange_window`), and serve the ``W`` rebalanced micro-batches + one at a time from a cursor. ``W = 1`` (default) is the historical per-micro-batch behavior, + byte-for-byte. ``W == GA`` (the gradient-accumulation count) is the desired setting — the single + window collective then lands once per optimizer step. + + **Cross-window prefetch overlap** (``overlap=True``): a background thread runs the *entire next + window's* exchange — the blocking Gloo metadata all-gather *and* the NCCL all-to-all (on a + dedicated CUDA stream) — while the main thread computes the current window's ``W`` micro-batches. + This is what turns the scalable-data-parallel read win into a net throughput win: one big + all-to-all is hidden behind ``W`` micro-batches of compute instead of sitting on the critical + path (for ``W = 1`` this is the one-step-ahead per-micro-batch overlap). The exchange's GPU work + is synchronized inside the worker before hand-off, so the main thread consumes fully materialized + GPU batches with no further wait. A ``maxsize=1`` queue keeps every rank exactly one *window* + ahead, so all ranks' Gloo/NCCL collectives stay in lockstep — re-synced each window by the + optimizer/DDP all-reduce barrier between consumptions. The worker uses the (separate) + ``dp_group_nccl`` exclusively, so its all-to-all never races the main thread's bridge/DDP + collectives on their own PGs (requires ``CUDA_DEVICE_MAX_CONNECTIONS != 1``). Two windows stay + resident (consuming + prefetched) ≈ ``2·W`` micro-batches — the headline memory cost. + + With ``overlap=False`` the window exchange runs synchronously in ``__next__`` (no thread). + """ + + def __init__( + self, + data_iterator: "Any", + *, + dp_rank: int, + dp_size: int, + n_groups: int, + cost_of: "Any", + dp_group_gloo: "Any", + dp_group_nccl: "Any", + overlap: bool = True, + image_count_of: "Any" = None, + window_size: int = 1, + ): + if window_size < 1: + raise ValueError(f"window_size must be >= 1, got {window_size}.") + self._it = data_iterator + self._dp_rank = dp_rank + self._dp_size = dp_size + self._n_groups = n_groups + self._cost_of = cost_of + self._image_count_of = image_count_of + self._gloo = dp_group_gloo + self._nccl = dp_group_nccl + self._window_size = window_size + self._cuda = torch.cuda.is_available() + self._stream = torch.cuda.Stream() if self._cuda else None + self._thread = None + # Sync window state (used whenever the prefetch thread is not running): the active window's + # rebalanced micro-batches and the cursor serving them one at a time. + self._active: List[Dict[str, Any]] = [] + self._cursor = 0 + # Count of window exchanges run, for throttling the §10.5 balance probe (first few + every 50th). + self._exchanges = 0 + # Cross-window prefetch: a thread exchanges window t+1 while the main thread computes + # window t. Started for any window_size (W == 1 is the one-step-ahead per-micro-batch case). + if overlap and self._cuda and dp_size > 1: + # The side-stream exchange only overlaps compute when it can use a separate hardware + # queue. With CUDA_DEVICE_MAX_CONNECTIONS=1 every stream serializes onto one queue, so + # the prefetch a2a runs *behind* the window's compute instead of beside it — correct, but + # silently no faster (often slower, for the extra thread + side PG). Warn loudly rather + # than appear to overlap. Not a hard error: a user may set =1 for TP/SP comm overlap and + # still want the (degraded) reorder to run. + if os.environ.get("CUDA_DEVICE_MAX_CONNECTIONS") == "1": + logger.warning( + "ReorderingBuffer overlap is enabled but CUDA_DEVICE_MAX_CONNECTIONS=1; the " + "exchange will serialize behind compute and give no overlap speedup. Unset it " + "or set it != 1 to get overlap, or pass overlap=False to skip the prefetch thread." + ) + self._device = torch.cuda.current_device() + self._queue: "queue.Queue" = queue.Queue(maxsize=1) + self._stop = threading.Event() + self._thread = threading.Thread(target=self._worker, name="mimo-reorder-prefetch", daemon=True) + self._thread.start() + + def __iter__(self): + return self + + def shutdown(self) -> None: + """Stop the overlap worker and wait briefly for it to exit.""" + if self._thread is None: + return + self._stop.set() + self._thread.join(timeout=1.0) + + def __del__(self) -> None: + """Best-effort cleanup for runs that drop the iterator without closing it.""" + try: + self.shutdown() + except Exception: + logger.debug("Ignoring exception while shutting down ReorderingBuffer.", exc_info=True) + + def _exchange_window(self, batches: List[Dict[str, Any]], stream: "Any") -> List[Dict[str, Any]]: + # Throttle the §10.5 balance probe: log the first few windows, then every 50th. + self._exchanges += 1 + probe = self._exchanges <= 3 or self._exchanges % 50 == 0 + return exchange_window( + batches, + dp_rank=self._dp_rank, + dp_size=self._dp_size, + n_groups=self._n_groups, + cost_of=self._cost_of, + dp_group_gloo=self._gloo, + dp_group_nccl=self._nccl, + nccl_stream=stream, + image_count_of=self._image_count_of, + probe=probe, + ) + + def _collect_window(self) -> List[Dict[str, Any]]: + """Pull up to ``window_size`` micro-batches from the source iterator (fewer at the tail).""" + batches: List[Dict[str, Any]] = [] + for _ in range(self._window_size): + try: + batches.append(next(self._it)) + except StopIteration: + break + return batches + + def _refill_window(self) -> None: + """Collect up to ``window_size`` micro-batches, exchange the window, reset the cursor. + + A short final window (fewer than ``window_size`` micro-batches before ``StopIteration``) is + exchanged at its actual size. ``self._active`` is left empty when the iterator is exhausted. + """ + batches = self._collect_window() + if not batches: + self._active = [] + elif self._dp_size <= 1: + # Nothing to exchange, but the paired dp>1 module is reordered: match its order locally. + self._active = reorder_window_local( + batches, n_groups=self._n_groups, cost_of=self._cost_of, image_count_of=self._image_count_of + ) + else: + self._active = self._exchange_window(batches, stream=self._stream) + self._cursor = 0 + + def _put_worker_item(self, item: Any) -> None: + """Put an item unless shutdown was requested while the bounded queue was full.""" + while not self._stop.is_set(): + try: + self._queue.put(item, timeout=0.1) + return + except queue.Full: + continue + + def _worker(self) -> None: + # Cross-window prefetch: exchange window t+1 while the main thread computes window t. + # Each queue item is the LIST of W rebalanced micro-batches for one window. Everything + # (Gloo metadata all-gather + NCCL all-to-all + deserialize/merge) runs here on self._stream; + # we synchronize that stream before queuing so the consumer gets ready GPU batches. + torch.cuda.set_device(self._device) + try: + while not self._stop.is_set(): + batches = self._collect_window() + if not batches: + break + with torch.cuda.stream(self._stream): + out = self._exchange_window(batches, stream=None) # a2a on the current (self._stream) + self._stream.synchronize() # blocks the WORKER (not the main thread) + self._put_worker_item(out) + except BaseException as exc: + self._put_worker_item(_WorkerException(exc)) + finally: + self._put_worker_item(_SENTINEL) + + def __next__(self) -> Dict[str, Any]: + if self._thread is None: + # Synchronous path (no prefetch thread): serve the active window from the cursor, + # refilling (collect W -> one exchange) when it is exhausted. + if self._cursor >= len(self._active): + self._refill_window() + if not self._active: + raise StopIteration + out = self._active[self._cursor] + self._cursor += 1 + return out + # Cross-window prefetch path: serve from the prefetched window; pull the next one (a list of + # W micro-batches) from the worker when the cursor is exhausted. + if self._cursor >= len(self._active): + item = self._queue.get() + if item is _SENTINEL: + raise StopIteration + if isinstance(item, _WorkerException): + raise item.exc + if torch.cuda.is_available(): + _record_stream_for_batch(item, torch.cuda.current_stream()) + self._active = item + self._cursor = 0 + out = self._active[self._cursor] + self._cursor += 1 + return out + + +def build_module_dp_process_groups(dp_group_nccl_main: "Any", *, overlap: bool) -> "Tuple[int, int, Any, Any]": + """Create the Gloo (metadata) and side NCCL (exchange) PGs for this rank's module DP group. + + ``dist.new_group`` is collective over the world PG, so every rank must create every module's + DP groups in the same order. We all-gather each rank's DP-group global ranks, deduplicate and + sort, then create a Gloo group (and, for overlap, a **separate** NCCL group so the exchange + all-to-all does not serialize against the gradient all-reduce) for each, keeping the ones this + rank belongs to. + + Args: + dp_group_nccl_main: This rank's main DP NCCL group (``grid.get_pg(["dp"])``). + overlap: When ``True``, allocate a separate NCCL PG for the side-stream exchange; + otherwise reuse the main DP NCCL group (sync mode). + + Returns: + ``(dp_rank, dp_size, dp_group_gloo, dp_group_nccl)``. + """ + import torch.distributed as dist + + dp_rank = dist.get_rank(dp_group_nccl_main) + dp_size = dist.get_world_size(dp_group_nccl_main) + my_ranks = list(dist.get_process_group_ranks(dp_group_nccl_main)) + + gathered: List[Any] = [None] * dist.get_world_size() + dist.all_gather_object(gathered, my_ranks) + seen = set() + unique: List[List[int]] = [] + for ranks in gathered: + key = tuple(sorted(ranks)) + if key not in seen: + seen.add(key) + unique.append(list(key)) + unique.sort() + + cur = dist.get_rank() + dp_group_gloo = None + for ranks in unique: + g = dist.new_group(ranks=ranks, backend="gloo") + if cur in ranks: + dp_group_gloo = g + + if overlap: + dp_group_nccl = None + for ranks in unique: + g = dist.new_group(ranks=ranks, backend="nccl") + if cur in ranks: + dp_group_nccl = g + else: + dp_group_nccl = dp_group_nccl_main + + return dp_rank, dp_size, dp_group_gloo, dp_group_nccl diff --git a/src/megatron/bridge/training/megatron_mimo_step.py b/src/megatron/bridge/training/megatron_mimo_step.py index 03173d5426d..91ddf3b799f 100644 --- a/src/megatron/bridge/training/megatron_mimo_step.py +++ b/src/megatron/bridge/training/megatron_mimo_step.py @@ -280,6 +280,9 @@ def forward_step( data_batch, packing_kwargs = pack_language_shard(data_batch, lengths=pack_lengths) if packing_kwargs is not None: data_batch["packing_kwargs"] = packing_kwargs + # The tokenizer padding mask is a data-path length source only (packing, reorder cost, MRoPE); + # the model runs causal attention and MimoModel rejects a dense mask under CP. + data_batch["attention_mask"] = None else: # Non-data stages consume hidden states from pipeline input tensors. data_batch = { diff --git a/src/megatron/bridge/training/train_megatron_mimo.py b/src/megatron/bridge/training/train_megatron_mimo.py index 600741e5bd8..817450cfe42 100644 --- a/src/megatron/bridge/training/train_megatron_mimo.py +++ b/src/megatron/bridge/training/train_megatron_mimo.py @@ -476,4 +476,9 @@ def train_megatron_mimo( timers("interval-time").stop() + # Release the reorder prefetch thread and its side process groups when the train iterator is a + # ReorderingBuffer; a plain iterator has no shutdown() and is left alone. + if hasattr(train_data_iterator, "shutdown"): + train_data_iterator.shutdown() + logger.info(f"Rank {dist.get_rank()}: MegatronMIMO training completed") diff --git a/tests/functional_tests/test_groups/megatron_mimo/test_reorder_exchange.py b/tests/functional_tests/test_groups/megatron_mimo/test_reorder_exchange.py new file mode 100644 index 00000000000..c4fe6225b5c --- /dev/null +++ b/tests/functional_tests/test_groups/megatron_mimo/test_reorder_exchange.py @@ -0,0 +1,175 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Functional test for the MegatronMIMO intra-microbatch reorder exchange on a real 2-rank +process group. + +The CPU unit tests (``tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py``) emulate the +``all_to_all_single``; this exercises the on-device path that those cannot — the real Gloo cost +all-gather + NCCL ragged ``all_to_all_single`` inside :func:`exchange_window`, plus the side-PG +setup in :func:`build_module_dp_process_groups`. The 2-rank world is treated as a single module's +``dp=2`` group (het-DP pairing needs two modules / 4 ranks and stays unit-emulated). + +Run by ``tests/functional_tests/launch_scripts/h100/active/L0_Launch_training_megatron_mimo.sh`` +(``torch.distributed.run --nproc_per_node=2 -m pytest .../test_groups/megatron_mimo``). +""" + +from __future__ import annotations + +import functools +from typing import Any, List, Tuple + +import pytest +import torch +import torch.distributed as dist + +from megatron.bridge.data.megatron_mimo.reorder_buffer import ( + balanced_assignment, + build_module_dp_process_groups, + exchange_window, + sample_cost, +) +from tests.functional_tests.utils import initialize_distributed + + +_VISION_START = 990 # synthetic vision-start token; one per image, counted by image_count_of +_SEQ = 8 +_DIM = 4 + + +def _make_shard(samples: List[Tuple[int, List[int]]]) -> dict: + """Build one rank's nested micro-batch from ``[(global_idx, [patches_per_image, ...]), ...]``. + + ``input_ids[i]`` is tagged with ``1 + g*100`` (so the global index is recoverable as + ``(first_token - 1)//100``) and carries one ``_VISION_START`` token per image. Each image + contributes ``p`` ``hidden_states`` rows all filled with ``float(g)`` so vision can be traced to + its owning sample; a text-only sample (no images) contributes empty ``[0, d]`` / ``[0, 3]``. + """ + b = len(samples) + rows = [] + grids: List[List[int]] = [] + hidden_blocks: List[torch.Tensor] = [] + for g, patches in samples: + row = [7] * _SEQ + row[0] = 1 + g * 100 + for k in range(len(patches)): + row[1 + k] = _VISION_START + rows.append(row) + for p in patches: + grids.append([1, 1, p]) + hidden_blocks.append(torch.full((p, _DIM), float(g), dtype=torch.float32)) + input_ids = torch.tensor(rows, dtype=torch.int64) + grid_thw = ( + torch.tensor(grids, dtype=torch.int64).reshape(-1, 3) if grids else torch.empty((0, 3), dtype=torch.int64) + ) + hidden = torch.cat(hidden_blocks, dim=0) if hidden_blocks else torch.empty((0, _DIM), dtype=torch.float32) + return { + "input_ids": input_ids, + "position_ids": torch.stack([input_ids, input_ids, input_ids], dim=0), # [3, B, S] MRoPE + "attention_mask": None, + "labels": input_ids + 1, + "loss_mask": torch.ones(b, _SEQ, dtype=torch.bool), + "modality_inputs": {"images": {"enc": {"hidden_states": hidden, "grid_thw": grid_thw}}}, + } + + +def _expected_owned(global_patches: List[List[int]], dp_rank: int, dp_size: int, n_groups: int) -> List[int]: + """Global indices this rank should own after the exchange, in canonical-position order.""" + costs = [float(sum(p)) for p in global_patches] + assignment = balanced_assignment(costs, n_groups, dp_size) + owned = sorted((pos, g) for g, (owner, pos) in enumerate(assignment) if owner == dp_rank) + return [g for _pos, g in owned] + + +def _run_exchange_case(global_samples: List[List[int]], image_count_of: "Any") -> None: + """Shard ``global_samples`` (patch lists per global sample) contiguously across 2 ranks, run the + real exchange, and assert this rank recovers exactly its cost-balanced shard with vision intact. + """ + dp_size = dist.get_world_size() + n_groups = dp_size + b = len(global_samples) + local = b // dp_size + rank = dist.get_rank() + + my_globals = list(range(rank * local, (rank + 1) * local)) + batch = _make_shard([(g, global_samples[g]) for g in my_globals]) + + main_pg = dist.new_group(ranks=list(range(dp_size)), backend="nccl") + dp_rank, ds, gloo, nccl = build_module_dp_process_groups(main_pg, overlap=False) + assert (ds, dp_rank) == (dp_size, rank) + + out = exchange_window( + [batch], + dp_rank=dp_rank, + dp_size=ds, + n_groups=n_groups, + cost_of=functools.partial(sample_cost, encoder_cost_weight=1.0), + dp_group_gloo=gloo, + dp_group_nccl=nccl, + image_count_of=image_count_of, + )[0] + + expected_globals = _expected_owned(global_samples, dp_rank, dp_size, n_groups) + + got_globals = [int((int(row[0].item()) - 1) // 100) for row in out["input_ids"].cpu()] + assert got_globals == expected_globals, (rank, got_globals, expected_globals) + + # Vision must travel with its sample through the real all-to-all: reconstruct the expected + # concatenation (per-sample patch blocks tagged float(g)) and compare byte-for-byte. + exp_hidden = [torch.full((sum(global_samples[g]), _DIM), float(g), dtype=torch.float32) for g in expected_globals] + exp_grid = [ + torch.tensor([[1, 1, p] for p in global_samples[g]], dtype=torch.int64).reshape(-1, 3) + for g in expected_globals + ] + exp_hidden_t = torch.cat(exp_hidden, dim=0) if exp_hidden else torch.empty((0, _DIM), dtype=torch.float32) + exp_grid_t = torch.cat(exp_grid, dim=0) if exp_grid else torch.empty((0, 3), dtype=torch.int64) + + enc = out["modality_inputs"]["images"]["enc"] + assert torch.equal(enc["hidden_states"].cpu().float(), exp_hidden_t), rank + assert torch.equal(enc["grid_thw"].cpu(), exp_grid_t), rank + + +@pytest.mark.run_only_on("GPU") +def test_reorder_exchange_dp2_single_image(): + """One image per sample; per-sample patch counts force a real cross-rank swap.""" + initialize_distributed() + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for this functional test") + if dist.get_world_size() != 2: + pytest.skip("This functional test requires exactly 2 ranks") + + # costs [4,3,1,2]: balanced owner is g0,g2 -> rank0 and g1,g3 -> rank1, but the contiguous + # initial shards are [g0,g1] / [g2,g3], so rank0 must SEND g1 and RECEIVE g2 (a real exchange). + global_samples = [[4], [3], [1], [2]] + image_count_of = lambda b: (b["input_ids"] == _VISION_START).sum(dim=1).to(torch.long) # noqa: E731 + _run_exchange_case(global_samples, image_count_of) + dist.barrier() + + +@pytest.mark.run_only_on("GPU") +def test_reorder_exchange_dp2_variable_images(): + """Mixed image counts per sample (text-only / single / multi) survive the real exchange.""" + initialize_distributed() + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for this functional test") + if dist.get_world_size() != 2: + pytest.skip("This functional test requires exactly 2 ranks") + + # counts [2,1,0,1] (g2 text-only, g0 multi-image), patches=2/image -> costs [4,2,0,2]: owner is + # g0,g2 -> rank0 and g1,g3 -> rank1; initial shards [g0,g1] / [g2,g3], so a multi-image sample + # stays, a single-image (g1) and a TEXT-ONLY (g2) sample swap across ranks. + global_samples = [[2, 2], [2], [], [2]] + image_count_of = lambda b: (b["input_ids"] == _VISION_START).sum(dim=1).to(torch.long) # noqa: E731 + _run_exchange_case(global_samples, image_count_of) + dist.barrier() diff --git a/tests/unit_tests/data/builders/test_direct_hf_sft.py b/tests/unit_tests/data/builders/test_direct_hf_sft.py index 8349c9952d4..45f0d5d10e9 100644 --- a/tests/unit_tests/data/builders/test_direct_hf_sft.py +++ b/tests/unit_tests/data/builders/test_direct_hf_sft.py @@ -626,3 +626,53 @@ def build(self, context): assert seen["config"] is config assert seen["context"].train_samples == 12 assert seen["context"].tokenizer is tokenizer + + +def _mimo_reorder_config(**overrides): + kwargs = dict( + seq_length=128, + source=HFDatasetSourceConfig(path_or_dataset="json"), + dataloader_type="single", + megatron_mimo_scalable_dp=True, + megatron_mimo_intra_microbatch_reorder=True, + ) + kwargs.update(overrides) + config = DirectHFSFTDatasetConfig(**kwargs) + config.drop_last = True + return config + + +def test_mimo_reorder_requires_scalable_dp(): + config = _mimo_reorder_config(megatron_mimo_scalable_dp=False) + with pytest.raises(ValueError, match="requires megatron_mimo_scalable_dp=True"): + config.validate() + + +def test_mimo_reorder_rejects_zero_cost(): + config = _mimo_reorder_config( + megatron_mimo_reorder_encoder_cost_weight=0.0, megatron_mimo_reorder_language_cost_weight=0.0 + ) + with pytest.raises(ValueError, match="encoder_cost_weight > 0"): + config.validate() + + +def test_mimo_reorder_rejects_negative_cost_weight(): + config = _mimo_reorder_config(megatron_mimo_reorder_language_cost_weight=-1.0) + with pytest.raises(ValueError, match="cost_weight must be >= 0"): + config.validate() + + +def test_mimo_reorder_rejects_window_below_one(): + config = _mimo_reorder_config(megatron_mimo_reorder_window_size=0) + with pytest.raises(ValueError, match="window_size must be >= 1"): + config.validate() + + +def test_mimo_reorder_defaults_validate(): + _mimo_reorder_config().validate() + assert ( + DirectHFSFTDatasetConfig( + seq_length=128, source=HFDatasetSourceConfig(path_or_dataset="json") + ).megatron_mimo_intra_microbatch_reorder + is False + ) diff --git a/tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py b/tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py new file mode 100644 index 00000000000..c16436efe3b --- /dev/null +++ b/tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py @@ -0,0 +1,1186 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the per-sample reorder-buffer building blocks (Phase 0, no process group). + +The all-to-all is emulated on CPU: each rank's ``send_buf`` is split by ``send_splits`` into +per-destination chunks, and each destination's ``recv_buf`` is the concatenation of the chunks +sent to it in source-rank order — exactly the contract of ``dist.all_to_all_single``. +""" + +import pytest +import torch + +from megatron.bridge.data.megatron_mimo.canonical_sampler import ( + build_canonical_group_batch_sampler, + canonical_grid_size, + covered_canonical_groups, +) +from megatron.bridge.data.megatron_mimo.reorder_buffer import ( + ReorderingBuffer, + _cu_img_from_counts, + _reorder_vision_by_images, + assert_intra_no_cross_slot, + balanced_assignment, + balanced_index_order, + build_window_route, + deserialize_sample, + intra_route, + merge_samples, + prepare_sample_exchange, + reassemble_window, + reorder_window_local, + sample_byte_size, + sample_cost, + sample_keys, + serialize_sample, + split_microbatch, + tensor_metadata, + window_cost_spread, +) + + +def _make_sample(global_idx: int, n_tokens: int, n_patches: int) -> dict: + """A ragged synthetic MIMO-like sample, content keyed off global_idx for identity checks.""" + return { + "input_ids": torch.arange(global_idx * 1000, global_idx * 1000 + n_tokens, dtype=torch.int64), + "loss_mask": torch.ones(n_tokens, dtype=torch.bool), + # vision patches: [n_patches, 4], 0 patches -> None (text-only) + "hidden_states": (torch.full((n_patches, 4), float(global_idx), dtype=torch.float32) if n_patches else None), + } + + +# --------------------------------------------------------------------------- +# Serialization round-trip +# --------------------------------------------------------------------------- + + +def test_serialize_deserialize_roundtrip(): + flat = _make_sample(global_idx=7, n_tokens=5, n_patches=3) + meta = tensor_metadata(flat) + keys = sample_keys(meta) + assert keys == ["hidden_states", "input_ids", "loss_mask"] # sorted, None dropped + + buf = serialize_sample(flat, keys) + # byte size from metadata alone must match the actual buffer + assert buf.numel() == sample_byte_size(meta, keys) + assert buf.numel() % 8 == 0 # 8-byte aligned + + out, offset = deserialize_sample(buf, 0, meta, keys) + assert offset == buf.numel() + for k in keys: + assert torch.equal(out[k], flat[k]), k + # None keys are restored + assert set(out.keys()) == set(flat.keys()) + + +def test_serialize_text_only_sample(): + flat = _make_sample(global_idx=2, n_tokens=4, n_patches=0) + meta = tensor_metadata(flat) + assert meta["hidden_states"] is None + keys = sample_keys(meta) + assert "hidden_states" not in keys + buf = serialize_sample(flat, keys) + out, _ = deserialize_sample(buf, 0, meta, keys) + assert out["hidden_states"] is None + assert torch.equal(out["input_ids"], flat["input_ids"]) + + +def test_non_tensor_metadata_roundtrip(): + flat = {"input_ids": torch.tensor([1, 2]), "sample_id": "abc", "optional": None} + meta = tensor_metadata(flat) + keys = sample_keys(meta) + assert keys == ["input_ids"] + + buf = serialize_sample(flat, keys) + out, offset = deserialize_sample(buf, 0, meta, keys) + + assert offset == buf.numel() + assert torch.equal(out["input_ids"], flat["input_ids"]) + assert out["sample_id"] == "abc" + assert out["optional"] is None + + +def test_unknown_metadata_dtype_raises(): + meta = {"input_ids": {"shape": [2], "dtype": "torch.not_a_dtype"}} + keys = sample_keys(meta) + with pytest.raises(ValueError, match="Unsupported tensor dtype"): + sample_byte_size(meta, keys) + with pytest.raises(ValueError, match="Unsupported tensor dtype"): + deserialize_sample(torch.empty(0, dtype=torch.uint8), 0, meta, keys) + + +def test_tensor_metadata_rejects_unsupported_dtype_early(): + # An out-of-map dtype (float64) must fail at metadata creation on the producing rank, + # before the all-gather — not only later as a receiver-side byte-size/deserialize error. + # Reuses the same _dtype_info check the decode path uses. + flat = {"input_ids": torch.tensor([1, 2]), "weird": torch.zeros(2, dtype=torch.float64)} + with pytest.raises(ValueError, match="Unsupported tensor dtype"): + tensor_metadata(flat) + + +# --------------------------------------------------------------------------- +# Assignment +# --------------------------------------------------------------------------- + + +def test_balanced_assignment_contiguous_blocks_and_balance(): + # 8 samples, dp=2, canonical n_groups=2. + costs = [10.0, 1.0, 9.0, 2.0, 8.0, 3.0, 7.0, 4.0] + n_groups, dp_size = 2, 2 + assignment = balanced_assignment(costs, n_groups, dp_size) + + n = len(costs) + local = n // dp_size + # each rank owns a contiguous canonical block [r*local, (r+1)*local) + per_rank = {0: [], 1: []} + pos_seen = set() + for global_idx, (owner, pos) in enumerate(assignment): + per_rank[owner].append(pos) + pos_seen.add(pos) + assert pos_seen == set(range(n)) # every canonical position used once + assert sorted(per_rank[0]) == list(range(0, local)) + assert sorted(per_rank[1]) == list(range(local, n)) + + # cost balance: each rank's total cost within ~max single item of the other + rank_cost = {0: 0.0, 1: 0.0} + for global_idx, (owner, _pos) in enumerate(assignment): + rank_cost[owner] += costs[global_idx] + assert abs(rank_cost[0] - rank_cost[1]) <= max(costs) + + +def test_balanced_assignment_matches_balanced_index_order(): + costs = [5.0, 1.0, 4.0, 2.0] + assignment = balanced_assignment(costs, n_groups=2, dp_size=2) + perm = balanced_index_order(costs, 2) + # canonical_pos in the assignment must invert the perm + pos2gidx = {pos: g for g, (_owner, pos) in enumerate(assignment)} + assert [pos2gidx[p] for p in range(len(costs))] == perm + + +def test_balanced_assignment_invalid_divisibility(): + for kwargs in ( + dict(costs=[1.0, 2.0, 3.0], n_groups=2, dp_size=2), # n_groups does not divide n + dict(costs=[1.0, 2.0, 3.0, 4.0], n_groups=4, dp_size=3), # dp does not divide n_groups + ): + try: + balanced_assignment(**kwargs) + except ValueError: + continue + raise AssertionError(f"expected ValueError for {kwargs}") + + +# --------------------------------------------------------------------------- +# Het-DP pairing: dp2 (vision) and dp4 (language) over the same costs/canonical groups +# --------------------------------------------------------------------------- + + +class TestBalancedIndexOrder: + def test_valid_permutation_and_equal_groups(self): + costs = [5.0, 1.0, 4.0, 2.0, 3.0, 6.0] + order = balanced_index_order(costs, n_groups=2) + assert sorted(order) == list(range(6)) # permutation + g0, g1 = order[:3], order[3:] + assert len(g0) == 3 and len(g1) == 3 + + def test_balances_contiguous_shards(self): + # Grouped heavies: a naive contiguous split would put all three 10s on one shard. + costs = [10.0, 10.0, 10.0, 1.0, 1.0, 1.0] + order = balanced_index_order(costs, n_groups=3) # 2 per group + groups = [order[0:2], order[2:4], order[4:6]] + totals = [sum(costs[i] for i in g) for g in groups] + assert max(totals) - min(totals) == 0.0 # each group one heavy + one light + + def test_determinism(self): + costs = [3.0, 1.0, 9.0, 0.5, 4.0, 2.0] + assert balanced_index_order(costs, 2) == balanced_index_order(costs, 2) + + def test_errors(self): + with pytest.raises(ValueError): + balanced_index_order([1.0, 2.0, 3.0], 2) # 3 not divisible by 2 + with pytest.raises(ValueError): + balanced_index_order([1.0], 0) + + +@pytest.mark.parametrize("module_dps", [(1, 2), (2, 4), (2, 3)]) +def test_het_dp_pairing_on_canonical_grid(module_dps): + """n_groups = canonical_grid_size(module dps) (the LCM), the geometry the scalable loader reads with. + + Every module must derive the same canonical order, and each module's owner rank must be the + rank that covers that canonical group under covered_canonical_groups. + """ + grid = canonical_grid_size(list(module_dps)) + n = grid * 3 # 3 samples per canonical group + costs = [float((7 * g) % 11 + 1) for g in range(n)] + cap = n // grid + assignments = {dp: balanced_assignment(costs, grid, dp_size=dp) for dp in module_dps} + reference = assignments[module_dps[0]] + for dp, assignment in assignments.items(): + for g in range(n): + owner, canonical_pos = assignment[g] + assert canonical_pos == reference[g][1], (dp, g) + assert canonical_pos // cap in covered_canonical_groups(owner, dp, grid), (dp, g) + + +def test_het_dp_pairing_consistency(): + # 12 samples, canonical n_groups = max dp = 4. Vision dp2, language dp4. + costs = [float(c) for c in [12, 1, 11, 2, 10, 3, 9, 4, 8, 5, 7, 6]] + n_groups = 4 + vis = balanced_assignment(costs, n_groups, dp_size=2) + lang = balanced_assignment(costs, n_groups, dp_size=4) + + # canonical_pos must be identical across modules (same balanced order) + for g in range(len(costs)): + assert vis[g][1] == lang[g][1], g + + # vision rank r must cover exactly language ranks {2r, 2r+1} + for g in range(len(costs)): + v_owner = vis[g][0] + l_owner = lang[g][0] + assert l_owner // 2 == v_owner, (g, v_owner, l_owner) + + +# --------------------------------------------------------------------------- +# Full per-sample all-to-all (emulated collective) +# --------------------------------------------------------------------------- + + +def _emulate_all_to_all(plans, dp_size): + """Emulate dist.all_to_all_single across ``dp_size`` ranks from their plans.""" + # split each source's send_buf into per-destination chunks + chunks = {} + for s in range(dp_size): + buf = plans[s]["send_buf"] + off = 0 + for d in range(dp_size): + nbytes = plans[s]["send_splits"][d] + chunks[(s, d)] = buf[off : off + nbytes] + off += nbytes + assert off == buf.numel(), f"rank {s} send_splits do not cover send_buf" + + # each destination receives chunks in source-rank order; verify split symmetry + recv_bufs = {} + for d in range(dp_size): + parts = [] + for s in range(dp_size): + assert plans[d]["recv_splits"][s] == chunks[(s, d)].numel(), ( + f"recv_splits[{d}][{s}] != bytes sent {s}->{d}" + ) + parts.append(chunks[(s, d)]) + recv_bufs[d] = torch.cat(parts) if parts else torch.empty(0, dtype=torch.uint8) + return recv_bufs + + +def _run_exchange(costs, dp_size, n_groups, sample_specs): + """End-to-end: read disjoint shards -> assign -> plan -> emulate a2a -> reconstruct. + + sample_specs[global_idx] = (n_tokens, n_patches). Initial (disjoint, contiguous) layout: + rank r holds global indices [r*local, (r+1)*local). + """ + n = len(costs) + local = n // dp_size + originals = {g: _make_sample(g, *sample_specs[g]) for g in range(n)} + + # initial disjoint layout + rank_global_indices = [list(range(r * local, (r + 1) * local)) for r in range(dp_size)] + rank_flats = [[originals[g] for g in rank_global_indices[r]] for r in range(dp_size)] + all_tensor_meta = [[tensor_metadata(f) for f in rank_flats[r]] for r in range(dp_size)] + + assignment = balanced_assignment(costs, n_groups, dp_size) + route = intra_route(assignment, src_slot=0) + pos2gidx = {pos: g for g, (_o, pos) in enumerate(assignment)} + + plans = [ + prepare_sample_exchange( + local_flats=rank_flats[r], + local_global_indices=rank_global_indices[r], + route=route, + all_global_indices=rank_global_indices, + all_tensor_meta=all_tensor_meta, + dp_rank=r, + dp_size=dp_size, + window_size=1, + ) + for r in range(dp_size) + ] + + recv_bufs = _emulate_all_to_all(plans, dp_size) + + # reconstruct each rank's balanced shard + for d in range(dp_size): + recovered = [] + for _dst_slot, canonical_pos, local_idx in plans[d]["local_samples"]: + recovered.append((canonical_pos, rank_flats[d][local_idx])) + buf = recv_bufs[d] + offset = 0 + for src in range(dp_size): + for _dst_slot, canonical_pos, src_local_idx in plans[d]["recv_schedule"].get(src, []): + meta = all_tensor_meta[src][src_local_idx] + flat, offset = deserialize_sample(buf, offset, meta, sample_keys(meta)) + recovered.append((canonical_pos, flat)) + assert offset == buf.numel(), f"rank {d} did not consume entire recv_buf" + + recovered.sort(key=lambda x: x[0]) + positions = [p for p, _ in recovered] + # rank d owns a contiguous canonical block + assert positions == list(range(d * local, (d + 1) * local)), (d, positions) + + # tensors survived the round-trip intact, matched by global index + for canonical_pos, flat in recovered: + g = pos2gidx[canonical_pos] + for k, v in originals[g].items(): + if v is None: + assert flat[k] is None, (g, k) + else: + assert torch.equal(flat[k], v), (g, k) + + +def test_exchange_dp2_mixed_modality(): + # 8 samples, some text-only (0 patches), varied token/patch loads + sample_specs = { + 0: (10, 5), + 1: (3, 0), + 2: (8, 2), + 3: (4, 0), + 4: (12, 7), + 5: (2, 1), + 6: (6, 0), + 7: (5, 3), + } + costs = [t + 4.0 * p for t, p in (sample_specs[g] for g in range(8))] + _run_exchange(costs, dp_size=2, n_groups=2, sample_specs=sample_specs) + + +def test_exchange_dp4_canonical(): + sample_specs = {g: (3 + g, (g % 3)) for g in range(12)} + costs = [t + 4.0 * p for t, p in (sample_specs[g] for g in range(12))] + _run_exchange(costs, dp_size=4, n_groups=4, sample_specs=sample_specs) + + +# --------------------------------------------------------------------------- +# MIMO data plane: split_microbatch / merge_samples +# --------------------------------------------------------------------------- + + +def _make_mimo_batch(patches_per_sample): + """Synthetic nested MIMO micro-batch: one image per sample, MRoPE position_ids.""" + b = len(patches_per_sample) + s = 6 + # Start at 1 (not 0) so no token collides with the default pad_token_id=0: every token is real, + # which keeps the real-token count == s for the joint-cost test. + input_ids = torch.stack([torch.arange(1 + i * 100, 1 + i * 100 + s, dtype=torch.int64) for i in range(b)]) + labels = input_ids + 1 + loss_mask = torch.ones(b, s, dtype=torch.bool) + position_ids = torch.stack([input_ids, input_ids, input_ids], dim=0) # [3, B, S] MRoPE + # one image per sample; grid [1, h, w] with prod(grid) = patches + grids = torch.tensor([[1, 1, p] for p in patches_per_sample], dtype=torch.int64) # [B, 3] + hidden = torch.cat( + [torch.full((p, 4), float(i), dtype=torch.float32) for i, p in enumerate(patches_per_sample)], dim=0 + ) + return { + "input_ids": input_ids, + "position_ids": position_ids, + "attention_mask": None, + "labels": labels, + "loss_mask": loss_mask, + "modality_inputs": {"images": {"enc": {"hidden_states": hidden, "grid_thw": grids}}}, + } + + +def _assert_batch_equal(a, b): + assert set(a.keys()) == set(b.keys()) + for k, va in a.items(): + vb = b[k] + if isinstance(va, dict): + _assert_batch_equal(va, vb) + elif va is None: + assert vb is None, k + else: + assert torch.equal(va, vb), k + + +def test_split_merge_roundtrip_identity(): + batch = _make_mimo_batch([2, 5, 1, 3]) + samples = split_microbatch(batch) + assert len(samples) == 4 + # per-sample flat dicts carry dotted vision keys + assert "modality_inputs.images.enc.hidden_states" in samples[0] + merged = merge_samples(samples) + _assert_batch_equal(merged, batch) + + +# --------------------------------------------------------------------------- +# merge_samples robustness to per-sample-varying vision (PR-A) +# --------------------------------------------------------------------------- + +_VKEY = "modality_inputs.images.enc" + + +def _flat_vision_sample(idx: int, n_tokens: int, n_patches: int, *, none_vision: bool = False) -> dict: + """Per-sample flat dict as ``split_microbatch`` produces: ``[1, T]`` input_ids and one image's + ``[p, 4]`` hidden_states + ``[1, 3]`` grid_thw. ``n_patches==0`` is text-only: empty ``[0, 4]`` / + ``[0, 3]`` (or ``None`` when ``none_vision`` to exercise the legacy/None merge path).""" + input_ids = torch.arange(idx * 100, idx * 100 + n_tokens, dtype=torch.int64).reshape(1, -1) + if n_patches == 0: + if none_vision: + hidden, grid = None, None + else: + hidden = torch.empty((0, 4), dtype=torch.float32) + grid = torch.empty((0, 3), dtype=torch.int64) + else: + hidden = torch.full((n_patches, 4), float(idx), dtype=torch.float32) + grid = torch.tensor([[1, 1, n_patches]], dtype=torch.int64) + return {"input_ids": input_ids, f"{_VKEY}.hidden_states": hidden, f"{_VKEY}.grid_thw": grid} + + +def test_merge_text_only_first(): + # Sample 0 is text-only with vision keys as None (the case the old ``ref = samples[0]`` + # classification silently dropped); samples 1-2 carry one image each. + flats = [ + _flat_vision_sample(0, n_tokens=4, n_patches=0, none_vision=True), + _flat_vision_sample(1, n_tokens=4, n_patches=3), + _flat_vision_sample(2, n_tokens=4, n_patches=2), + ] + enc = merge_samples(flats)["modality_inputs"]["images"]["enc"] + # 0 + 3 + 2 = 5 patches, 0 + 1 + 1 = 2 image rows — no silent loss of sample 1/2 vision. + assert enc["hidden_states"].shape == (5, 4) + assert enc["grid_thw"].shape == (2, 3) + assert torch.equal(enc["hidden_states"][:3], torch.full((3, 4), 1.0)) + assert torch.equal(enc["hidden_states"][3:], torch.full((2, 4), 2.0)) + + +def test_merge_all_text_only(): + flats = [_flat_vision_sample(i, n_tokens=4, n_patches=0) for i in range(3)] + enc = merge_samples(flats)["modality_inputs"]["images"]["enc"] + assert enc["hidden_states"].shape == (0, 4) + assert enc["grid_thw"].shape == (0, 3) + + +def test_serialize_empty_vision_roundtrip(): + flat = _flat_vision_sample(0, n_tokens=4, n_patches=0) # [0,4] / [0,3] empties + meta = tensor_metadata(flat) + keys = sample_keys(meta) + buf = serialize_sample(flat, keys) + assert buf.numel() == sample_byte_size(meta, keys) # empties contribute 0 bytes + out, offset = deserialize_sample(buf, 0, meta, keys) + assert offset == buf.numel() + assert out[f"{_VKEY}.hidden_states"].shape == (0, 4) + assert out[f"{_VKEY}.grid_thw"].shape == (0, 3) + assert torch.equal(out["input_ids"], flat["input_ids"]) + + +# --------------------------------------------------------------------------- +# Variable images per sample: split / merge / exchange (PR-B) +# --------------------------------------------------------------------------- + + +def _make_mimo_batch_var(images_per_sample: list[int]) -> dict: + """Nested MIMO micro-batch with a variable number of images per sample (0/1/N). + + ``images_per_sample[s]`` images for sample ``s``; each image gets a distinct patch count and its + ``hidden_states`` rows are tagged with its global image index so reorder/merge can be content-checked. + """ + b = len(images_per_sample) + seq = 6 + input_ids = torch.stack([torch.arange(i * 100, i * 100 + seq, dtype=torch.int64) for i in range(b)]) + position_ids = torch.stack([input_ids, input_ids, input_ids], dim=0) # [3, B, S] MRoPE + grids: list[list[int]] = [] + hidden_blocks: list[torch.Tensor] = [] + img_gid = 0 + for n_img in images_per_sample: + for _ in range(n_img): + p = img_gid + 2 # distinct, identifiable patch count per image + grids.append([1, 1, p]) + hidden_blocks.append(torch.full((p, 4), float(img_gid), dtype=torch.float32)) + img_gid += 1 + grid_thw = ( + torch.tensor(grids, dtype=torch.int64).reshape(-1, 3) if grids else torch.empty((0, 3), dtype=torch.int64) + ) + hidden = torch.cat(hidden_blocks, dim=0) if hidden_blocks else torch.empty((0, 4), dtype=torch.float32) + return { + "input_ids": input_ids, + "position_ids": position_ids, + "attention_mask": None, + "labels": input_ids + 1, + "loss_mask": torch.ones(b, seq, dtype=torch.bool), + "modality_inputs": {"images": {"enc": {"hidden_states": hidden, "grid_thw": grid_thw}}}, + } + + +def test_cu_img_from_counts(): + assert _cu_img_from_counts([2, 0, 1, 3]) == [0, 2, 2, 3, 6] + assert _cu_img_from_counts(torch.tensor([2, 0, 1, 3])) == [0, 2, 2, 3, 6] + assert _cu_img_from_counts([0, 0]) == [0, 0, 0] + + +def test_image_count_sourcing(): + vision_start = 248053 + input_ids = torch.tensor( + [ + [vision_start, 5, vision_start, 6, 7, 8], # 2 images + [1, 2, 3, 4, 5, 6], # text-only + [vision_start, 9, 9, 9, 9, 9], # 1 image + ], + dtype=torch.int64, + ) + counts = (input_ids == vision_start).sum(dim=1) + assert counts.tolist() == [2, 0, 1] + batch = _make_mimo_batch_var([2, 0, 1]) + assert int(counts.sum()) == batch["modality_inputs"]["images"]["enc"]["grid_thw"].shape[0] + + +def test_split_merge_roundtrip_multi_image(): + images_per_sample = [2, 0, 1, 3] + batch = _make_mimo_batch_var(images_per_sample) + samples = split_microbatch(batch, cu_img=_cu_img_from_counts(images_per_sample)) + assert len(samples) == 4 + merged = merge_samples(samples) + _assert_batch_equal(merged, batch) + + +def test_split_all_text_only(): + batch = _make_mimo_batch_var([0, 0]) + samples = split_microbatch(batch, cu_img=_cu_img_from_counts([0, 0])) + for smp in samples: + assert smp[f"{_VKEY}.hidden_states"].shape == (0, 4) + assert smp[f"{_VKEY}.grid_thw"].shape == (0, 3) + _assert_batch_equal(merge_samples(samples), batch) + + +def test_soft_validation_mismatch(): + batch = _make_mimo_batch_var([1, 1]) # 2 images + with pytest.raises(ValueError, match="mis-sourced"): + split_microbatch(batch, cu_img=[0, 1, 5]) # sums to 5 != 2 images + + +def test_split_legacy_no_cu_img(): + # one-image-per-sample works with cu_img=None (legacy path) + batch = _make_mimo_batch([2, 5, 1, 3]) + _assert_batch_equal(merge_samples(split_microbatch(batch)), batch) + # multi-image without cu_img raises a clear error + multi = _make_mimo_batch_var([2, 1]) # 3 images, 2 samples + with pytest.raises(ValueError, match="one image per sample"): + split_microbatch(multi) + + +def _run_exchange_nested(images_per_sample, dp_size, n_groups): + """Full split→assign→exchange→reconstruct on nested variable-image batches (emulated a2a).""" + n = len(images_per_sample) + local = n // dp_size + full = _make_mimo_batch_var(images_per_sample) + all_flats = split_microbatch(full, cu_img=_cu_img_from_counts(images_per_sample)) + costs = [sample_cost(f, encoder_cost_weight=1.0) for f in all_flats] + + rank_global_indices = [list(range(r * local, (r + 1) * local)) for r in range(dp_size)] + rank_flats = [[all_flats[g] for g in rank_global_indices[r]] for r in range(dp_size)] + all_tensor_meta = [[tensor_metadata(f) for f in rank_flats[r]] for r in range(dp_size)] + assignment = balanced_assignment(costs, n_groups, dp_size) + route = intra_route(assignment, src_slot=0) + pos2gidx = {pos: g for g, (_o, pos) in enumerate(assignment)} + + plans = [ + prepare_sample_exchange( + local_flats=rank_flats[r], + local_global_indices=rank_global_indices[r], + route=route, + all_global_indices=rank_global_indices, + all_tensor_meta=all_tensor_meta, + dp_rank=r, + dp_size=dp_size, + window_size=1, + ) + for r in range(dp_size) + ] + recv_bufs = _emulate_all_to_all(plans, dp_size) + + for d in range(dp_size): + recovered = [(pos, rank_flats[d][li]) for _slot, pos, li in plans[d]["local_samples"]] + buf = recv_bufs[d] + offset = 0 + for src in range(dp_size): + for _dst_slot, canonical_pos, src_local_idx in plans[d]["recv_schedule"].get(src, []): + meta = all_tensor_meta[src][src_local_idx] + flat, offset = deserialize_sample(buf, offset, meta, sample_keys(meta)) + recovered.append((canonical_pos, flat)) + assert offset == buf.numel(), (d, offset, buf.numel()) + recovered.sort(key=lambda x: x[0]) + assert [p for p, _ in recovered] == list(range(d * local, (d + 1) * local)), d + for canonical_pos, flat in recovered: + g = pos2gidx[canonical_pos] + for k, v in all_flats[g].items(): + if v is None: + assert flat[k] is None, (g, k) + else: + assert torch.equal(flat[k], v), (g, k) + + +def test_exchange_var_images_dp2(): + # 8 samples: text-only (0), single (1), and multi-image (2/3) mixed so a rank owns a variety. + _run_exchange_nested([2, 0, 1, 0, 3, 1, 0, 2], dp_size=2, n_groups=2) + + +# --------------------------------------------------------------------------- +# 3D routing seam + per-slot reassembly (reassemble_window): Phase-1 gates §10.1.1 / §10.1.2 +# --------------------------------------------------------------------------- + + +def _emulate_window_exchange(rank_flats, rank_global_indices, all_tensor_meta, route, *, dp_size, window_size, local): + """Emulate the on-device exchange on CPU and reassemble each rank's window via the production + :func:`reassemble_window`. Returns ``per_rank_batches[d][slot]``.""" + plans = [ + prepare_sample_exchange( + local_flats=rank_flats[r], + local_global_indices=rank_global_indices[r], + route=route, + all_global_indices=rank_global_indices, + all_tensor_meta=all_tensor_meta, + dp_rank=r, + dp_size=dp_size, + window_size=window_size, + ) + for r in range(dp_size) + ] + recv_bufs = _emulate_all_to_all(plans, dp_size) + return [ + reassemble_window( + rank_flats[d], + plans[d], + recv_bufs[d], + all_tensor_meta, + dp_size=dp_size, + window_size=window_size, + local=local, + ) + for d in range(dp_size) + ] + + +def test_reassemble_window_w1_matches_canonical_merge(): + """§10.1.1: with the intra route and W=1, ``reassemble_window`` rebuilds exactly the canonically + ordered balanced micro-batch — byte-identical to merging the globally-balanced samples directly.""" + images_per_sample = [2, 0, 1, 0, 3, 1, 0, 2] + dp_size, n_groups = 2, 2 + n = len(images_per_sample) + local = n // dp_size + + full = _make_mimo_batch_var(images_per_sample) + all_flats = split_microbatch(full, cu_img=_cu_img_from_counts(images_per_sample)) + costs = [sample_cost(f, encoder_cost_weight=1.0) for f in all_flats] + assignment = balanced_assignment(costs, n_groups, dp_size) + route = intra_route(assignment, src_slot=0) + pos2gidx = {pos: g for g, (_o, pos) in enumerate(assignment)} + + rank_global_indices = [list(range(r * local, (r + 1) * local)) for r in range(dp_size)] + rank_flats = [[all_flats[g] for g in rank_global_indices[r]] for r in range(dp_size)] + all_tensor_meta = [[tensor_metadata(f) for f in rank_flats[r]] for r in range(dp_size)] + + per_rank = _emulate_window_exchange( + rank_flats, rank_global_indices, all_tensor_meta, route, dp_size=dp_size, window_size=1, local=local + ) + + for d in range(dp_size): + assert len(per_rank[d]) == 1 # W=1 -> one micro-batch + owned_positions = sorted(pos for g, (owner, pos) in enumerate(assignment) if owner == d) + expected = merge_samples([all_flats[pos2gidx[p]] for p in owned_positions]) + _assert_batch_equal(per_rank[d][0], expected) + + +def test_reassemble_window_w2_3d_route(): + """§10.1.2: a hand-built W=2 3D route (no balancer) routes samples into specific slots; assert each + slot reassembles its sample, a text-only sample lands as an empty-vision slot, vision is intact.""" + # 4 samples, dp=2, W=2 slots, local=1 per (slot, rank). Layout: rank0 holds globals [0,1], + # rank1 holds [2,3]. g2 is text-only (0 patches); others carry distinct image patch counts. + dp_size, window_size, local = 2, 2, 1 + originals = [ + _flat_vision_sample(0, n_tokens=4, n_patches=3), + _flat_vision_sample(1, n_tokens=4, n_patches=2), + _flat_vision_sample(2, n_tokens=4, n_patches=0), # text-only + _flat_vision_sample(3, n_tokens=4, n_patches=1), + ] + rank_global_indices = [[0, 1], [2, 3]] + rank_flats = [[originals[g] for g in rank_global_indices[r]] for r in range(dp_size)] + all_tensor_meta = [[tensor_metadata(f) for f in rank_flats[r]] for r in range(dp_size)] + + # route[g] = (dst_slot, owner_rank, canonical_pos). Cross-rank moves into a chosen slot: + # g0 -> slot0,rank0 g1 -> slot1,rank1 g2 -> slot0,rank1 g3 -> slot1,rank0 + route = [(0, 0, 0), (1, 1, 0), (0, 1, 0), (1, 0, 0)] + + per_rank = _emulate_window_exchange( + rank_flats, rank_global_indices, all_tensor_meta, route, dp_size=dp_size, window_size=window_size, local=local + ) + + def first_token(batch): + return int(batch["input_ids"][0, 0].item()) + + # rank0: slot0 <- g0, slot1 <- g3 (received). rank1: slot0 <- g2 (text-only), slot1 <- g1. + assert [first_token(b) for b in per_rank[0]] == [0, 300] + assert [first_token(b) for b in per_rank[1]] == [200, 100] + + # Vision attribution per slot: g3 (1 patch) on rank0 slot1; g2 (text-only -> [0,4]) on rank1 slot0. + r0_slot1_enc = per_rank[0][1]["modality_inputs"]["images"]["enc"] + assert r0_slot1_enc["hidden_states"].shape == (1, 4) + assert torch.equal(r0_slot1_enc["hidden_states"], torch.full((1, 4), 3.0)) + r1_slot0_enc = per_rank[1][0]["modality_inputs"]["images"]["enc"] + assert r1_slot0_enc["hidden_states"].shape == (0, 4) + + +def test_reassemble_window_slot_shape_guard(): + """The §10.0 slot-shape invariant fires loudly when a slot ends up with the wrong sample count.""" + # Route both rank0's samples into the same slot of rank0 -> slot0 has 2, slot1 has 0 (expected 1). + dp_size, window_size, local = 2, 2, 1 + originals = [_flat_vision_sample(i, n_tokens=4, n_patches=1) for i in range(4)] + rank_global_indices = [[0, 1], [2, 3]] + rank_flats = [[originals[g] for g in rank_global_indices[r]] for r in range(dp_size)] + all_tensor_meta = [[tensor_metadata(f) for f in rank_flats[r]] for r in range(dp_size)] + # g0,g1 both -> slot0/rank0; g2,g3 -> slot0/rank1. Slot0 over-fills (2) -> guard fires on slot 0. + route = [(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1)] + with pytest.raises(RuntimeError, match="slot 0 reassembled 2 samples, expected 1"): + _emulate_window_exchange( + rank_flats, + rank_global_indices, + all_tensor_meta, + route, + dp_size=dp_size, + window_size=window_size, + local=local, + ) + + +# --------------------------------------------------------------------------- +# Window route (build_window_route) + multi-slot window exchange: Phase-2 gates +# --------------------------------------------------------------------------- + + +def test_build_window_route_layout(): + """The window route is slot-major: route[s*B + j] has dst_slot == s and a per-slot canonical_pos.""" + dp_size, n_groups = 2, 2 + costs_per_slot = [[4.0, 1.0, 3.0, 2.0], [1.0, 2.0, 3.0, 4.0], [2.0, 2.0, 2.0, 2.0]] + b_global = len(costs_per_slot[0]) + route = build_window_route(costs_per_slot, n_groups, dp_size) + assert len(route) == len(costs_per_slot) * b_global + for s in range(len(costs_per_slot)): + slot_entries = route[s * b_global : (s + 1) * b_global] + # intra invariant: every entry in this slot's block routes to slot s. + assert all(dst == s for dst, _owner, _pos in slot_entries) + # canonical positions are a permutation of [0, B) within the slot. + assert sorted(pos for _dst, _owner, pos in slot_entries) == list(range(b_global)) + # matches an independent per-slot intra_route. + assert slot_entries == intra_route(balanced_assignment(costs_per_slot[s], n_groups, dp_size), src_slot=s) + + +def _run_window_exchange(images_per_slot, dp_size, n_groups): + """Full window exchange over W micro-batches (emulated a2a) using the real window indexing + (build_window_route + reassemble_window). Asserts each slot reassembles to exactly the + independent per-slot intra-balanced shard — i.e. slots never mix (intra invariant).""" + window_size = len(images_per_slot) + # Per-slot full batches + flats + costs (slot-local global index j in [0, B)). + slot_full = [_make_mimo_batch_var(imgs) for imgs in images_per_slot] + slot_flats = [ + split_microbatch(slot_full[s], cu_img=_cu_img_from_counts(images_per_slot[s])) for s in range(window_size) + ] + b_global = len(slot_flats[0]) + local = b_global // dp_size + costs_per_slot = [[sample_cost(f, encoder_cost_weight=1.0) for f in slot_flats[s]] for s in range(window_size)] + + # Window-local layout: rank r holds slot-major [slot0 shard, slot1 shard, ...]. + rank_global_indices = [ + [s * b_global + r * local + i for s in range(window_size) for i in range(local)] for r in range(dp_size) + ] + rank_flats = [ + [slot_flats[s][r * local + i] for s in range(window_size) for i in range(local)] for r in range(dp_size) + ] + all_tensor_meta = [[tensor_metadata(f) for f in rank_flats[r]] for r in range(dp_size)] + + route = build_window_route(costs_per_slot, n_groups, dp_size) + per_rank = _emulate_window_exchange( + rank_flats, rank_global_indices, all_tensor_meta, route, dp_size=dp_size, window_size=window_size, local=local + ) + + # Expected: each slot independently balanced (W=1 semantics applied per slot). + for s in range(window_size): + assignment_s = balanced_assignment(costs_per_slot[s], n_groups, dp_size) + pos2gidx_s = {pos: g for g, (_o, pos) in enumerate(assignment_s)} + for d in range(dp_size): + owned = sorted(pos for g, (owner, pos) in enumerate(assignment_s) if owner == d) + expected = merge_samples([slot_flats[s][pos2gidx_s[p]] for p in owned]) + _assert_batch_equal(per_rank[d][s], expected) + + +def test_window_exchange_two_slots_dp2(): + # W=2 window, dp=2; distinct image mixes per slot so each slot triggers a real cross-rank swap. + _run_window_exchange([[2, 0, 1, 3], [1, 2, 0, 1]], dp_size=2, n_groups=2) + + +def test_window_exchange_w3_dp2_text_only_and_multi(): + # W=3 window with text-only + multi-image samples; proves §4a vision attribution per slot at W>1. + _run_window_exchange([[0, 2, 1, 0], [3, 0, 0, 1], [1, 1, 1, 1]], dp_size=2, n_groups=2) + + +def test_window_exchange_het_dp_paired_slots(): + # Het-DP at the window level: dp2 and dp4 over the same per-slot costs must agree per slot. + images_per_slot = [[2, 0, 1, 3, 1, 2, 0, 1], [1, 1, 2, 0, 3, 0, 1, 2]] + for dp_size in (2, 4): + _run_window_exchange(images_per_slot, dp_size=dp_size, n_groups=4) + + +def test_window_cost_spread_tightens_after_balance(): + """§10.5 balance probe: the per-rank cost spread must tighten (and never widen) after balancing — + the quantitative evidence that reordering evens the per-rank load (§10.3).""" + dp_size, n_groups = 4, 4 + # Heavily skewed costs: the natural contiguous shards are very imbalanced (rank0 holds the big + # samples, rank3 the small), so balancing should shrink the per-rank max/min spread. + # Both slots are skewed across the contiguous shards (rank0 heavy, rank1 light) so each must + # move samples and strictly tighten — slot 1 differs from slot 0 to exercise per-slot balancing. + costs_per_slot = [ + [100.0, 90.0, 80.0, 70.0, 8.0, 6.0, 4.0, 2.0, 50.0, 40.0, 30.0, 20.0, 9.0, 7.0, 5.0, 3.0], + [80.0, 75.0, 70.0, 65.0, 10.0, 9.0, 8.0, 7.0, 45.0, 40.0, 35.0, 30.0, 12.0, 11.0, 10.0, 9.0], + ] + route = build_window_route(costs_per_slot, n_groups, dp_size) + spreads = window_cost_spread(costs_per_slot, route, dp_size) + assert len(spreads) == 2 + for sp in spreads: + before = sp["before_max"] - sp["before_min"] + after = sp["after_max"] - sp["after_min"] + assert after <= before, sp # balancing never widens the spread + assert after < before # and on these skewed inputs it strictly tightens + # these skewed inputs require real movement: 0 < remote <= slot size. + assert all(0 < sp["remote"] <= len(costs_per_slot[0]) for sp in spreads) + + +def test_assert_intra_no_cross_slot_passes_for_window_route(): + # §10.0 guard: build_window_route output (intra) always satisfies dst_slot == src_slot. + dp_size, n_groups = 2, 2 + costs_per_slot = [[4.0, 1.0, 3.0, 2.0], [1.0, 2.0, 3.0, 4.0]] + route = build_window_route(costs_per_slot, n_groups, dp_size) # calls the guard internally + assert_intra_no_cross_slot(route, b_global=4) # explicit re-check is a no-op (no raise) + + +def test_assert_intra_no_cross_slot_trips_on_cross_slot_route(): + # A hand-built route that sends a slot-0 sample (g=1 -> src_slot 0) into slot 1 must trip the guard. + b_global = 2 # 2 slots of 2 samples; g//b_global is the source slot + route = [(0, 0, 0), (1, 1, 0), (1, 0, 0), (1, 1, 1)] # g=1 leaks slot0 -> slot1 + with pytest.raises(RuntimeError, match="cross-slot leak at global index 1"): + assert_intra_no_cross_slot(route, b_global=b_global) + + +def test_window_cost_spread_already_balanced_no_remote(): + # Uniform costs are already perfectly balanced: spread stays 0 and nothing needs to move. + dp_size, n_groups = 2, 2 + costs_per_slot = [[5.0, 5.0, 5.0, 5.0]] + route = build_window_route(costs_per_slot, n_groups, dp_size) + (sp,) = window_cost_spread(costs_per_slot, route, dp_size) + assert sp["before_max"] == sp["before_min"] == sp["after_max"] == sp["after_min"] == 10.0 + + +# --------------------------------------------------------------------------- +# ReorderingBuffer at dp_size == 1: local canonical reorder (no PG) + window cursor logic +# --------------------------------------------------------------------------- + + +def test_reordering_buffer_window_cursor_single_rank(): + # n_groups=1 + uniform cost -> identity order; exercises windows of [3, 3, 1] over 7 items. + items = [{"input_ids": torch.tensor([[i]])} for i in range(7)] + buf = ReorderingBuffer( + iter(items), + dp_rank=0, + dp_size=1, # single rank -> local reorder, no exchange / no PG needed + n_groups=1, + cost_of=lambda f: 0.0, + dp_group_gloo=None, + dp_group_nccl=None, + overlap=False, + window_size=3, + ) + got = [int(b["input_ids"].item()) for b in buf] + assert got == list(range(7)) # all items served, short final window handled + + +def test_reordering_buffer_overlap_single_rank_no_thread(): + # overlap=True but dp_size=1: no prefetch thread. + items = [{"input_ids": torch.tensor([[i]])} for i in range(5)] + buf = ReorderingBuffer( + iter(items), + dp_rank=0, + dp_size=1, + n_groups=1, + cost_of=lambda f: 0.0, + dp_group_gloo=None, + dp_group_nccl=None, + overlap=True, + window_size=2, + ) + assert buf._thread is None # dp_size <= 1 never starts the prefetch thread + assert [int(b["input_ids"].item()) for b in buf] == list(range(5)) + + +def test_reorder_window_local_matches_paired_dp2_order(): + # dp=1 local order must equal the paired dp=2 module's post-exchange order (rank 0's, then rank 1's). + n_groups = 2 + costs = [1.0, 5.0, 2.0, 4.0] # LPT over 2 groups: [1, 0] and [3, 2] -> canonical order [1, 0, 3, 2] + expected = [1, 0, 3, 2] + batch = { + "input_ids": torch.arange(4).view(4, 1), + "sample_ids": ["s0", "s1", "s2", "s3"], + # one image / one patch per sample + "vision": {"hidden_states": torch.arange(8.0).view(4, 2), "grid_thw": torch.tensor([[1, 1, 1]] * 4)}, + } + + def cost_of(flat): + return costs[int(flat["input_ids"].item())] + + (out,) = reorder_window_local( + [batch], + n_groups=n_groups, + cost_of=cost_of, + image_count_of=lambda b: torch.ones(4, dtype=torch.long), + ) + assert out["input_ids"].view(-1).tolist() == expected + assert out["sample_ids"] == [f"s{i}" for i in expected] + assert out["vision"]["hidden_states"].tolist() == [[2.0 * i, 2.0 * i + 1] for i in expected] + assert out["vision"]["grid_thw"].shape == (4, 3) + + # dp=2 order: rank 0's samples by canonical_pos, then rank 1's. + assignment = balanced_assignment(costs, n_groups, dp_size=2) + dp2_order = [ + g + for r in range(2) + for g in sorted((g for g in range(4) if assignment[g][0] == r), key=lambda g: assignment[g][1]) + ] + assert dp2_order == expected + + +def test_reordering_buffer_single_rank_serves_canonical_order(): + # Each micro-batch of a window is reordered independently; overlap=True runs inline at dp_size=1. + costs = [1.0, 5.0, 2.0, 4.0] + batches = [{"input_ids": (torch.arange(4) + 10 * s).view(4, 1)} for s in range(2)] + + def cost_of(flat): + return costs[int(flat["input_ids"].item()) % 10] + + buf = ReorderingBuffer( + iter(batches), + dp_rank=0, + dp_size=1, + n_groups=2, + cost_of=cost_of, + dp_group_gloo=None, + dp_group_nccl=None, + overlap=True, + window_size=2, + ) + assert buf._thread is None + assert [b["input_ids"].view(-1).tolist() for b in buf] == [[1, 0, 3, 2], [11, 10, 13, 12]] + + +def test_reordering_buffer_window_size_validation(): + with pytest.raises(ValueError, match="window_size must be >= 1"): + ReorderingBuffer( + iter([]), + dp_rank=0, + dp_size=1, + n_groups=1, + cost_of=lambda f: 0.0, + dp_group_gloo=None, + dp_group_nccl=None, + overlap=False, + window_size=0, + ) + + +# --------------------------------------------------------------------------- +# Joint per-sample cost +# --------------------------------------------------------------------------- + + +def test_sample_cost_patches_only(): + from megatron.bridge.data.megatron_mimo.reorder_buffer import sample_cost + + batch = _make_mimo_batch([2, 5]) # sample 0: 2 patches, sample 1: 5 patches; 6 tokens each + samples = split_microbatch(batch) + # default language_cost_weight=0.0 -> patch-only cost: encoder_cost_weight * sum(prod(grid_thw)); token length not counted + c0 = sample_cost(samples[0], encoder_cost_weight=1.0) + c1 = sample_cost(samples[1], encoder_cost_weight=1.0) + assert c0 == 2 + assert c1 == 5 + # encoder_cost_weight scales the patch count + assert sample_cost(samples[0], encoder_cost_weight=3.0) == 3 * 2 + + +def test_sample_cost_text_only_is_zero(): + from megatron.bridge.data.megatron_mimo.reorder_buffer import sample_cost + + flat = {"input_ids": torch.tensor([[5, 5, 0, 0]], dtype=torch.int64)} # no vision -> zero cost + c = sample_cost(flat, encoder_cost_weight=1.0) + assert c == 0 + + +def test_sample_cost_joint_vit_and_lm(): + from megatron.bridge.data.megatron_mimo.reorder_buffer import sample_cost + + # The LM term reads the real length from attention_mask (all 6 tokens real here). + batch = _make_mimo_batch([2, 5]) # patches 2 and 5; 6 real tokens each + samples = split_microbatch(batch) + for sample in samples: + sample["attention_mask"] = torch.ones(1, 6, dtype=torch.int64) + # cost = encoder_cost_weight*patches + language_cost_weight*real_tokens + assert sample_cost(samples[0], encoder_cost_weight=1.0, language_cost_weight=0.5) == 2 + 0.5 * 6 + assert sample_cost(samples[1], encoder_cost_weight=1.0, language_cost_weight=0.5) == 5 + 0.5 * 6 + # encoder_cost_weight=0 -> token-only cost + assert sample_cost(samples[0], encoder_cost_weight=0.0, language_cost_weight=2.0) == 2.0 * 6 + + +def test_sample_cost_lm_term_ignores_padding(): + from megatron.bridge.data.megatron_mimo.reorder_buffer import sample_cost + + # Same 3 real tokens but different padded widths -> identical LM cost (collation-independent). + short = { + "input_ids": torch.tensor([[5, 6, 7, 0]], dtype=torch.int64), + "attention_mask": torch.tensor([[1, 1, 1, 0]], dtype=torch.int64), + } + long = { + "input_ids": torch.tensor([[5, 6, 7, 0, 0, 0, 0, 0]], dtype=torch.int64), + "attention_mask": torch.tensor([[1, 1, 1, 0, 0, 0, 0, 0]], dtype=torch.int64), + } + assert sample_cost(short, encoder_cost_weight=1.0, language_cost_weight=1.0) == 3 + assert sample_cost(long, encoder_cost_weight=1.0, language_cost_weight=1.0) == 3 + + +def test_sample_cost_lm_term_does_not_compare_pad_id(): + from megatron.bridge.data.megatron_mimo.reorder_buffer import sample_cost + + # A real token equal to the pad id (pad_token == eos_token tokenizers) still counts: the length + # comes from attention_mask, never from comparing input_ids against a pad id. + flat = { + "input_ids": torch.tensor([[5, 0, 7, 0]], dtype=torch.int64), + "attention_mask": torch.tensor([[1, 1, 1, 0]], dtype=torch.int64), + } + assert sample_cost(flat, encoder_cost_weight=1.0, language_cost_weight=1.0) == 3 + + +def test_sample_cost_lm_term_requires_attention_mask(): + from megatron.bridge.data.megatron_mimo.reorder_buffer import sample_cost + + flat = {"input_ids": torch.tensor([[5, 6, 7, 0]], dtype=torch.int64)} + # Patch-only cost never needs the mask. + assert sample_cost(flat, encoder_cost_weight=1.0) == 0 + with pytest.raises(ValueError, match="attention_mask"): + sample_cost(flat, encoder_cost_weight=1.0, language_cost_weight=1.0) + + +def test_sample_cost_image_token_id_recovers_patches_from_input_ids(): + from megatron.bridge.data.megatron_mimo.reorder_buffer import sample_cost + + img = 99 # image-placeholder token id + # 3 placeholder tokens * square_merge_size(4) = 12 patches. + flat = {"input_ids": torch.tensor([[1, img, img, img, 2, 0]], dtype=torch.int64)} + assert sample_cost(flat, encoder_cost_weight=1.0, image_token_id=img, square_merge_size=4) == 12 + # encoder_cost_weight scales the recovered patch count. + assert sample_cost(flat, encoder_cost_weight=2.0, image_token_id=img, square_merge_size=4) == 24 + + +def test_sample_cost_image_token_id_matches_across_modules(): + """Bug #1: vision (has grid_thw) and language (grid_thw nulled) must derive the SAME cost. + + With image_token_id wired, the patch cost comes from the module-independent input_ids image-token + count, so a vision-style flat and a language-style flat with identical input_ids cost the same even + though the language flat has no grid_thw (nulled by the #4442 rank-aware metadata collate). + """ + from megatron.bridge.data.megatron_mimo.reorder_buffer import sample_cost + + img = 99 + input_ids = torch.tensor([[1, img, img, 2, 0, 0]], dtype=torch.int64) # 2 image tokens + # Vision-style flat: carries grid_thw (would dominate the legacy path). + vision_flat = { + "input_ids": input_ids, + f"{_VKEY}.grid_thw": torch.tensor([[1, 4, 2]], dtype=torch.int64), # prod = 8 patches + } + # Language-style flat: grid_thw nulled, only input_ids remains. + language_flat = {"input_ids": input_ids} + + vc = sample_cost(vision_flat, encoder_cost_weight=1.0, image_token_id=img, square_merge_size=4) + lc = sample_cost(language_flat, encoder_cost_weight=1.0, image_token_id=img, square_merge_size=4) + assert vc == lc == 2 * 4 # 2 image tokens * square_merge_size, independent of grid_thw + # Without image_token_id the two diverge (the bug): vision reads grid_thw, language gets 0. + assert sample_cost(vision_flat, encoder_cost_weight=1.0) == 8 + assert sample_cost(language_flat, encoder_cost_weight=1.0) == 0 + + +# --------------------------------------------------------------------------- +# F14 — single D2H on the vision reorder path (Task 4.3) +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestReorderVisionSingleD2H: + def test_reorder_vision_single_d2h(self, monkeypatch): + # Count Tensor.item() calls during the reorder; the cumsum should be a single .tolist(). + calls = {"item": 0} + orig_item = torch.Tensor.item + + def counting_item(self): + calls["item"] += 1 + return orig_item(self) + + grid = torch.tensor([[1, 2, 2], [1, 4, 4], [1, 2, 3]]) # patches 4,16,6 + hidden = torch.arange(26, dtype=torch.float32).reshape(26, 1) + monkeypatch.setattr(torch.Tensor, "item", counting_item) + hs, g = _reorder_vision_by_images(hidden, grid, [2, 0, 1]) + # no per-image .item() D2H sync in the loop (cumsum uses one .tolist()). + assert calls["item"] == 0 + # output bit-identical to the expected reorder + expected = torch.cat([hidden[20:26], hidden[0:4], hidden[4:20]]) + torch.testing.assert_close(hs, expected) + + +class _IndexDataset: + def __init__(self, n: int): + self.n = n + + def __len__(self) -> int: + return self.n + + def __getitem__(self, idx: int) -> int: + return idx + + +@pytest.mark.parametrize("dataloader_type", ["single", "cyclic"]) +def test_canonical_grid_shards_are_contiguous_blocks_of_one_global_order(dataloader_type): + """exchange_window assumes rank r holds global positions [r*local, (r+1)*local) of one ordered + micro-batch shared by every module. Under the canonical grid that holds for both samplers: a + rank's window is the concatenation of its groups' windows in group order, and the global order + is the concatenation of all groups' windows. + """ + module_dps = [1, 2] + grid = canonical_grid_size(module_dps) # 2 + micro_batch = 8 + dataset = _IndexDataset(64) + + def rank_window(dp_size: int, dp_rank: int): + sampler = build_canonical_group_batch_sampler( + dataloader_type=dataloader_type, + dataset=dataset, + consumed_samples=0, + micro_batch_size=micro_batch, + grid_size=grid, + groups=covered_canonical_groups(dp_rank, dp_size, grid), + data_sharding=True, + drop_last=True, + ) + return next(iter(sampler)) + + global_order = rank_window(dp_size=1, dp_rank=0) # vision (dp=1) reads every group + assert len(global_order) == micro_batch + local = micro_batch // 2 + for r in range(2): # language (dp=2) + assert rank_window(dp_size=2, dp_rank=r) == global_order[r * local : (r + 1) * local] diff --git a/tests/unit_tests/examples/test_qwen35_vl_mimo_finetune.py b/tests/unit_tests/examples/test_qwen35_vl_mimo_finetune.py index 5de99fd6ff4..03e0a375c86 100644 --- a/tests/unit_tests/examples/test_qwen35_vl_mimo_finetune.py +++ b/tests/unit_tests/examples/test_qwen35_vl_mimo_finetune.py @@ -160,6 +160,11 @@ def test_dataset_config_enables_only_requested_or_known_validation_splits( do_validation=do_validation, pack_sequences_in_batch=False, scalable_dp=False, + intra_microbatch_reorder=False, + reorder_encoder_cost_weight=1.0, + reorder_language_cost_weight=0.0, + no_overlap_intra_microbatch_reorder=False, + reorder_window_size=1, ) ) config.validate() @@ -280,3 +285,44 @@ def test_qwen35_vl_mimo_rejects_truncated_visual_tokens(): module._adapt_qwen35_hf_batch(batch, spec, seq_length=4, pad_to_seq_length=True) finally: sys.modules.pop(name, None) + + +def test_batch_spec_language_non_first_stage_ships_cheap_fields(monkeypatch): + # Non-first language stage: input_ids/position_ids always shipped (reorder cost, MRoPE); labels and + # loss_mask only on the last stage. forward_step decides what the model receives. + name = "qwen35_vl_mimo_finetune_batch_spec_under_test" + module = _load_example_module(name) + try: + + class _PG: + def rank(self): + return 1 # stage 1 of 3: neither first nor last + + grid = SimpleNamespace(dim_names=("pp",), shape=[3], get_pg=lambda dims: _PG()) + monkeypatch.setattr(module, "_rank_grid_and_module", lambda grids: (grid, module.MIMO_LANGUAGE_MODULE_KEY)) + cfg = SimpleNamespace(model=SimpleNamespace(_grids={module.MIMO_LANGUAGE_MODULE_KEY: grid})) + spec = module._batch_spec_for_rank(cfg) + assert spec.input_ids is True and spec.position_ids is True + assert spec.labels is False and spec.loss_mask is False + assert spec.modality_inputs is False + finally: + sys.modules.pop(name, None) + + +def test_batch_spec_encoder_ships_inputs_and_pixels(monkeypatch): + name = "qwen35_vl_mimo_finetune_encoder_spec_under_test" + module = _load_example_module(name) + try: + + class _PG: + def rank(self): + return 0 + + grid = SimpleNamespace(dim_names=("pp",), shape=[1], get_pg=lambda dims: _PG()) + monkeypatch.setattr(module, "_rank_grid_and_module", lambda grids: (grid, "images")) + cfg = SimpleNamespace(model=SimpleNamespace(_grids={"images": grid})) + spec = module._batch_spec_for_rank(cfg) + assert spec.input_ids is True and spec.modality_inputs is True + assert spec.labels is False and spec.loss_mask is False and spec.position_ids is False + finally: + sys.modules.pop(name, None) diff --git a/tests/unit_tests/training/megatron_mimo/test_megatron_mimo_step.py b/tests/unit_tests/training/megatron_mimo/test_megatron_mimo_step.py index 943fe330543..1158f932058 100644 --- a/tests/unit_tests/training/megatron_mimo/test_megatron_mimo_step.py +++ b/tests/unit_tests/training/megatron_mimo/test_megatron_mimo_step.py @@ -290,6 +290,43 @@ def test_forward_step_packs_language_shard_with_mask_lengths(self, mock_unwrap, assert call_kwargs["position_ids"].tolist() == [[0, 1, 2, 0, 1]] assert call_kwargs["attention_mask"] is None + @patch("megatron.bridge.training.megatron_mimo_step._get_module_dp_info") + @patch("megatron.bridge.training.megatron_mimo_step.get_batch") + @patch("megatron.bridge.training.megatron_mimo_step.unwrap_megatron_mimo_model") + def test_forward_step_nulls_attention_mask_before_model(self, mock_unwrap, mock_get_batch, mock_dp_info): + """The tokenizer padding mask is a data-path length source only; the model must get None.""" + from megatron.bridge.training.megatron_mimo_step import forward_step + + mock_state = MagicMock() + mock_state.cfg.dataset = SimpleNamespace( + enable_in_batch_packing=False, + defer_in_batch_packing_to_step=False, + megatron_mimo_scalable_dp=True, + ) + mock_model = MagicMock() + mock_role = MagicMock() + mock_role.has_language_module = True + mock_role.has_modality_modules = False + mock_role.is_first_stage.return_value = True + mock_role.is_last_stage.return_value = False + mock_model.role = mock_role + mock_model.return_value = (torch.tensor([1.0]), None) + mock_unwrap.return_value = mock_model + mock_dp_info.return_value = (0, 1) + + mock_get_batch.return_value = { + "input_ids": torch.arange(16).reshape(4, 4), + "position_ids": torch.arange(4).repeat(4, 1), + "attention_mask": torch.ones(4, 4, dtype=torch.long), + "labels": None, + "loss_mask": None, + "modality_inputs": None, + } + + forward_step(mock_state, iter([]), mock_model) + + assert mock_model.call_args.kwargs["attention_mask"] is None + @patch("megatron.bridge.training.megatron_mimo_step._get_module_dp_info") @patch("megatron.bridge.training.megatron_mimo_step.get_batch") @patch("megatron.bridge.training.megatron_mimo_step.unwrap_megatron_mimo_model") From 4e2de5cb9f37470785b404fe52f9d6aaa733a077 Mon Sep 17 00:00:00 2001 From: "kayeon.song" Date: Wed, 30 Sep 2026 07:44:18 +0000 Subject: [PATCH 02/13] fix(data): reject intra-microbatch reorder with pad_to_seq_length=False; tidy reorder comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The reorder merge concatenates exchanged samples along the batch dim, so every rank must pad to the same sequence length; reject the combination at argument validation instead of failing inside torch.cat. Comment/docstring cleanup in reorder_buffer.py and its tests: drop internal design-doc references, fix a stale cross-reference, correct the resident-window memory note (3·W), and replace "legacy/historical" wording. Signed-off-by: kayeon.song Signed-off-by: Yoonsik Kim Signed-off-by: Chanwoo Park Co-authored-by: Yoonsik Kim Co-authored-by: Chanwoo Park <45866990+sailor1493@users.noreply.github.com> --- .../qwen35_vl/finetune_qwen35_vl.py | 6 ++ .../data/megatron_mimo/reorder_buffer.py | 62 +++++++++---------- .../bridge/training/train_megatron_mimo.py | 4 +- .../data/megatron_mimo/test_reorder_buffer.py | 36 +++++------ 4 files changed, 57 insertions(+), 51 deletions(-) diff --git a/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py b/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py index 7f60c25607d..b80ebc18f63 100644 --- a/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py +++ b/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py @@ -334,6 +334,12 @@ def _validate_mimo_batch_sizes( f"--micro-batch-size ({args.micro_batch_size})." ) + if args.intra_microbatch_reorder and not args.pad_to_seq_length: + raise ValueError( + "--intra-microbatch-reorder requires --pad-to-seq-length true: exchanged samples are concatenated " + "back into a micro-batch, so every rank must pad to the same sequence length." + ) + summaries = [] for name, parallelism in parallelism_config.module_parallelisms.items(): dp = parallelism.data_parallel_size diff --git a/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py b/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py index 816c337dbce..55873b1724e 100644 --- a/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py +++ b/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py @@ -27,7 +27,7 @@ - :func:`balanced_assignment` — the per-sample plan: for each global sample in a micro-batch, the module-local DP rank that should own it and its canonical position. - Built on :func:`megatron.bridge.data.datasets.packing_utils.balanced_index_order`, + Built on :func:`balanced_index_order` (same module), with het-DP handled via canonical ``n_groups = canonical_grid_size(module dp_sizes)`` (their LCM) (so the vision and language modules, which may have different DP sizes, derive an identical assignment and the ``BridgeCommunicator`` keeps vision replica *r* paired with @@ -210,9 +210,9 @@ def balanced_assignment(costs: List[float], n_groups: int, dp_size: int) -> List def intra_route(assignment: List[Tuple[int, int]], *, src_slot: int = 0) -> List[Tuple[int, int, int]]: - """Lift a single-slot 2D balanced assignment to the vehicle's 3D route. + """Lift a single-slot 2D balanced assignment to the 3D window route. - The inter-ready vehicle routes every sample by **three** coordinates + The window transport routes every sample by **three** coordinates ``route[g] = (dst_slot, owner_rank, canonical_pos)`` (the destination micro-batch slot, the module-local DP rank that owns it, and its canonical intra-slot position). The **intra** balancer never moves a sample across micro-batch slots, so ``dst_slot == src_slot`` for every @@ -257,18 +257,18 @@ def build_window_route(costs_per_slot: List[List[float]], n_groups: int, dp_size route: List[Tuple[int, int, int]] = [] for slot, costs in enumerate(costs_per_slot): route.extend(intra_route(balanced_assignment(costs, n_groups, dp_size), src_slot=slot)) - # §10.0 safety net: the intra balancer must never move a sample across slots. Cheap, always-on. + # Safety net: the intra balancer must never move a sample across slots. Cheap, always-on. assert_intra_no_cross_slot(route, len(costs_per_slot[0]) if costs_per_slot else 0) return route def assert_intra_no_cross_slot(route: List[Tuple[int, int, int]], b_global: int) -> None: - """§10.0 "no cross-slot leak" guard: every sample's ``dst_slot`` equals its source slot. + """Guard against cross-slot leaks: every sample's ``dst_slot`` equals its source slot. For the slot-major window layout (``g = src_slot · b_global + j``) the source slot of global index ``g`` is ``g // b_global``. The **intra** balancer keeps every sample in its own slot, so - a mismatch means slot plumbing (route build / binning) corrupted the slot dimension — the §12 - "primary safety net while W>1 plumbing is new". A future inter (cross-slot) balancer would route + a mismatch means slot plumbing (route build / binning) corrupted the slot dimension. A future + inter (cross-slot) balancer would route across slots by design and must **not** run this guard. Args: @@ -292,13 +292,13 @@ def assert_intra_no_cross_slot(route: List[Tuple[int, int, int]], b_global: int) def window_cost_spread( costs_per_slot: List[List[float]], route: List[Tuple[int, int, int]], dp_size: int ) -> List[Dict[str, float]]: - """Per-slot pre/post-balance per-rank cost spread + remote-sample count (the §10.5 balance probe). + """Per-slot pre/post-balance per-rank cost spread + remote-sample count (the balance probe). For each slot, totals the per-rank cost **before** balancing (the natural contiguous shard: rank ``r`` holds slot-local positions ``[r·local, (r+1)·local)``) and **after** balancing (each sample credited to its routed ``owner_rank``), plus how many samples change rank. Pure and cheap — used only to log how much the reorder tightens the per-rank load (``after`` max/min should narrow - toward 1.0), the quantitative balance evidence the design's §10.3 validation looks for. + toward 1.0). Args: costs_per_slot: ``costs_per_slot[s]`` = slot ``s``'s per-sample costs by slot-local position. @@ -475,7 +475,7 @@ def prepare_sample_exchange( Each sample carries its destination micro-batch slot (``dst_slot``) end-to-end so the receiver can bin it into the right slot bucket on reassembly (:func:`reassemble_window`); this is the - only genuinely new plumbing vs. the single-slot path (§4a of the inter-ready design). Sender + only genuinely new plumbing vs. the single-slot path. Sender and receiver order the moves identically by ``(dst_slot, canonical_pos, src_local_idx)`` so the packed bytes line up. @@ -593,7 +593,7 @@ def prepare_sample_exchange( # grid_thw: [n_img,3]}}}} # # ``split_microbatch`` turns it into B per-sample **flat** dicts (nesting flattened to -# dotted keys) so the Phase-0 serializer/all-to-all can move individual samples; after the +# dotted keys) so the serializer/all-to-all can move individual samples; after the # exchange ``merge_samples`` rebuilds the (rebalanced) micro-batch. Sample slicing/vision # attribution uses :func:`_apply_sample_dispatch` (LM by batch dim, vision by image via # ``image_counts``). @@ -686,7 +686,7 @@ def _reorder_vision_by_images( # (avoids ``torch.cat([])`` and preserves dtype/device/requires_grad). return hidden_states[:0], grid_thw[:0] patches = patches_per_image(grid_thw) # [n_images] - # Single host copy of the cumulative offsets (F14): materialize the cumsum once with one + # Single host copy of the cumulative offsets: materialize the cumsum once with one # .tolist(), then index with plain Python ints — no per-image .item() D2H sync in the loop. cu = torch.cat([patches.new_zeros(1), patches.cumsum(0)]).tolist() blocks = [hidden_states[cu[i] : cu[i + 1]] for i in image_perm] @@ -710,15 +710,15 @@ def _gather_vision_subdict( :func:`_reorder_vision_by_images` (which keeps ``hidden_states`` patch-dim and ``grid_thw`` image-dim aligned). A text-only sample contributes an empty ``[0, d]`` / ``[0, 3]`` block. - When ``cu_img`` is ``None`` the legacy one-image-per-sample mapping (image ``i`` ↔ sample ``i``) + When ``cu_img`` is ``None`` the one-image-per-sample fallback (image ``i`` ↔ sample ``i``) is used and requires ``n_images == n_samples``. Args: value: A sub-dict carrying ``hidden_states`` and ``grid_thw``. sample_indices: Per-sample permutation/selection (global sample ids). n_samples: Global batch sample count ``B``. - cu_img: Cumulative per-sample image offsets (length ``B + 1``), or ``None`` for the legacy - one-image-per-sample path. + cu_img: Cumulative per-sample image offsets (length ``B + 1``), or ``None`` for the + one-image-per-sample fallback. Returns: A new sub-dict with permuted ``hidden_states`` / ``grid_thw`` (other keys carried). @@ -763,7 +763,7 @@ def _apply_sample_dispatch( Used by :func:`split_microbatch` (per-sample local-shard selection); ``n_samples`` is the canonical global batch size ``B`` so list/vision per-sample detection is consistent. ``cu_img`` (cumulative per-sample image offsets) is forwarded to :func:`_gather_vision_subdict` for the - variable-images-per-sample reorder; ``None`` keeps the legacy one-image-per-sample path. + variable-images-per-sample reorder; ``None`` selects the one-image-per-sample fallback. """ idx = torch.as_tensor(sample_indices, dtype=torch.long) out: Dict[str, Any] = {} @@ -795,8 +795,8 @@ def split_microbatch(batch: Dict[str, Any], *, cu_img: "List[int] | None" = None so :func:`merge_samples` is a plain concat. ``cu_img`` (cumulative per-sample image offsets, length ``B + 1``) lets the vision gather support - a variable number of images per sample (0 = text-only, 1, or N). When ``None``, the legacy - one-image-per-sample path is used. + a variable number of images per sample (0 = text-only, 1, or N). When ``None``, the + one-image-per-sample fallback is used. """ b = _batch_size(batch) return [_flatten_nested(_apply_sample_dispatch(batch, [i], n_samples=b, cu_img=cu_img)) for i in range(b)] @@ -922,10 +922,10 @@ def sample_cost( reorder is active, so vision and language derive the same cost. This is the correct source: the rank-aware metadata collate (#4442) nulls ``modality_inputs``/``grid_thw`` on language shards, so reading ``grid_thw`` there yields ``p = 0`` and mispairs the vision↔language fan-out. - - Otherwise ``p = Σ prod(grid_thw)`` (the legacy patch-only path, used by mock/tests where no + - Otherwise ``p = Σ prod(grid_thw)`` (the ``grid_thw`` fallback, used by mock/tests where no ``image_token_id`` is wired and ``grid_thw`` is present on the sample). - With the default ``language_cost_weight=0.0`` the cost is patch-only (the historical behavior). + With the default ``language_cost_weight=0.0`` the cost is patch-only. Args: flat: One per-sample flat (dotted-key) dict. @@ -977,7 +977,7 @@ def reassemble_window( window_size: int, local: int, ) -> List[Dict[str, Any]]: - """Rebuild this rank's ``W`` micro-batches from kept-local + received samples (§4a). + """Rebuild this rank's ``W`` micro-batches from kept-local + received samples. Generalizes the single-micro-batch reassembler to ``W`` slot-buckets: kept-local samples and deserialized received samples are binned by ``dst_slot``, sorted **within each slot** by @@ -985,7 +985,7 @@ def reassemble_window( via :func:`merge_samples`. ``deserialize_sample`` / ``merge_samples`` are reused verbatim — the only new step is the per-slot binning. - The §10.0 invariant asserts are baked in (cheap, always on): they turn a silent slot mispair or + The slot invariant asserts are baked in (cheap, always on): they turn a silent slot mispair or a dropped/duplicated sample into a loud failure at the exact rank. Args: @@ -1045,7 +1045,7 @@ def exchange_by_route( ) -> List[Dict[str, Any]]: """Transport-only per-sample exchange: route -> one ``all_to_all_single`` -> ``W`` micro-batches. - The inter-ready seam (§3.2): given a 3D ``route`` and the all-gathered layout/metadata, build the + Given a 3D ``route`` and the all-gathered layout/metadata, build the plan (:func:`prepare_sample_exchange`), run **one** ragged ``all_to_all_single``, and reassemble ``W`` micro-batches per slot (:func:`reassemble_window`). Knows nothing about the balancer — an intra route (``dst_slot == src_slot``, see :func:`intra_route`) and a future inter route flow @@ -1140,9 +1140,9 @@ def exchange_window( nccl_stream: Optional CUDA stream to run the all-to-all on (cross-step overlap). image_count_of: Optional ``callable(batch) -> LongTensor[local]`` giving a shard's per-sample image count, enabling a variable number of images per sample (0/1/N). When ``None``, the - legacy one-image-per-sample path is used. + one-image-per-sample fallback is used. probe: When ``True``, log a one-line per-slot cost-spread digest (:func:`window_cost_spread`) - on ``dp_rank == 0`` — the §10.5 balance probe. Off by default (the caller throttles it). + on ``dp_rank == 0`` (the balance probe). Off by default (the caller throttles it). Returns: ``W`` rebalanced micro-batches for this rank, slot ``s`` at index ``s``. @@ -1291,8 +1291,8 @@ class ReorderingBuffer: **Window buffering** (``window_size = W``): collect ``W`` micro-batches, cost-balance them in one exchange over the window (:func:`exchange_window`), and serve the ``W`` rebalanced micro-batches - one at a time from a cursor. ``W = 1`` (default) is the historical per-micro-batch behavior, - byte-for-byte. ``W == GA`` (the gradient-accumulation count) is the desired setting — the single + one at a time from a cursor. ``W = 1`` (default) exchanges every micro-batch on its own. + ``W == GA`` (the gradient-accumulation count) is the desired setting — the single window collective then lands once per optimizer step. **Cross-window prefetch overlap** (``overlap=True``): a background thread runs the *entire next @@ -1306,8 +1306,8 @@ class ReorderingBuffer: ahead, so all ranks' Gloo/NCCL collectives stay in lockstep — re-synced each window by the optimizer/DDP all-reduce barrier between consumptions. The worker uses the (separate) ``dp_group_nccl`` exclusively, so its all-to-all never races the main thread's bridge/DDP - collectives on their own PGs (requires ``CUDA_DEVICE_MAX_CONNECTIONS != 1``). Two windows stay - resident (consuming + prefetched) ≈ ``2·W`` micro-batches — the headline memory cost. + collectives on their own PGs (requires ``CUDA_DEVICE_MAX_CONNECTIONS != 1``). Up to three windows + stay resident (consuming + queued + being exchanged) ≈ ``3·W`` micro-batches — the headline memory cost. With ``overlap=False`` the window exchange runs synchronously in ``__next__`` (no thread). """ @@ -1344,7 +1344,7 @@ def __init__( # rebalanced micro-batches and the cursor serving them one at a time. self._active: List[Dict[str, Any]] = [] self._cursor = 0 - # Count of window exchanges run, for throttling the §10.5 balance probe (first few + every 50th). + # Count of window exchanges run, for throttling the balance probe (first few + every 50th). self._exchanges = 0 # Cross-window prefetch: a thread exchanges window t+1 while the main thread computes # window t. Started for any window_size (W == 1 is the one-step-ahead per-micro-batch case). @@ -1385,7 +1385,7 @@ def __del__(self) -> None: logger.debug("Ignoring exception while shutting down ReorderingBuffer.", exc_info=True) def _exchange_window(self, batches: List[Dict[str, Any]], stream: "Any") -> List[Dict[str, Any]]: - # Throttle the §10.5 balance probe: log the first few windows, then every 50th. + # Throttle the balance probe: log the first few windows, then every 50th. self._exchanges += 1 probe = self._exchanges <= 3 or self._exchanges % 50 == 0 return exchange_window( diff --git a/src/megatron/bridge/training/train_megatron_mimo.py b/src/megatron/bridge/training/train_megatron_mimo.py index 817450cfe42..9546b7bd10c 100644 --- a/src/megatron/bridge/training/train_megatron_mimo.py +++ b/src/megatron/bridge/training/train_megatron_mimo.py @@ -476,8 +476,8 @@ def train_megatron_mimo( timers("interval-time").stop() - # Release the reorder prefetch thread and its side process groups when the train iterator is a - # ReorderingBuffer; a plain iterator has no shutdown() and is left alone. + # Stop the reorder prefetch thread when the train iterator is a ReorderingBuffer; a plain + # iterator has no shutdown() and is left alone. if hasattr(train_data_iterator, "shutdown"): train_data_iterator.shutdown() diff --git a/tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py b/tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py index c16436efe3b..7a1e0c48599 100644 --- a/tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py +++ b/tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for the per-sample reorder-buffer building blocks (Phase 0, no process group). +"""Unit tests for the per-sample reorder-buffer building blocks (no process group). The all-to-all is emulated on CPU: each rank's ``send_buf`` is split by ``send_splits`` into per-destination chunks, and each destination's ``recv_buf`` is the concatenation of the chunks @@ -421,7 +421,7 @@ def test_split_merge_roundtrip_identity(): # --------------------------------------------------------------------------- -# merge_samples robustness to per-sample-varying vision (PR-A) +# merge_samples robustness to per-sample-varying vision # --------------------------------------------------------------------------- _VKEY = "modality_inputs.images.enc" @@ -430,7 +430,7 @@ def test_split_merge_roundtrip_identity(): def _flat_vision_sample(idx: int, n_tokens: int, n_patches: int, *, none_vision: bool = False) -> dict: """Per-sample flat dict as ``split_microbatch`` produces: ``[1, T]`` input_ids and one image's ``[p, 4]`` hidden_states + ``[1, 3]`` grid_thw. ``n_patches==0`` is text-only: empty ``[0, 4]`` / - ``[0, 3]`` (or ``None`` when ``none_vision`` to exercise the legacy/None merge path).""" + ``[0, 3]`` (or ``None`` when ``none_vision`` to exercise the None-vision merge path).""" input_ids = torch.arange(idx * 100, idx * 100 + n_tokens, dtype=torch.int64).reshape(1, -1) if n_patches == 0: if none_vision: @@ -481,7 +481,7 @@ def test_serialize_empty_vision_roundtrip(): # --------------------------------------------------------------------------- -# Variable images per sample: split / merge / exchange (PR-B) +# Variable images per sample: split / merge / exchange # --------------------------------------------------------------------------- @@ -564,8 +564,8 @@ def test_soft_validation_mismatch(): split_microbatch(batch, cu_img=[0, 1, 5]) # sums to 5 != 2 images -def test_split_legacy_no_cu_img(): - # one-image-per-sample works with cu_img=None (legacy path) +def test_split_one_image_per_sample_no_cu_img(): + # one-image-per-sample works with cu_img=None (fallback path) batch = _make_mimo_batch([2, 5, 1, 3]) _assert_batch_equal(merge_samples(split_microbatch(batch)), batch) # multi-image without cu_img raises a clear error @@ -631,7 +631,7 @@ def test_exchange_var_images_dp2(): # --------------------------------------------------------------------------- -# 3D routing seam + per-slot reassembly (reassemble_window): Phase-1 gates §10.1.1 / §10.1.2 +# 3D routing + per-slot reassembly (reassemble_window) # --------------------------------------------------------------------------- @@ -667,7 +667,7 @@ def _emulate_window_exchange(rank_flats, rank_global_indices, all_tensor_meta, r def test_reassemble_window_w1_matches_canonical_merge(): - """§10.1.1: with the intra route and W=1, ``reassemble_window`` rebuilds exactly the canonically + """With the intra route and W=1, ``reassemble_window`` rebuilds exactly the canonically ordered balanced micro-batch — byte-identical to merging the globally-balanced samples directly.""" images_per_sample = [2, 0, 1, 0, 3, 1, 0, 2] dp_size, n_groups = 2, 2 @@ -697,7 +697,7 @@ def test_reassemble_window_w1_matches_canonical_merge(): def test_reassemble_window_w2_3d_route(): - """§10.1.2: a hand-built W=2 3D route (no balancer) routes samples into specific slots; assert each + """A hand-built W=2 3D route (no balancer) routes samples into specific slots; assert each slot reassembles its sample, a text-only sample lands as an empty-vision slot, vision is intact.""" # 4 samples, dp=2, W=2 slots, local=1 per (slot, rank). Layout: rank0 holds globals [0,1], # rank1 holds [2,3]. g2 is text-only (0 patches); others carry distinct image patch counts. @@ -736,7 +736,7 @@ def first_token(batch): def test_reassemble_window_slot_shape_guard(): - """The §10.0 slot-shape invariant fires loudly when a slot ends up with the wrong sample count.""" + """The slot-shape invariant fires loudly when a slot ends up with the wrong sample count.""" # Route both rank0's samples into the same slot of rank0 -> slot0 has 2, slot1 has 0 (expected 1). dp_size, window_size, local = 2, 2, 1 originals = [_flat_vision_sample(i, n_tokens=4, n_patches=1) for i in range(4)] @@ -758,7 +758,7 @@ def test_reassemble_window_slot_shape_guard(): # --------------------------------------------------------------------------- -# Window route (build_window_route) + multi-slot window exchange: Phase-2 gates +# Window route (build_window_route) + multi-slot window exchange # --------------------------------------------------------------------------- @@ -823,7 +823,7 @@ def test_window_exchange_two_slots_dp2(): def test_window_exchange_w3_dp2_text_only_and_multi(): - # W=3 window with text-only + multi-image samples; proves §4a vision attribution per slot at W>1. + # W=3 window with text-only + multi-image samples; vision attribution per slot at W>1. _run_window_exchange([[0, 2, 1, 0], [3, 0, 0, 1], [1, 1, 1, 1]], dp_size=2, n_groups=2) @@ -835,8 +835,8 @@ def test_window_exchange_het_dp_paired_slots(): def test_window_cost_spread_tightens_after_balance(): - """§10.5 balance probe: the per-rank cost spread must tighten (and never widen) after balancing — - the quantitative evidence that reordering evens the per-rank load (§10.3).""" + """Balance probe: the per-rank cost spread must tighten (and never widen) after balancing — + the quantitative evidence that reordering evens the per-rank load.""" dp_size, n_groups = 4, 4 # Heavily skewed costs: the natural contiguous shards are very imbalanced (rank0 holds the big # samples, rank3 the small), so balancing should shrink the per-rank max/min spread. @@ -859,7 +859,7 @@ def test_window_cost_spread_tightens_after_balance(): def test_assert_intra_no_cross_slot_passes_for_window_route(): - # §10.0 guard: build_window_route output (intra) always satisfies dst_slot == src_slot. + # build_window_route output (intra) always satisfies dst_slot == src_slot. dp_size, n_groups = 2, 2 costs_per_slot = [[4.0, 1.0, 3.0, 2.0], [1.0, 2.0, 3.0, 4.0]] route = build_window_route(costs_per_slot, n_groups, dp_size) # calls the guard internally @@ -1090,7 +1090,7 @@ def test_sample_cost_image_token_id_recovers_patches_from_input_ids(): def test_sample_cost_image_token_id_matches_across_modules(): - """Bug #1: vision (has grid_thw) and language (grid_thw nulled) must derive the SAME cost. + """Vision (has grid_thw) and language (grid_thw nulled) must derive the SAME cost. With image_token_id wired, the patch cost comes from the module-independent input_ids image-token count, so a vision-style flat and a language-style flat with identical input_ids cost the same even @@ -1100,7 +1100,7 @@ def test_sample_cost_image_token_id_matches_across_modules(): img = 99 input_ids = torch.tensor([[1, img, img, 2, 0, 0]], dtype=torch.int64) # 2 image tokens - # Vision-style flat: carries grid_thw (would dominate the legacy path). + # Vision-style flat: carries grid_thw (would dominate the grid_thw fallback). vision_flat = { "input_ids": input_ids, f"{_VKEY}.grid_thw": torch.tensor([[1, 4, 2]], dtype=torch.int64), # prod = 8 patches @@ -1117,7 +1117,7 @@ def test_sample_cost_image_token_id_matches_across_modules(): # --------------------------------------------------------------------------- -# F14 — single D2H on the vision reorder path (Task 4.3) +# Single D2H on the vision reorder path # --------------------------------------------------------------------------- From 4058f1b117d5c36d61955a0ec2e1411eaeb65077 Mon Sep 17 00:00:00 2001 From: "kayeon.song" Date: Wed, 30 Sep 2026 07:50:54 +0000 Subject: [PATCH 03/13] refactor(data): name the serialization byte alignment; trim reorder comments Rename _ALIGN/_pad_to_align to _SERIALIZED_TENSOR_BYTE_ALIGNMENT/ _alignment_padding_bytes and shorten the comments added with the reorder wiring. No behavior change. Signed-off-by: kayeon.song Signed-off-by: Yoonsik Kim Signed-off-by: Chanwoo Park Co-authored-by: Yoonsik Kim Co-authored-by: Chanwoo Park <45866990+sailor1493@users.noreply.github.com> --- .../qwen35_vl/finetune_qwen35_vl.py | 19 ++++------- .../bridge/data/builders/direct_hf_sft.py | 3 +- .../data/megatron_mimo/reorder_buffer.py | 33 +++++++++---------- .../bridge/training/megatron_mimo_step.py | 3 +- 4 files changed, 24 insertions(+), 34 deletions(-) diff --git a/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py b/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py index b80ebc18f63..3e6465829fb 100644 --- a/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py +++ b/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py @@ -12,8 +12,7 @@ resulting Qwen batch is adapted into the MIMO forward shape: - language inputs: ``input_ids``, MRoPE ``position_ids``, labels, loss mask, and the - tokenizer's ``attention_mask`` (a data-path length source; the forward step drops it - before the model) + tokenizer's ``attention_mask`` - image inputs: ``modality_inputs["images"]["qwen_visual"]`` Example 2-GPU smoke: @@ -281,9 +280,7 @@ def _batch_spec_for_rank(cfg: Any) -> MIMOBatchSpec: pp_size = _grid_dim_size(grid, "pp", 1) is_last_pp = pp_rank == pp_size - 1 - # Cheap [B, S] fields (input_ids, position_ids, attention_mask) are always shipped: the data path - # (reorder cost, packing lengths, MRoPE) reads them on every stage, and forward_step alone decides - # what the model receives. Only the expensive or last-stage-only fields are gated here. + # Cheap [B, S] fields are always shipped; forward_step decides what reaches the model. if module_name == MIMO_LANGUAGE_MODULE_KEY: return MIMOBatchSpec( input_ids=True, @@ -293,7 +290,6 @@ def _batch_spec_for_rank(cfg: Any) -> MIMOBatchSpec: modality_inputs=False, ) - # Only an encoder's first PP stage loads data, so nothing is stage-gated here. return MIMOBatchSpec( input_ids=True, position_ids=False, @@ -336,8 +332,8 @@ def _validate_mimo_batch_sizes( if args.intra_microbatch_reorder and not args.pad_to_seq_length: raise ValueError( - "--intra-microbatch-reorder requires --pad-to-seq-length true: exchanged samples are concatenated " - "back into a micro-batch, so every rank must pad to the same sequence length." + "--intra-microbatch-reorder requires --pad-to-seq-length true (exchanged samples must share one " + "sequence length)." ) summaries = [] @@ -959,8 +955,7 @@ def _build_data_iterators(cfg, _megatron_mimo_infra, *, train_state=None): # return an iterator (DataLoader is iterable but not itself an iterator). loader_iter: Iterator[dict[str, Any]] = iter(train_loader) - # Gate on the canonical grid (same on every rank), not this rank's DP size: dp=1 modules must - # still join the world-collective PG creation below and reorder locally to stay paired. + # Gate on the canonical grid so every rank (dp=1 included) takes the same branch. reorder_on = ( scalable_dp and bool(getattr(cfg.dataset, "megatron_mimo_intra_microbatch_reorder", False)) @@ -1335,13 +1330,13 @@ def _parse_args() -> argparse.Namespace: "--reorder-encoder-cost-weight", type=float, default=1.0, - help="Reorder cost weight of the encoder term (per image patch).", + help="Reorder cost weight per image patch.", ) parser.add_argument( "--reorder-language-cost-weight", type=float, default=0.0, - help="Reorder cost weight of the language term (per sequence token, image tokens included); 0 = encoder-only.", + help="Reorder cost weight per real token (image tokens included); 0 = encoder-only.", ) parser.add_argument( "--no-overlap-intra-microbatch-reorder", diff --git a/src/megatron/bridge/data/builders/direct_hf_sft.py b/src/megatron/bridge/data/builders/direct_hf_sft.py index 2a81b09cf5b..cb7c81b9568 100644 --- a/src/megatron/bridge/data/builders/direct_hf_sft.py +++ b/src/megatron/bridge/data/builders/direct_hf_sft.py @@ -95,8 +95,7 @@ class DirectHFSFTDatasetConfig(DataloaderConfig): in_batch_packing_pad_to_multiple_of: int = 1 megatron_mimo_scalable_dp: bool = False megatron_mimo_intra_microbatch_reorder: bool = False - # Reorder cost = encoder_weight * encoder work (image patches) + language_weight * sequence tokens - # (image placeholder tokens included); only the ratio matters. + # Reorder cost = encoder_weight * image patches + language_weight * real tokens; only the ratio matters. megatron_mimo_reorder_encoder_cost_weight: float = 1.0 megatron_mimo_reorder_language_cost_weight: float = 0.0 megatron_mimo_reorder_overlap: bool = True diff --git a/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py b/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py index 55873b1724e..c906c7e7bd4 100644 --- a/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py +++ b/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py @@ -79,12 +79,12 @@ # All sample tensors are serialized into uint8 byte buffers padded to this alignment, so # the receiver can view-cast any dtype back from the received bytes without alignment # errors and compute exact offsets from metadata alone. -_ALIGN = 8 +_SERIALIZED_TENSOR_BYTE_ALIGNMENT = 8 -def _pad_to_align(nbytes: int) -> int: - """Zero-padding bytes to round ``nbytes`` up to the next :data:`_ALIGN` multiple.""" - return (-nbytes) % _ALIGN +def _alignment_padding_bytes(nbytes: int) -> int: + """Zero-padding bytes to round ``nbytes`` up to the next :data:`_SERIALIZED_TENSOR_BYTE_ALIGNMENT` multiple.""" + return (-nbytes) % _SERIALIZED_TENSOR_BYTE_ALIGNMENT # String -> (torch.dtype, byte width), for reconstructing tensors from metadata shared via @@ -381,7 +381,7 @@ def serialize_sample(flat: Dict[str, Any], keys: List[str]) -> torch.Tensor: Returns: 1D uint8 CPU tensor with each tensor's raw bytes concatenated in ``keys`` order, - each padded to :data:`_ALIGN`. + each padded to :data:`_SERIALIZED_TENSOR_BYTE_ALIGNMENT`. """ parts = [] for key in keys: @@ -389,7 +389,7 @@ def serialize_sample(flat: Dict[str, Any], keys: List[str]) -> torch.Tensor: if t is None: continue raw = t.contiguous().view(torch.uint8).reshape(-1) - pad = _pad_to_align(raw.numel()) + pad = _alignment_padding_bytes(raw.numel()) if pad: raw = torch.cat([raw, torch.zeros(pad, dtype=torch.uint8, device=raw.device)]) parts.append(raw) @@ -405,7 +405,7 @@ def sample_byte_size(meta: Dict[str, Any], keys: List[str]) -> int: continue _, elem = _dtype_info(info["dtype"]) nbytes = math.prod(info["shape"]) * elem - nbytes += _pad_to_align(nbytes) + nbytes += _alignment_padding_bytes(nbytes) total += nbytes return total @@ -437,7 +437,7 @@ def deserialize_sample( dtype, elem = _dtype_info(info["dtype"]) nbytes = math.prod(info["shape"]) * elem flat[key] = buf[cursor : cursor + nbytes].view(dtype).reshape(info["shape"]).clone() - cursor += nbytes + _pad_to_align(nbytes) + cursor += nbytes + _alignment_padding_bytes(nbytes) for key, info in meta.items(): if key not in flat: flat[key] = info.get("non_tensor") if isinstance(info, dict) and "non_tensor" in info else None @@ -900,13 +900,12 @@ def sample_cost( ) -> float: """Joint per-sample cost: ``encoder_cost_weight·p + language_cost_weight·t``. - Only the ratio of the two weights matters (the balancer is scale-invariant). Each term is - linear in an intrinsic per-sample workload count: + Only the ratio of the two weights matters. Each term is linear in a per-sample count: - ``p`` is the vision patch count (mirrors vision-encoder FLOPs). - ``t`` is the **real (non-pad) token count** of ``input_ids`` (mirrors LM FLOPs), counted - via :func:`real_token_lengths` (from the batch's ``attention_mask``). Image placeholder - tokens are part of the language sequence and are included. + via :func:`real_token_lengths` (from the batch's ``attention_mask``), image placeholder + tokens included. Both terms must be **collation-independent** so vision and language derive an identical assignment with no cross-module communication: the patch count is intrinsic to the image, @@ -929,9 +928,8 @@ def sample_cost( Args: flat: One per-sample flat (dotted-key) dict. - encoder_cost_weight: Weight of the encoder term (per patch here; patch-grid encoders only, - other encoders need their own ``cost_of``). - language_cost_weight: Weight of the language term (per sequence token). ``0.0`` disables it. + encoder_cost_weight: Weight per vision patch. + language_cost_weight: Weight per real token; ``0.0`` disables the term. image_token_id: Placeholder token id whose count in ``input_ids`` is the module-independent patch proxy. ``None`` falls back to the ``grid_thw`` patch sum. square_merge_size: ``spatial_merge_size²`` — patches represented by one placeholder token. @@ -1232,8 +1230,7 @@ def reorder_window_local( ) -> List[Dict[str, Any]]: """Reorder a window into canonical order on a ``dp_size == 1`` module (no exchange needed). - Same costs + ``n_groups`` as the paired dp>1 module, so the result equals that module's - post-exchange order (rank 0's samples, then rank 1's, ...), keeping the batch-dim pairing. + Uses the same costs and ``n_groups`` as the paired dp>1 module, so both end in the same order. Args: batches: This rank's ``W`` whole micro-batches. @@ -1421,7 +1418,7 @@ def _refill_window(self) -> None: if not batches: self._active = [] elif self._dp_size <= 1: - # Nothing to exchange, but the paired dp>1 module is reordered: match its order locally. + # dp_size == 1: nothing to exchange, apply the canonical order locally. self._active = reorder_window_local( batches, n_groups=self._n_groups, cost_of=self._cost_of, image_count_of=self._image_count_of ) diff --git a/src/megatron/bridge/training/megatron_mimo_step.py b/src/megatron/bridge/training/megatron_mimo_step.py index 91ddf3b799f..5e2622a7a20 100644 --- a/src/megatron/bridge/training/megatron_mimo_step.py +++ b/src/megatron/bridge/training/megatron_mimo_step.py @@ -280,8 +280,7 @@ def forward_step( data_batch, packing_kwargs = pack_language_shard(data_batch, lengths=pack_lengths) if packing_kwargs is not None: data_batch["packing_kwargs"] = packing_kwargs - # The tokenizer padding mask is a data-path length source only (packing, reorder cost, MRoPE); - # the model runs causal attention and MimoModel rejects a dense mask under CP. + # The padding mask only feeds data-path lengths; the model runs causal attention and must get None. data_batch["attention_mask"] = None else: # Non-data stages consume hidden states from pipeline input tensors. From 67d27f9f05e034c3c657b6052b56d4ad1515536a Mon Sep 17 00:00:00 2001 From: "kayeon.song" Date: Wed, 30 Sep 2026 08:02:49 +0000 Subject: [PATCH 04/13] docs(training): add MegatronMIMO intra-microbatch reordering page Signed-off-by: kayeon.song Signed-off-by: Yoonsik Kim Signed-off-by: Chanwoo Park Co-authored-by: Yoonsik Kim Co-authored-by: Chanwoo Park <45866990+sailor1493@users.noreply.github.com> --- docs/training/README.md | 1 + .../training/mimo-intra-microbatch-reorder.md | 126 ++++++++++++++++++ 2 files changed, 127 insertions(+) create mode 100644 docs/training/mimo-intra-microbatch-reorder.md diff --git a/docs/training/README.md b/docs/training/README.md index 9f16c453308..b7c3b70ce17 100644 --- a/docs/training/README.md +++ b/docs/training/README.md @@ -62,6 +62,7 @@ This directory contains comprehensive documentation for training and customizing |----------|---------|--------------| | **[PEFT](peft.md)** | Parameter-Efficient Fine-Tuning (LoRA, etc.) | Fine-tuning with limited resources | | **[Packed Sequences](packed-sequences.md)** | Sequence packing for efficiency | Optimizing data loading | +| **[MegatronMIMO Intra-Microbatch Reordering](mimo-intra-microbatch-reorder.md)** | Balance per-sample vision load across module DP ranks | Non-colocated MIMO VLM finetuning with uneven image load | | **[Megatron FSDP](megatron-fsdp.md)** | Stable overview of Megatron FSDP | Choosing an FSDP path | | **[Distillation](../modelopt/distillation.md)** | Knowledge distillation techniques | Transferring knowledge between models | | **[Checkpointing](checkpointing.md)** | Checkpoint saving, loading, and resuming | Managing training state | diff --git a/docs/training/mimo-intra-microbatch-reorder.md b/docs/training/mimo-intra-microbatch-reorder.md new file mode 100644 index 00000000000..09385ab2ca4 --- /dev/null +++ b/docs/training/mimo-intra-microbatch-reorder.md @@ -0,0 +1,126 @@ +# MegatronMIMO Intra-Microbatch Reordering + +Intra-microbatch reordering rebalances the per-sample vision load of each +micro-batch across a module's data-parallel ranks in non-colocated +MegatronMIMO training. It removes the straggler that appears when one +data-parallel rank happens to draw the heaviest images every step. + +This page is the stable overview of what the feature does, how to enable it, +and which constraints are durable. The implementation lives in +`megatron.bridge.data.megatron_mimo.reorder_buffer`; the Qwen3.5-VL example +under `examples/megatron_mimo/qwen35_vl/` wires it end to end. + +## What It Is + +In a non-colocated MegatronMIMO run the vision encoder and the language model +run on disjoint rank groups, each with its own data-parallel size. With +scalable data parallelism (`megatron_mimo_scalable_dp`), every rank reads only +its own shard of each global micro-batch, so the images a rank must encode are +whatever the sampler happened to assign to it. Image sizes vary widely, so the +per-rank vision work varies too, and the slowest rank gates every step. + +The reorder buffer wraps the training data iterator and, for each micro-batch +(or window of micro-batches), does three things: + +1. Computes a per-sample cost on every rank from fields that are identical on + the vision and language modules: the image-placeholder token count in + `input_ids` (proportional to the number of vision patches) and, optionally, + the real token count from `attention_mask`. +2. All-gathers the costs over the module's data-parallel group and computes one + balanced assignment of samples to ranks. The assignment is a deterministic + function of the costs, so the vision and language modules compute the same + one without talking to each other. +3. Exchanges samples with a single ragged `all_to_all` so each rank ends up + holding its assigned samples in a canonical order. + +Because both modules apply the same assignment, sample *i* on vision replica +*r* is still sample *i* on language replica *r*, and the `BridgeCommunicator` +fan-out stays aligned. Balancing is done over the canonical grid (the least +common multiple of the module data-parallel sizes), so modules with different +data-parallel sizes, including a module with data-parallel size 1, agree on +the same order. A data-parallel-1 module has nothing to exchange and simply +permutes its whole micro-batch locally. + +By default the exchange for the next window runs on a background thread with a +dedicated NCCL process group, overlapping with the current window's compute. + +## When to Use It + +Reordering is a good fit when all of the following are true: + +- the run is non-colocated MegatronMIMO with `megatron_mimo_scalable_dp=True` +- at least one module has data-parallel size 2 or more +- image load varies across samples (document, receipt, or mixed text-and-image + SFT data), so per-rank vision time is the per-step straggler + +It does not help text-only training, colocated single-module runs, or data +where every sample carries the same number of similarly sized images. + +## Bridge Configuration + +All fields live on `DirectHFSFTDatasetConfig` and default to off. + +| Field | Default | Meaning | +|---|---|---| +| `megatron_mimo_scalable_dp` | `False` | Read only this rank's canonical-grid shard. Required by reordering. | +| `megatron_mimo_intra_microbatch_reorder` | `False` | Enable the reorder buffer. | +| `megatron_mimo_reorder_encoder_cost_weight` | `1.0` | Cost weight per vision patch. | +| `megatron_mimo_reorder_language_cost_weight` | `0.0` | Cost weight per real language token; `0.0` balances on vision load only. | +| `megatron_mimo_reorder_overlap` | `True` | Exchange the next window on a background thread. | +| `megatron_mimo_reorder_window_size` | `1` | Micro-batches balanced together in one exchange. | + +Only the ratio of the two cost weights matters; at least one must be positive. + +The Qwen3.5-VL example exposes the same fields as `--scalable-dp`, +`--intra-microbatch-reorder`, `--reorder-encoder-cost-weight`, +`--reorder-language-cost-weight`, `--no-overlap-intra-microbatch-reorder`, and +`--reorder-window-size`: + +```bash +python examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py \ + --component "language=tp=1,pp=1,dp=2,rank_offset=0" \ + --component "images=tp=1,dp=2,rank_offset=2" \ + --scalable-dp --intra-microbatch-reorder \ + ... +``` + +Overlap only hides the exchange when the side stream can run on its own +hardware queue; with `CUDA_DEVICE_MAX_CONNECTIONS=1` the buffer logs a warning +and the exchange runs serially behind compute. + +## Stable Constraints + +| Configuration | Status | +|---|---| +| Heterogeneous module data parallelism (for example vision dp 1 or 2 with language dp 2 or 4) | Supported. Balancing uses the canonical grid; a dp-1 module reorders locally. | +| Language pipeline parallelism > 1 | Supported. Every language stage receives `input_ids` so all stages compute the same assignment. | +| Encoder pipeline parallelism > 1 | Not supported; the example raises before training starts. | +| `--pad-to-seq-length false` (variable batch width) | Not supported with reordering; the example raises before training starts. Exchanged samples are concatenated back into one micro-batch, so all ranks must pad to the same length. | +| `micro_batch_size` not divisible by the canonical grid size | Rejected by the canonical loader. | +| `drop_last=False`, or a dataloader type other than `single` / `cyclic` | Rejected by scalable data parallelism. | +| Expert tensor parallelism > 1 | Rejected by scalable data parallelism. | +| Tensor parallelism > 1, context parallelism > 1 | Untested. | + +A real exchange needs a module with data-parallel size 2 or more, which +exceeds the two-GPU functional-test budget. CI covers the exchange arithmetic +with CPU-emulated all-to-all unit tests and a two-GPU smoke test; multi-rank +runs are validated manually. + +## Feature Interactions + +- **In-batch sequence packing** (`enable_in_batch_packing`): reordering runs in + the data iterator, before the forward step packs each language shard, so the + two compose. Packing drops padding, so the exchange transfers padded rows but + model compute is unaffected. +- **Checkpoint resume**: reordering changes which rank holds a sample, not how + many samples are consumed, so the resume arithmetic of the canonical sampler + is unchanged. +- **Cost weights and the language term**: setting + `megatron_mimo_reorder_language_cost_weight > 0` counts real tokens from + `attention_mask`; the example ships `attention_mask` on every stage for this + purpose and the forward step drops it before the model. + +## Related Docs + +- [Packed Sequences](packed-sequences.md) +- Example: `examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py` From 0477a8d93eec9a532968a4d7d9ed2ce51d1a815a Mon Sep 17 00:00:00 2001 From: "kayeon.song" Date: Wed, 30 Sep 2026 09:18:57 +0000 Subject: [PATCH 05/13] fix(training): unblock MegatronMIMO Qwen3.5 training on current MCore - Qwen3VLGPTModel.forward: accept `mtp_input_mask` and forward it to `_postprocess`, matching GPTModel.forward. MCore's MimoModel passes this kwarg unconditionally since NVIDIA/Megatron-LM#6770, so the overridden forward raised TypeError on the first step. - training_log: read `cuda_graph_impl` with getattr. MegatronMIMOProvider is not a TransformerConfig and has no such field, so the memory report added in #6003 raised AttributeError after iteration 1. Signed-off-by: kayeon.song Signed-off-by: Yoonsik Kim Signed-off-by: Chanwoo Park Co-authored-by: Yoonsik Kim Co-authored-by: Chanwoo Park <45866990+sailor1493@users.noreply.github.com> --- .../bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py | 2 ++ src/megatron/bridge/training/utils/train_utils.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py index 79dcba93e75..e9a5b93a269 100644 --- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py +++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py @@ -166,6 +166,7 @@ def forward( *, inference_params: Optional[BaseInferenceContext] = None, loss_mask: Optional[Tensor] = None, + mtp_input_mask: Optional[Tensor] = None, padding_mask: Optional[Tensor] = None, # args for deepstack visual_pos_masks: Optional[torch.Tensor] = None, @@ -265,6 +266,7 @@ def _sp_scatter_embedding(input_ids, position_ids): rotary_pos_sin=rotary_pos_sin, mtp_in_postprocess=self.mtp_process, loss_mask=loss_mask, + mtp_input_mask=mtp_input_mask, decoder_input=decoder_input, attention_mask=attention_mask, padding_mask=padding_mask, diff --git a/src/megatron/bridge/training/utils/train_utils.py b/src/megatron/bridge/training/utils/train_utils.py index 800a6da97f4..db6c96e337e 100644 --- a/src/megatron/bridge/training/utils/train_utils.py +++ b/src/megatron/bridge/training/utils/train_utils.py @@ -1214,7 +1214,7 @@ def training_log( if torch.distributed.get_rank(group=pg_collection.dp) == 0: print("[Rank {}] {}".format(torch.distributed.get_rank(), memory_string), flush=True) cuda_graphs_enabled = ( - config.model.cuda_graph_impl != "none" + getattr(config.model, "cuda_graph_impl", "none") != "none" or config.optimizer.optimizer_cuda_graph or getattr(config.model, "vision_cuda_graph_impl", None) == "transformer_engine" ) From 4a241bbd84c89e3e91a4065872d1fa935561858d Mon Sep 17 00:00:00 2001 From: "kayeon.song" Date: Sat, 3 Oct 2026 02:44:16 +0000 Subject: [PATCH 06/13] fix(data): refuse intra-microbatch reorder overlap under CUDA_DEVICE_MAX_CONNECTIONS=1 With a single hardware queue the prefetch thread's side-PG all-to-all and the main thread's DDP all-reduce execute strictly in issue order, and the two threads issue in rank-dependent order. One rank then queues [all-to-all, all-reduce] while a peer queues [all-reduce, all-to-all], each NCCL kernel waits for a peer kernel that cannot start, and training hangs. Reproduced twice: iteration 290 on a 4-rank dp2/dp2 run and iteration 33 on a 6-rank TP=2 run; the non-overlapped path completed 1500 iterations under the same variable. The previous warning ("correct but no faster") was wrong. - ReorderingBuffer: raise ValueError for overlap=True under the variable (dp_size > 1) instead of logging a warning; docstring updated. - Qwen3.5-VL example: reject the combination in _validate_mimo_batch_sizes so every rank fails before model build; the --no-overlap path is allowed. - Docs: replace the warning description with the refusal and the TP choice. - Unit tests for the buffer guard (refuse / allow without overlap / allow without the variable) and the example validate. Signed-off-by: kayeon.song Signed-off-by: Yoonsik Kim Signed-off-by: Chanwoo Park Co-authored-by: Yoonsik Kim Co-authored-by: Chanwoo Park <45866990+sailor1493@users.noreply.github.com> --- .../training/mimo-intra-microbatch-reorder.md | 13 +++++-- .../qwen35_vl/finetune_qwen35_vl.py | 10 ++++++ .../data/megatron_mimo/reorder_buffer.py | 28 ++++++++------- .../data/megatron_mimo/test_reorder_buffer.py | 23 ++++++++++++ .../examples/test_qwen35_vl_mimo_finetune.py | 36 +++++++++++++++++++ 5 files changed, 94 insertions(+), 16 deletions(-) diff --git a/docs/training/mimo-intra-microbatch-reorder.md b/docs/training/mimo-intra-microbatch-reorder.md index 09385ab2ca4..d346f690140 100644 --- a/docs/training/mimo-intra-microbatch-reorder.md +++ b/docs/training/mimo-intra-microbatch-reorder.md @@ -84,9 +84,16 @@ python examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py \ ... ``` -Overlap only hides the exchange when the side stream can run on its own -hardware queue; with `CUDA_DEVICE_MAX_CONNECTIONS=1` the buffer logs a warning -and the exchange runs serially behind compute. +Overlap is refused under `CUDA_DEVICE_MAX_CONNECTIONS=1`: with a single +hardware queue the prefetch thread's all-to-all and the DDP all-reduce execute +in issue order, that order differs across ranks, and training deadlocks +(reproduced on a 4-rank run and on a 6-rank TP=2 run). Use +`--no-overlap-intra-microbatch-reorder` in that case; it issues both +collectives from the training thread and completed 1500 iterations under the +same setting. Megatron Bridge does not require the variable for tensor +parallelism; it is a Megatron-LM sequence-parallel performance setting, so TP +users choose between that setting with the non-overlapped reorder, or overlap +with the variable unset. ## Stable Constraints diff --git a/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py b/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py index 3e6465829fb..10f21d6e518 100644 --- a/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py +++ b/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py @@ -335,6 +335,16 @@ def _validate_mimo_batch_sizes( "--intra-microbatch-reorder requires --pad-to-seq-length true (exchanged samples must share one " "sequence length)." ) + if ( + args.intra_microbatch_reorder + and not args.no_overlap_intra_microbatch_reorder + and os.environ.get("CUDA_DEVICE_MAX_CONNECTIONS") == "1" + ): + raise ValueError( + "--intra-microbatch-reorder with overlap deadlocks under CUDA_DEVICE_MAX_CONNECTIONS=1 (single " + "hardware queue serializes the prefetch all-to-all and the DDP all-reduce in rank-dependent " + "order). Add --no-overlap-intra-microbatch-reorder or unset the variable." + ) summaries = [] for name, parallelism in parallelism_config.module_parallelisms.items(): diff --git a/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py b/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py index c906c7e7bd4..e25f031df66 100644 --- a/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py +++ b/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py @@ -1303,7 +1303,8 @@ class ReorderingBuffer: ahead, so all ranks' Gloo/NCCL collectives stay in lockstep — re-synced each window by the optimizer/DDP all-reduce barrier between consumptions. The worker uses the (separate) ``dp_group_nccl`` exclusively, so its all-to-all never races the main thread's bridge/DDP - collectives on their own PGs (requires ``CUDA_DEVICE_MAX_CONNECTIONS != 1``). Up to three windows + collectives on their own PGs. ``CUDA_DEVICE_MAX_CONNECTIONS=1`` is refused for ``overlap=True``: + it serializes both onto one hardware queue and the rank-dependent issue order deadlocks. Up to three windows stay resident (consuming + queued + being exchanged) ≈ ``3·W`` micro-batches — the headline memory cost. With ``overlap=False`` the window exchange runs synchronously in ``__next__`` (no thread). @@ -1345,19 +1346,20 @@ def __init__( self._exchanges = 0 # Cross-window prefetch: a thread exchanges window t+1 while the main thread computes # window t. Started for any window_size (W == 1 is the one-step-ahead per-micro-batch case). + if overlap and dp_size > 1 and os.environ.get("CUDA_DEVICE_MAX_CONNECTIONS") == "1": + # With CUDA_DEVICE_MAX_CONNECTIONS=1 every stream maps onto one hardware queue, so the + # worker's side-PG all-to-all and the main thread's DDP all-reduce execute strictly in + # issue order. The two threads issue in rank-dependent order, so one rank can queue + # [a2a, all-reduce] while a peer queues [all-reduce, a2a]; each NCCL kernel then waits + # for a peer kernel that cannot start, and training deadlocks (reproduced at iteration + # 290 on 4 ranks and at iteration 33 with TP=2 on 6 ranks). Refuse rather than degrade; + # the overlap=False path issues both collectives from one thread and is safe. + raise ValueError( + "ReorderingBuffer overlap=True deadlocks under CUDA_DEVICE_MAX_CONNECTIONS=1 (single " + "hardware queue: the side-PG all-to-all and the DDP all-reduce can wait on each other). " + "Pass overlap=False (--no-overlap-intra-microbatch-reorder) or unset the variable." + ) if overlap and self._cuda and dp_size > 1: - # The side-stream exchange only overlaps compute when it can use a separate hardware - # queue. With CUDA_DEVICE_MAX_CONNECTIONS=1 every stream serializes onto one queue, so - # the prefetch a2a runs *behind* the window's compute instead of beside it — correct, but - # silently no faster (often slower, for the extra thread + side PG). Warn loudly rather - # than appear to overlap. Not a hard error: a user may set =1 for TP/SP comm overlap and - # still want the (degraded) reorder to run. - if os.environ.get("CUDA_DEVICE_MAX_CONNECTIONS") == "1": - logger.warning( - "ReorderingBuffer overlap is enabled but CUDA_DEVICE_MAX_CONNECTIONS=1; the " - "exchange will serialize behind compute and give no overlap speedup. Unset it " - "or set it != 1 to get overlap, or pass overlap=False to skip the prefetch thread." - ) self._device = torch.cuda.current_device() self._queue: "queue.Queue" = queue.Queue(maxsize=1) self._stop = threading.Event() diff --git a/tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py b/tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py index 7a1e0c48599..0368ebc9db8 100644 --- a/tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py +++ b/tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py @@ -998,6 +998,29 @@ def test_reordering_buffer_window_size_validation(): ) +def _buffer_kwargs(): + return dict(dp_rank=0, dp_size=2, n_groups=2, cost_of=lambda f: 0.0, dp_group_gloo=None, dp_group_nccl=None) + + +def test_reordering_buffer_overlap_refuses_single_hardware_queue(monkeypatch): + # CUDA_DEVICE_MAX_CONNECTIONS=1 serializes the worker's all-to-all behind the DDP all-reduce in + # rank-dependent order (reproduced deadlock); overlap=True must be refused before any thread starts. + monkeypatch.setenv("CUDA_DEVICE_MAX_CONNECTIONS", "1") + with pytest.raises(ValueError, match="CUDA_DEVICE_MAX_CONNECTIONS=1"): + ReorderingBuffer(iter([]), overlap=True, **_buffer_kwargs()) + + +def test_reordering_buffer_no_overlap_allowed_under_single_hardware_queue(monkeypatch): + monkeypatch.setenv("CUDA_DEVICE_MAX_CONNECTIONS", "1") + buf = ReorderingBuffer(iter([]), overlap=False, **_buffer_kwargs()) + assert buf._thread is None + + +def test_reordering_buffer_overlap_allowed_without_single_hardware_queue(monkeypatch): + monkeypatch.delenv("CUDA_DEVICE_MAX_CONNECTIONS", raising=False) + ReorderingBuffer(iter([]), overlap=True, **_buffer_kwargs()) # no raise (thread only starts on CUDA) + + # --------------------------------------------------------------------------- # Joint per-sample cost # --------------------------------------------------------------------------- diff --git a/tests/unit_tests/examples/test_qwen35_vl_mimo_finetune.py b/tests/unit_tests/examples/test_qwen35_vl_mimo_finetune.py index 03e0a375c86..086689b1ef6 100644 --- a/tests/unit_tests/examples/test_qwen35_vl_mimo_finetune.py +++ b/tests/unit_tests/examples/test_qwen35_vl_mimo_finetune.py @@ -326,3 +326,39 @@ def rank(self): assert spec.labels is False and spec.loss_mask is False and spec.position_ids is False finally: sys.modules.pop(name, None) + + +def _reorder_args(**overrides): + args = SimpleNamespace( + micro_batch_size=4, + global_batch_size=8, + intra_microbatch_reorder=True, + no_overlap_intra_microbatch_reorder=False, + pad_to_seq_length=True, + ) + for key, value in overrides.items(): + if not hasattr(args, key): + raise ValueError(f"args has no field '{key}'") + setattr(args, key, value) + return args + + +@pytest.mark.parametrize("no_overlap", [False, True]) +def test_validate_rejects_reorder_overlap_under_single_hardware_queue(monkeypatch, no_overlap): + # Overlapped reorder deadlocks under CUDA_DEVICE_MAX_CONNECTIONS=1; validate must refuse it on + # every rank before model build. The non-overlapped path stays allowed. + name = "qwen35_vl_mimo_finetune_validate_under_test" + module = _load_example_module(name) + try: + monkeypatch.setenv("CUDA_DEVICE_MAX_CONNECTIONS", "1") + empty = SimpleNamespace(module_parallelisms={}) + args = _reorder_args(no_overlap_intra_microbatch_reorder=no_overlap) + if no_overlap: + assert module._validate_mimo_batch_sizes(empty, args) == [] + else: + with pytest.raises(ValueError, match="CUDA_DEVICE_MAX_CONNECTIONS=1"): + module._validate_mimo_batch_sizes(empty, args) + monkeypatch.delenv("CUDA_DEVICE_MAX_CONNECTIONS") + assert module._validate_mimo_batch_sizes(empty, _reorder_args()) == [] + finally: + sys.modules.pop(name, None) From aeab43dd822cb9d576c7b0cecced9d8edb849d3c Mon Sep 17 00:00:00 2001 From: "kayeon.song" Date: Sat, 3 Oct 2026 04:04:28 +0000 Subject: [PATCH 07/13] refactor(data): trim MegatronMIMO reorder docstrings and comments Rewrite the reorder_buffer docstrings in the Google style used by the neighbouring megatron_mimo modules (summary, Args, Returns, Raises) and drop the design narrative; keep inline comments only where they state a constraint the code depends on. Same trim for the reorder unit and functional tests. No code change (AST identical). Signed-off-by: kayeon.song Signed-off-by: Yoonsik Kim Signed-off-by: Chanwoo Park Co-authored-by: Yoonsik Kim Co-authored-by: Chanwoo Park <45866990+sailor1493@users.noreply.github.com> --- .../data/megatron_mimo/reorder_buffer.py | 722 ++++++------------ .../megatron_mimo/test_reorder_exchange.py | 38 +- .../data/megatron_mimo/test_reorder_buffer.py | 198 ++--- 3 files changed, 292 insertions(+), 666 deletions(-) diff --git a/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py b/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py index e25f031df66..f1d33717bcb 100644 --- a/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py +++ b/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py @@ -12,50 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Per-sample reorder buffer for MegatronMIMO intra-microbatch reordering (post-batch). - -Each MegatronMIMO data rank reads only its **disjoint** scalable-data-parallel shard of the global -micro-batch, and the per-microbatch cost imbalance between shards is corrected by a -**communication** step (cost all-gather + ragged sample all-to-all over the module DP -group) that can be overlapped with compute — instead of every rank reading the full -global micro-batch and slicing it locally. - -Contents, bottom-up: - -1. Pure, communication-free building blocks (all CPU-unit-testable without a process - group): - - - :func:`balanced_assignment` — the per-sample plan: for each global sample in a - micro-batch, the module-local DP rank that should own it and its canonical position. - Built on :func:`balanced_index_order` (same module), - with het-DP handled via canonical ``n_groups = canonical_grid_size(module dp_sizes)`` (their LCM) (so the vision - and language modules, which may have different DP sizes, derive an identical - assignment and the ``BridgeCommunicator`` keeps vision replica *r* paired with - language replica *r*). - - Ragged (de)serialization of a sample's tensors to an 8-byte-aligned ``uint8`` buffer - (:func:`serialize_sample` / :func:`deserialize_sample`), so any dtype/shape mix can - ride a single ``all_to_all_single`` with per-rank byte splits; the receiver - reconstructs from metadata alone. - - :func:`prepare_sample_exchange` — from this rank's local samples plus the all-gathered - plan and metadata, partition kept-vs-exchanged samples, pack the contiguous send buffer, - and build the ``send_splits``/``recv_splits`` + ordered recv schedule for one - ``all_to_all_single``. - -2. The MIMO data plane: :func:`split_microbatch` / :func:`merge_samples` convert a nested - micro-batch to and from ``B`` per-sample flat (dotted-key) dicts, and :func:`sample_cost` - is the (collation-independent) per-sample cost — ``encoder_cost_weight·patches + language_cost_weight·real_tokens`` — - that drives the balanced assignment. - -3. The distributed glue: :func:`exchange_window` ties the above into one cost-balancing round - over a window of micro-batches on a module DP group, :class:`ReorderingBuffer` wraps a data - iterator to run that round per window (optionally on a background prefetch thread with a dedicated CUDA - stream + side NCCL PG, so step ``t+1``'s exchange overlaps step ``t``'s compute), and - :func:`build_module_dp_process_groups` creates the Gloo (metadata) and side NCCL - (exchange) process groups the buffer needs. - -The unit of exchange is a **single sample** (finer than a whole-microbatch swap), and the -assignment reuses the existing MIMO cost-balanced ordering rather than a separate -group-then-LPT pass. +"""Per-sample reorder buffer for MegatronMIMO intra-microbatch reordering. + +Each data rank reads only its disjoint scalable-data-parallel shard of the global micro-batch. +The per-microbatch cost imbalance between shards is corrected by a cost all-gather plus one +ragged per-sample ``all_to_all_single`` over the module DP group, optionally overlapped with +compute. Paired modules derive the same assignment from the same costs and canonical grid, so +the ``BridgeCommunicator`` positional routing stays valid without cross-module communication. """ import logging @@ -76,20 +39,17 @@ logger = logging.getLogger(__name__) -# All sample tensors are serialized into uint8 byte buffers padded to this alignment, so -# the receiver can view-cast any dtype back from the received bytes without alignment -# errors and compute exact offsets from metadata alone. +# Serialized fields are padded to this alignment so the receiver can view-cast any dtype back +# from the byte buffer and compute exact offsets from metadata alone. _SERIALIZED_TENSOR_BYTE_ALIGNMENT = 8 def _alignment_padding_bytes(nbytes: int) -> int: - """Zero-padding bytes to round ``nbytes`` up to the next :data:`_SERIALIZED_TENSOR_BYTE_ALIGNMENT` multiple.""" + """Return the zero-padding bytes that round ``nbytes`` up to the serialization alignment.""" return (-nbytes) % _SERIALIZED_TENSOR_BYTE_ALIGNMENT -# String -> (torch.dtype, byte width), for reconstructing tensors from metadata shared via -# all_gather_object. The itemsize is precomputed once so the (de)serialize byte-size loops do not -# allocate a throwaway tensor per field just to read ``element_size()``. +# Item sizes are precomputed so the byte-size loops do not allocate a throwaway tensor per field. _DTYPE_INFO = { name: (dt, torch.empty((), dtype=dt).element_size()) for name, dt in { @@ -105,7 +65,7 @@ def _alignment_padding_bytes(nbytes: int) -> int: def _dtype_info(dtype_name: str) -> Tuple[torch.dtype, int]: - """Return ``(dtype, byte width)`` for a dtype encoded in sample metadata, rejecting unknown values.""" + """Return ``(dtype, byte width)`` for a dtype name from sample metadata.""" try: return _DTYPE_INFO[dtype_name] except KeyError as exc: @@ -120,24 +80,15 @@ def _dtype_info(dtype_name: str) -> Tuple[torch.dtype, int]: def balanced_index_order(costs: List[float], n_groups: int) -> List[int]: """Permute indices so contiguous equal-size groups have balanced total cost. - For intra-microbatch reordering when the downstream distributor slices the global - batch into ``n_groups`` **contiguous, equal-size** shards (e.g. - :func:`megatron.bridge.data.megatron_mimo.dp_utils.slice_batch_for_megatron_mimo`). - Returns a permutation of ``range(len(costs))`` laid out group-by-group, so that - taking contiguous block ``r`` (size ``len(costs)//n_groups``) yields a - cost-balanced shard for DP rank ``r``. - - Equal-cardinality LPT: items in decreasing cost are each placed into the - lowest-cost group that still has capacity (``len(costs)//n_groups``). Deterministic - (stable tie-break by group index) so paired modules reorder identically. + Equal-cardinality LPT with a stable tie-break by group index, so paired modules that + feed the same costs produce the same permutation. Args: costs: Per-item cost. - n_groups: Number of contiguous shards (e.g. DP size). Must divide ``len(costs)``. + n_groups: Number of contiguous shards; must divide ``len(costs)``. Returns: - A permutation of ``range(len(costs))``; the concatenation of the ``n_groups`` - balanced groups in order. + A permutation of ``range(len(costs))`` laid out group-by-group. Raises: ValueError: If ``n_groups <= 0`` or does not divide ``len(costs)``. @@ -151,7 +102,6 @@ def balanced_index_order(costs: List[float], n_groups: int) -> List[int]: groups: List[List[int]] = [[] for _ in range(n_groups)] group_cost = [0.0] * n_groups for i in sorted(range(n), key=lambda j: (-costs[j], j)): - # lowest-cost group that still has capacity (tie-break by group index) best = min( (g for g in range(n_groups) if len(groups[g]) < cap), key=lambda g: (group_cost[g], g), @@ -162,34 +112,21 @@ def balanced_index_order(costs: List[float], n_groups: int) -> List[int]: def balanced_assignment(costs: List[float], n_groups: int, dp_size: int) -> List[Tuple[int, int]]: - """Per-sample exchange plan for one micro-batch. - - Runs the canonical cost-balanced ordering over ``n_groups`` equal-cardinality groups, - then maps each global sample to the **module-local DP rank** that should own it and to - its **canonical position** in the globally balanced order. + """Map each sample of one micro-batch to its owner rank and canonical position. - The owner mapping uses canonical groups (``n_groups = canonical_grid_size(module dp_sizes)`` (their LCM)): a module - of size ``dp_size`` covers ``n_groups // dp_size`` contiguous canonical groups, so - ``owner = canonical_group // (n_groups // dp_size)``. Because the vision and language - modules feed the same ``costs`` and ``n_groups``, they produce an identical plan and - paired replicas end up with the same samples in the same (canonical) order — required - by the ``BridgeCommunicator`` batch-dim fan-out. + ``n_groups`` is the canonical grid (LCM of the module DP sizes); a module of size ``dp_size`` + owns ``n_groups // dp_size`` consecutive groups, so modules with different DP sizes agree. Args: - costs: Per-sample cost, indexed by the sample's global position in the micro-batch - (the patch-only cost from :func:`sample_cost`). - n_groups: Canonical group count (LCM of the module DP sizes). Must divide ``len(costs)``. - dp_size: This module's DP size. Must divide ``n_groups``. + costs: Per-sample cost indexed by global position in the micro-batch. + n_groups: Canonical group count; must divide ``len(costs)``. + dp_size: This module's DP size; must divide ``n_groups``. Returns: - For each global sample index ``g``: ``(owner_rank, canonical_pos)`` where - ``owner_rank`` is in ``[0, dp_size)`` and ``canonical_pos`` is in ``[0, len(costs))``. - Sorting a rank's owned samples by ``canonical_pos`` yields the order paired modules - agree on. + ``(owner_rank, canonical_pos)`` per global sample index. Raises: - ValueError: If ``n_groups`` does not divide ``len(costs)`` or ``dp_size`` does not - divide ``n_groups``. + ValueError: If ``n_groups`` does not divide ``len(costs)`` or ``dp_size`` does not divide it. """ n = len(costs) if n_groups <= 0 or n % n_groups != 0: @@ -197,7 +134,7 @@ def balanced_assignment(costs: List[float], n_groups: int, dp_size: int) -> List if dp_size <= 0 or n_groups % dp_size != 0: raise ValueError(f"dp_size ({dp_size}) must be positive and divide n_groups ({n_groups}).") - perm = balanced_index_order(costs, n_groups) # global indices laid out group-by-group + perm = balanced_index_order(costs, n_groups) cap = n // n_groups groups_per_rank = n_groups // dp_size @@ -210,73 +147,49 @@ def balanced_assignment(costs: List[float], n_groups: int, dp_size: int) -> List def intra_route(assignment: List[Tuple[int, int]], *, src_slot: int = 0) -> List[Tuple[int, int, int]]: - """Lift a single-slot 2D balanced assignment to the 3D window route. + """Lift a single-slot assignment to the ``(dst_slot, owner_rank, canonical_pos)`` window route. - The window transport routes every sample by **three** coordinates - ``route[g] = (dst_slot, owner_rank, canonical_pos)`` (the destination micro-batch slot, the - module-local DP rank that owns it, and its canonical intra-slot position). The **intra** - balancer never moves a sample across micro-batch slots, so ``dst_slot == src_slot`` for every - sample; a future inter (cross-slot) balancer fills ``dst_slot`` freely behind this same seam. + Intra routing never moves a sample across slots, so ``dst_slot == src_slot`` for every sample. Args: - assignment: Per-slot output of :func:`balanced_assignment` — ``(owner_rank, canonical_pos)`` - per global sample index within the slot. - src_slot: The slot these samples were read from (and, for intra, stay in). Defaults to ``0`` - (the single-slot ``window_size == 1`` case). + assignment: Output of :func:`balanced_assignment` for one slot. + src_slot: Slot the samples were read from. Returns: - For each global sample index ``g``: ``(dst_slot, owner_rank, canonical_pos)`` with - ``dst_slot == src_slot``. + ``(dst_slot, owner_rank, canonical_pos)`` per global sample index. """ return [(src_slot, owner_rank, canonical_pos) for owner_rank, canonical_pos in assignment] def build_window_route(costs_per_slot: List[List[float]], n_groups: int, dp_size: int) -> List[Tuple[int, int, int]]: - """Concatenate the per-slot intra assignments of a window into one 3D route. + """Build the slot-major window route by balancing each slot independently. - The window holds ``W = len(costs_per_slot)`` micro-batch slots, each with ``B_global`` samples - contiguous-sharded across ``dp_size`` ranks. The **intra** balancer runs **independently per - slot** (each slot cost-balanced across ranks); slots never mix. The window global index of slot - ``s``'s slot-local sample ``j`` is ``g = s · B_global + j`` (slots laid out back-to-back), so - concatenating each slot's :func:`intra_route` (stamped with ``src_slot = s``) in slot order - yields ``route[g]`` directly. ``dst_slot == src_slot == s`` for every sample (intra invariant). - - A future inter (cross-slot) balancer replaces this builder with one that fills ``dst_slot`` - across slots; the transport (:func:`exchange_by_route`) is unchanged. + The window global index of slot ``s``, slot-local position ``j`` is ``g = s * B_global + j``. Args: - costs_per_slot: ``costs_per_slot[s]`` = slot ``s``'s per-sample costs indexed by slot-local - global position ``j`` (length ``B_global``, identical on paired modules). - n_groups: Canonical group count (LCM of the module DP sizes); must divide each ``B_global``. + costs_per_slot: Per-slot per-sample costs indexed by slot-local position. + n_groups: Canonical group count; must divide each slot's sample count. dp_size: This module's DP size. Returns: - ``route`` of length ``W · B_global`` — ``(dst_slot, owner_rank, canonical_pos)`` per window - global index, slot-major. + ``(dst_slot, owner_rank, canonical_pos)`` per window global index. """ route: List[Tuple[int, int, int]] = [] for slot, costs in enumerate(costs_per_slot): route.extend(intra_route(balanced_assignment(costs, n_groups, dp_size), src_slot=slot)) - # Safety net: the intra balancer must never move a sample across slots. Cheap, always-on. assert_intra_no_cross_slot(route, len(costs_per_slot[0]) if costs_per_slot else 0) return route def assert_intra_no_cross_slot(route: List[Tuple[int, int, int]], b_global: int) -> None: - """Guard against cross-slot leaks: every sample's ``dst_slot`` equals its source slot. - - For the slot-major window layout (``g = src_slot · b_global + j``) the source slot of global - index ``g`` is ``g // b_global``. The **intra** balancer keeps every sample in its own slot, so - a mismatch means slot plumbing (route build / binning) corrupted the slot dimension. A future - inter (cross-slot) balancer would route - across slots by design and must **not** run this guard. + """Raise if any sample's ``dst_slot`` differs from its source slot ``g // b_global``. Args: - route: The window route (``(dst_slot, owner_rank, canonical_pos)`` per global index). - b_global: Samples per micro-batch slot (``local · dp_size``); ``0`` skips the check. + route: ``(dst_slot, owner_rank, canonical_pos)`` per window global index. + b_global: Samples per micro-batch slot; ``0`` skips the check. Raises: - RuntimeError: If any sample's ``dst_slot`` differs from its source slot. + RuntimeError: If a sample's ``dst_slot`` differs from its source slot. """ if b_global <= 0: return @@ -292,21 +205,15 @@ def assert_intra_no_cross_slot(route: List[Tuple[int, int, int]], b_global: int) def window_cost_spread( costs_per_slot: List[List[float]], route: List[Tuple[int, int, int]], dp_size: int ) -> List[Dict[str, float]]: - """Per-slot pre/post-balance per-rank cost spread + remote-sample count (the balance probe). - - For each slot, totals the per-rank cost **before** balancing (the natural contiguous shard: rank - ``r`` holds slot-local positions ``[r·local, (r+1)·local)``) and **after** balancing (each sample - credited to its routed ``owner_rank``), plus how many samples change rank. Pure and cheap — used - only to log how much the reorder tightens the per-rank load (``after`` max/min should narrow - toward 1.0). + """Compute per-slot per-rank cost max/min before and after balancing and the remote-sample count. Args: - costs_per_slot: ``costs_per_slot[s]`` = slot ``s``'s per-sample costs by slot-local position. - route: The window route (:func:`build_window_route`), indexed by ``g = s·B_global + j``. + costs_per_slot: Per-slot per-sample costs indexed by slot-local position. + route: Window route from :func:`build_window_route`. dp_size: This module's DP size. Returns: - One dict per slot: ``{before_max, before_min, after_max, after_min, remote}``. + One ``{before_max, before_min, after_max, after_min, remote}`` dict per slot. """ spreads: List[Dict[str, float]] = [] for slot, costs in enumerate(costs_per_slot): @@ -340,30 +247,34 @@ def window_cost_spread( def sample_keys(meta: Dict[str, Any]) -> List[str]: - """Sorted non-None tensor keys for a sample, computed identically on both sides.""" + """Return the sorted tensor keys of a sample's metadata, identical on sender and receiver. + + Args: + meta: Sample metadata from :func:`tensor_metadata`. + + Returns: + Sorted keys whose metadata entry carries a ``dtype``. + """ return sorted(k for k, v in meta.items() if isinstance(v, dict) and "dtype" in v) def tensor_metadata(flat: Dict[str, Any]) -> Dict[str, Any]: - """Per-field metadata for a flattened sample, keyed identically on both sides. + """Build per-field ``{shape, dtype}`` / ``{non_tensor: value}`` / ``None`` metadata for one sample. - Each key maps to ``{shape, dtype}`` for a tensor, ``{non_tensor: value}`` for a non-None - scalar/global field, or ``None``. Shared across ranks via ``all_gather_object`` so a - receiver can pre-allocate buffers and reconstruct the sample without seeing the data. + Args: + flat: One flattened sample dict. - A tensor whose dtype is not in :data:`_DTYPE_INFO` is rejected **here** — at metadata creation, - before the all-gather — via the same :func:`_dtype_info` check the deserializer uses, - so an unsupported dtype fails loudly and uniformly on the producing rank rather than only - surfacing later as a byte-size/deserialize ``ValueError`` on the receiver. + Returns: + Metadata keyed like ``flat``, sufficient to size and decode the sample without its data. Raises: - ValueError: If a tensor field has a dtype not supported by the ragged (de)serializer. + ValueError: If a tensor field has a dtype the serializer does not support. """ meta: Dict[str, Any] = {} for key, t in flat.items(): if isinstance(t, torch.Tensor): dtype_name = str(t.dtype) - _dtype_info(dtype_name) # reject unsupported dtypes early (same check as decode) + _dtype_info(dtype_name) # fail on the producing rank, before the metadata all-gather meta[key] = {"shape": list(t.shape), "dtype": dtype_name} elif t is not None: meta[key] = {"non_tensor": t} @@ -373,15 +284,14 @@ def tensor_metadata(flat: Dict[str, Any]) -> Dict[str, Any]: def serialize_sample(flat: Dict[str, Any], keys: List[str]) -> torch.Tensor: - """Serialize one sample's tensors into a 1D uint8 buffer (8-byte aligned, sorted keys). + """Serialize one sample's tensors into a 1D aligned ``uint8`` buffer in ``keys`` order. Args: - flat: Flattened tensor dict; ``None`` values are skipped. - keys: Sorted non-None tensor keys (from :func:`sample_keys`), same on both sides. + flat: Flattened sample dict; ``None`` values are skipped. + keys: Sorted tensor keys from :func:`sample_keys`. Returns: - 1D uint8 CPU tensor with each tensor's raw bytes concatenated in ``keys`` order, - each padded to :data:`_SERIALIZED_TENSOR_BYTE_ALIGNMENT`. + 1D ``uint8`` tensor with each field padded to the serialization alignment. """ parts = [] for key in keys: @@ -397,7 +307,15 @@ def serialize_sample(flat: Dict[str, Any], keys: List[str]) -> torch.Tensor: def sample_byte_size(meta: Dict[str, Any], keys: List[str]) -> int: - """Byte size of a serialized sample from metadata alone (matches :func:`serialize_sample`).""" + """Return the serialized byte size of a sample from its metadata alone. + + Args: + meta: Sample metadata from :func:`tensor_metadata`. + keys: Sorted tensor keys from :func:`sample_keys`. + + Returns: + Byte count matching :func:`serialize_sample` for the same sample. + """ total = 0 for key in keys: info = meta.get(key) @@ -413,19 +331,16 @@ def sample_byte_size(meta: Dict[str, Any], keys: List[str]) -> int: def deserialize_sample( buf: torch.Tensor, offset: int, meta: Dict[str, Any], keys: List[str] ) -> Tuple[Dict[str, Any], int]: - """Inverse of :func:`serialize_sample`: reconstruct tensors from a uint8 buffer. + """Reconstruct one sample from a ``uint8`` buffer (inverse of :func:`serialize_sample`). Args: - buf: Received 1D uint8 buffer. - offset: Starting byte offset of this sample in ``buf``. - meta: This sample's field metadata (from :func:`tensor_metadata`: - ``{key: {shape, dtype} | {non_tensor: value} | None}``). - keys: Sorted non-None tensor keys (from :func:`sample_keys`). + buf: Received 1D ``uint8`` buffer. + offset: Byte offset of this sample in ``buf``. + meta: This sample's metadata from :func:`tensor_metadata`. + keys: Sorted tensor keys from :func:`sample_keys`. Returns: - ``(flat, new_offset)`` — the reconstructed tensor dict (tensors cloned off ``buf``, - with ``None`` and non-tensor keys restored from ``meta``) and the byte offset just - past this sample. + ``(flat, new_offset)``: the rebuilt sample and the byte offset just past it in ``buf``. """ flat: Dict[str, Any] = {} cursor = offset @@ -460,57 +375,34 @@ def prepare_sample_exchange( *, window_size: int = 1, ) -> Dict[str, Any]: - """Prepare everything one per-sample ``all_to_all_single`` needs, in three parts. + """Build the kept-local partition, packed send buffer, and splits for one ``all_to_all_single``. - Pure and communication-free: all cross-rank information (which rank currently holds which - global sample, and each sample's tensor metadata) is passed in, having been shared by a - prior ``all_gather_object``. This single pass produces three distinct things: - - 1. **Partition** — split this rank's samples into those it already owns (kept locally) and - those it must send/receive (``local_samples`` + the send/recv schedules). - 2. **Pack** — serialize the outgoing samples into one contiguous ``send_buf`` ordered by - destination rank. - 3. **Plan** — the per-rank byte ``send_splits``/``recv_splits`` and the ordered - ``recv_schedule`` that drive (and decode) the ``all_to_all_single``. - - Each sample carries its destination micro-batch slot (``dst_slot``) end-to-end so the receiver - can bin it into the right slot bucket on reassembly (:func:`reassemble_window`); this is the - only genuinely new plumbing vs. the single-slot path. Sender - and receiver order the moves identically by ``(dst_slot, canonical_pos, src_local_idx)`` so the - packed bytes line up. + Sender and receiver sort moves by ``(dst_slot, canonical_pos, src_local_idx)``, so the packed + bytes and ``recv_schedule`` line up without further communication. Args: - local_flats: This rank's flattened samples (tensor dicts), in local order. - local_global_indices: Global index of each entry in ``local_flats`` (same length). - route: Output of :func:`intra_route` (or a future inter balancer) — - ``(dst_slot, owner_rank, canonical_pos)`` per global sample index. - all_global_indices: ``all_global_indices[r]`` = the global indices currently held by - rank ``r`` (the all-gathered ``local_global_indices``). - all_tensor_meta: ``all_tensor_meta[r][i]`` = tensor metadata of rank ``r``'s ``i``-th - local sample (the all-gathered :func:`tensor_metadata`). + local_flats: This rank's flattened samples in local order. + local_global_indices: Global index of each entry in ``local_flats``. + route: ``(dst_slot, owner_rank, canonical_pos)`` per global sample index. + all_global_indices: Global indices held by each rank, all-gathered. + all_tensor_meta: :func:`tensor_metadata` of each rank's samples, all-gathered. dp_rank: This rank's module-local DP rank. dp_size: This module's DP size. - window_size: Number of micro-batch slots ``W`` in the window; ``dst_slot`` must be in - ``[0, W)``. Defaults to ``1`` (the single-slot, byte-for-byte-with-intra case). + window_size: Number of micro-batch slots; every ``dst_slot`` must be in ``[0, window_size)``. Returns: - Dict with: - - (partition) ``local_samples``: ``[(dst_slot, canonical_pos, local_idx)]`` kept on this - rank. - - (pack) ``send_buf``: 1D uint8 tensor on the same device as the sample tensors (GPU - when CUDA is available), outgoing samples serialized contiguously by dst. - - (plan) ``send_splits`` / ``recv_splits``: per-rank byte counts for - ``all_to_all_single`` (``0`` for self). - - (plan) ``recv_schedule``: ``{src_rank: [(dst_slot, canonical_pos, src_local_idx)]}`` in - recv order, so the receiver deserializes ``recv_buf`` in the same order it was packed. + Dict with ``local_samples`` (``(dst_slot, canonical_pos, local_idx)`` kept here), ``send_buf``, + per-rank byte ``send_splits`` / ``recv_splits``, and ``recv_schedule`` + (``{src_rank: [(dst_slot, canonical_pos, src_local_idx)]}`` in receive order). + + Raises: + ValueError: If a ``dst_slot`` is outside ``[0, window_size)``. """ - # Map global index -> (current owner rank, current local index) from the all-gather. location: Dict[int, Tuple[int, int]] = {} for r, gidxs in enumerate(all_global_indices): for i, g in enumerate(gidxs): location[g] = (r, i) - # Samples this rank must end up with: every global index whose owner is dp_rank. local_samples: List[Tuple[int, int, int]] = [] send_schedule: Dict[int, List[Tuple[int, int, int]]] = {r: [] for r in range(dp_size)} recv_schedule: Dict[int, List[Tuple[int, int, int]]] = {r: [] for r in range(dp_size)} @@ -522,7 +414,6 @@ def prepare_sample_exchange( raise ValueError(f"dst_slot ({dst_slot}) for global sample {global_idx} out of range [0, {window_size}).") src_rank, src_local_idx = location[global_idx] if owner_rank == src_rank: - # Stays put; only this rank records it (when it is the owner/holder). if owner_rank == dp_rank: local_samples.append((dst_slot, canonical_pos, local_pos[global_idx])) continue @@ -531,13 +422,12 @@ def prepare_sample_exchange( if owner_rank == dp_rank: recv_schedule[src_rank].append((dst_slot, canonical_pos, src_local_idx)) - # Deterministic order (by slot then canonical position) so sender and receiver agree. + # Sender and receiver must pack and decode in the same order. for r in range(dp_size): send_schedule[r].sort() recv_schedule[r].sort() - # Device of the (possibly GPU-resident) sample tensors, so the send buffer and any empty - # placeholders live on the same device for torch.cat / all_to_all_single. + # Empty placeholders must share the sample tensors' device for torch.cat / all_to_all_single. dev = torch.device("cpu") for f in local_flats: t = next((v for v in f.values() if isinstance(v, torch.Tensor)), None) @@ -545,7 +435,6 @@ def prepare_sample_exchange( dev = t.device break - # Build the contiguous send buffer ordered by destination rank. send_parts: List[torch.Tensor] = [] send_splits: List[int] = [] for dst in range(dp_size): @@ -560,7 +449,6 @@ def prepare_sample_exchange( send_parts.append(chunk) send_splits.append(int(chunk.numel())) - # Recv byte counts from metadata alone. recv_splits: List[int] = [] for src in range(dp_size): if src == dp_rank or not recv_schedule[src]: @@ -583,37 +471,24 @@ def prepare_sample_exchange( # --------------------------------------------------------------------------- -# MIMO data plane: micro-batch <-> per-sample flat dicts (flatten_fn / rewrap_fn) +# MIMO data plane: micro-batch <-> per-sample flat dicts # --------------------------------------------------------------------------- -# -# The exchange unit is a single sample. A MegatronMIMO micro-batch is nested: -# -# {input_ids: [B,S], position_ids: [3,B,S] (MRoPE), labels/loss_mask: [B,S] | None, -# attention_mask: None, modality_inputs: {modality: {encoder: {hidden_states: [ΣP,d], -# grid_thw: [n_img,3]}}}} -# -# ``split_microbatch`` turns it into B per-sample **flat** dicts (nesting flattened to -# dotted keys) so the serializer/all-to-all can move individual samples; after the -# exchange ``merge_samples`` rebuilds the (rebalanced) micro-batch. Sample slicing/vision -# attribution uses :func:`_apply_sample_dispatch` (LM by batch dim, vision by image via -# ``image_counts``). def _flatten_nested(d: Dict[str, Any], prefix: str = "") -> Dict[str, Any]: - """Flatten a (possibly nested) batch dict to dotted keys; tensors and ``None`` preserved.""" + """Flatten a nested batch dict to dotted keys, preserving tensors and ``None``.""" flat: Dict[str, Any] = {} for key, val in d.items(): path = f"{prefix}{key}" if isinstance(val, dict): flat.update(_flatten_nested(val, prefix=f"{path}.")) else: - # tensors and None pass through; the buffer's batch has no per-sample list fields. flat[path] = val return flat def _unflatten_nested(flat: Dict[str, Any]) -> Dict[str, Any]: - """Inverse of :func:`_flatten_nested`: rebuild the nested dict from dotted keys.""" + """Rebuild the nested dict from dotted keys (inverse of :func:`_flatten_nested`).""" out: Dict[str, Any] = {} for path, val in flat.items(): parts = path.split(".") @@ -625,11 +500,7 @@ def _unflatten_nested(flat: Dict[str, Any]) -> Dict[str, Any]: def _cu_img_from_counts(image_counts: "Any") -> List[int]: - """Cumulative per-sample image-offset table from per-sample image counts. - - ``image_counts`` is a length-``B`` int tensor/sequence; returns ``[0, c0, c0+c1, ...]`` (length - ``B + 1``) so sample ``s`` owns ``grid_thw`` rows ``[cu_img[s], cu_img[s + 1])``. - """ + """Return the cumulative image-offset table ``[0, c0, c0+c1, ...]`` from per-sample image counts.""" counts = image_counts.tolist() if isinstance(image_counts, torch.Tensor) else list(image_counts) cu = [0] for c in counts: @@ -640,12 +511,17 @@ def _cu_img_from_counts(image_counts: "Any") -> List[int]: def empty_like_vision( ref_hidden_states: torch.Tensor, ref_grid_thw: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: - """Empty vision tensors matching a reference sample's dtype/device. + """Return empty ``[0, d]`` hidden states and ``[0, 3]`` grid matching a reference sample's dtype and device. + + Text-only samples contribute these placeholders so every sample carries the same vision + keys and the merge concat and byte-split metadata stay symmetric across ranks. - Returns ``([0, d] hidden_states, [0, 3] grid_thw)`` — the placeholder a **text-only** sample - (0 images) contributes so the per-sample reorder/merge keeps uniform vision keys across samples - and ranks (the merge concat and the cross-rank byte-split metadata both stay symmetric). Used by - :func:`merge_samples`. + Args: + ref_hidden_states: A ``[patches, d]`` hidden-states tensor to match. + ref_grid_thw: A ``[n_images, 3]`` grid tensor to match. + + Returns: + ``(hidden_states, grid_thw)`` with zero rows. """ empty_hidden_states = ref_hidden_states.new_empty((0, ref_hidden_states.shape[1])) empty_grid_thw = ref_grid_thw.new_empty((0, ref_grid_thw.shape[1])) @@ -653,10 +529,13 @@ def empty_like_vision( def patches_per_image(grid_thw: torch.Tensor) -> torch.Tensor: - """Per-image vision patch count ``prod(t, h, w)`` for a ``[n_images, 3]`` grid. + """Return the per-image patch count ``prod(t, h, w)`` for a ``[n_images, 3]`` grid. + + Args: + grid_thw: Per-image ``(t, h, w)`` grid. - Centralizes the Qwen-VL "patches per image = ``prod(grid_thw[i])``" convention shared by - the vision reorder and cost paths so a future grid layout change is a single edit. + Returns: + ``[n_images]`` patch counts. """ return torch.prod(grid_thw, dim=1) @@ -666,28 +545,12 @@ def _reorder_vision_by_images( grid_thw: torch.Tensor, image_perm: list[int], ) -> Tuple[torch.Tensor, torch.Tensor]: - """Permute vision ``hidden_states`` (patch-dim) and ``grid_thw`` (image-dim) by image. - - The patch block of image ``i`` (rows ``cu[i]:cu[i+1]`` where ``cu`` is the cumulative - patch count from ``prod(grid_thw[i])``) is gathered in ``image_perm`` order so the two - stay aligned — the patch-dim equivalent of ``grid_thw[image_perm]``. Reorders whole images - (keeping each image's patch block intact) rather than slicing across image boundaries. - - Args: - hidden_states: Vision patch features ``[Σ patches, dim]``. - grid_thw: Per-image grid ``[n_images, 3]``. - image_perm: Permutation of ``range(n_images)``. - - Returns: - ``(hidden_states_perm, grid_thw_perm)``. - """ + """Permute ``hidden_states`` (patch dim) and ``grid_thw`` (image dim) by whole image.""" if not image_perm: - # No images selected (e.g. a text-only sample): return empty ``[0, d]`` / ``[0, 3]`` views - # (avoids ``torch.cat([])`` and preserves dtype/device/requires_grad). + # Avoid torch.cat([]) for a text-only sample while keeping dtype/device/requires_grad. return hidden_states[:0], grid_thw[:0] - patches = patches_per_image(grid_thw) # [n_images] - # Single host copy of the cumulative offsets: materialize the cumsum once with one - # .tolist(), then index with plain Python ints — no per-image .item() D2H sync in the loop. + patches = patches_per_image(grid_thw) + # One host copy of the offsets instead of a per-image .item() sync inside the loop. cu = torch.cat([patches.new_zeros(1), patches.cumsum(0)]).tolist() blocks = [hidden_states[cu[i] : cu[i + 1]] for i in image_perm] perm_idx = torch.as_tensor(image_perm, dtype=torch.long, device=grid_thw.device) @@ -701,32 +564,7 @@ def _gather_vision_subdict( *, cu_img: "list[int] | None" = None, ) -> Dict[str, Any]: - """Reorder/gather a vision ``{hidden_states, grid_thw}`` sub-dict by sample. - - Supports a **variable number of images per sample** (0 = text-only, 1, or N). ``cu_img`` is the - cumulative per-sample image-offset table (length ``n_samples + 1``), so sample ``s`` owns - ``grid_thw`` rows ``[cu_img[s], cu_img[s + 1])``. The image permutation for ``sample_indices`` is - the concatenation of each selected sample's image range, gathered by - :func:`_reorder_vision_by_images` (which keeps ``hidden_states`` patch-dim and ``grid_thw`` - image-dim aligned). A text-only sample contributes an empty ``[0, d]`` / ``[0, 3]`` block. - - When ``cu_img`` is ``None`` the one-image-per-sample fallback (image ``i`` ↔ sample ``i``) - is used and requires ``n_images == n_samples``. - - Args: - value: A sub-dict carrying ``hidden_states`` and ``grid_thw``. - sample_indices: Per-sample permutation/selection (global sample ids). - n_samples: Global batch sample count ``B``. - cu_img: Cumulative per-sample image offsets (length ``B + 1``), or ``None`` for the - one-image-per-sample fallback. - - Returns: - A new sub-dict with permuted ``hidden_states`` / ``grid_thw`` (other keys carried). - - Raises: - ValueError: If ``cu_img`` is ``None`` and the batch is not one-image-per-sample, or if - ``cu_img`` does not sum to ``n_images`` (mis-sourced image counts). - """ + """Gather a vision ``{hidden_states, grid_thw}`` sub-dict by sample, supporting 0..N images per sample.""" hidden_states, grid_thw = value["hidden_states"], value["grid_thw"] n_images = grid_thw.shape[0] if cu_img is None: @@ -758,13 +596,7 @@ def _apply_sample_dispatch( n_samples: int, cu_img: "list[int] | None" = None, ) -> Dict[str, Any]: - """Shared per-sample gather/permute visitor for tensors, vision sub-dicts, and lists. - - Used by :func:`split_microbatch` (per-sample local-shard selection); ``n_samples`` is the - canonical global batch size ``B`` so list/vision per-sample detection is consistent. ``cu_img`` - (cumulative per-sample image offsets) is forwarded to :func:`_gather_vision_subdict` for the - variable-images-per-sample reorder; ``None`` selects the one-image-per-sample fallback. - """ + """Gather ``sample_indices`` from every tensor, vision sub-dict, and per-sample list in a batch.""" idx = torch.as_tensor(sample_indices, dtype=torch.long) out: Dict[str, Any] = {} for key, value in batch.items(): @@ -777,7 +609,7 @@ def _apply_sample_dispatch( else: out[key] = _apply_sample_dispatch(value, sample_indices, n_samples=n_samples, cu_img=cu_img) elif isinstance(value, list): - # Per-sample list -> permuted gather; global-metadata / empty list -> passthrough. + # A list of exactly n_samples entries is per-sample; any other length is global metadata. if len(value) == n_samples: out[key] = [value[i] for i in sample_indices] else: @@ -788,15 +620,14 @@ def _apply_sample_dispatch( def split_microbatch(batch: Dict[str, Any], *, cu_img: "List[int] | None" = None) -> List[Dict[str, Any]]: - """Split a micro-batch into ``B`` per-sample flat (dotted-key) dicts. + """Split a micro-batch into ``B`` per-sample flat (dotted-key) dicts with size-1 batch dims. - Each sample is gathered with :func:`_apply_sample_dispatch` (so its image travels with it) - and then flattened to dotted keys for serialization. Leading batch dims are kept (size 1), - so :func:`merge_samples` is a plain concat. + Args: + batch: Nested micro-batch. + cu_img: Cumulative per-sample image offsets (length ``B + 1``); ``None`` assumes one image per sample. - ``cu_img`` (cumulative per-sample image offsets, length ``B + 1``) lets the vision gather support - a variable number of images per sample (0 = text-only, 1, or N). When ``None``, the - one-image-per-sample fallback is used. + Returns: + ``B`` flat sample dicts in batch order. """ b = _batch_size(batch) return [_flatten_nested(_apply_sample_dispatch(batch, [i], n_samples=b, cu_img=cu_img)) for i in range(b)] @@ -805,9 +636,14 @@ def split_microbatch(batch: Dict[str, Any], *, cu_img: "List[int] | None" = None def merge_samples(flats: List[Dict[str, Any]]) -> Dict[str, Any]: """Rebuild a micro-batch from per-sample flat dicts (inverse of :func:`split_microbatch`). - LM tensors concat on their batch dim (:func:`dp_utils._batch_dim_for_tensor`); vision - ``hidden_states`` concat on the patch dim and ``grid_thw`` on the image dim; ``None`` and - scalar/global fields are taken from the first sample. + Args: + flats: Per-sample flat dicts in the desired batch order. + + Returns: + The nested micro-batch, tensors concatenated on their batch, patch, or image dim. + + Raises: + ValueError: If ``flats`` is empty. """ if not flats: raise ValueError("merge_samples requires at least one sample.") @@ -816,11 +652,8 @@ def merge_samples(flats: List[Dict[str, Any]]) -> Dict[str, Any]: def _merge_nested(samples: List[Dict[str, Any]]) -> Dict[str, Any]: - # Iterate the union of keys (not just ``samples[0]``) and classify a vision sub-dict by scanning - # *all* samples: a text-only sample (0 images) carries the sub-dict as empty/None, so keying off - # ``samples[0]`` alone would misclassify it as a plain nested dict and silently drop the other - # samples' vision. With the empty-tensor convention every sample presents the sub-dict as a - # tensor, so the empty substitution below is a defensive backstop. + # Walk the union of keys and classify a vision sub-dict from any sample: a text-only sample + # may carry it as None, and keying off samples[0] would drop the other samples' vision. keys: List[str] = [] seen = set() for s in samples: @@ -834,8 +667,6 @@ def _merge_nested(samples: List[Dict[str, Any]]) -> Dict[str, Any]: vals = [s.get(key) for s in samples] vis_ref = next((v for v in vals if _is_patch_packed_visual_dict(v)), None) if vis_ref is not None: - # Vision sub-dict: concat hidden_states (patch dim) + grid_thw (image dim), substituting - # an empty [0,d]/[0,3] for any text-only sample (None/absent) so the concat stays aligned. ref_hidden_states, ref_grid_thw = vis_ref["hidden_states"], vis_ref["grid_thw"] hidden_parts: List[torch.Tensor] = [] grid_parts: List[torch.Tensor] = [] @@ -859,12 +690,12 @@ def _merge_nested(samples: List[Dict[str, Any]]) -> Dict[str, Any]: elif isinstance(rep, torch.Tensor): out[key] = torch.cat(vals, dim=_batch_dim_for_tensor(key, rep)) else: - out[key] = rep # None / scalar / global metadata + out[key] = rep return out def _batch_to_cuda(batch: Dict[str, Any]) -> Dict[str, Any]: - """Recursively move a (possibly nested) batch dict's tensors to CUDA (non-blocking).""" + """Move a nested batch dict's tensors to CUDA (non-blocking).""" out: Dict[str, Any] = {} for key, val in batch.items(): if isinstance(val, torch.Tensor): @@ -877,7 +708,7 @@ def _batch_to_cuda(batch: Dict[str, Any]) -> Dict[str, Any]: def _batch_size(batch: Dict[str, Any]) -> int: - """Sample count of a micro-batch from a known LM tensor (handles MRoPE position_ids).""" + """Return the sample count of a micro-batch from one of its LM tensors.""" for key in ("input_ids", "labels", "loss_mask", "position_ids"): t = batch.get(key) if isinstance(t, torch.Tensor): @@ -898,43 +729,21 @@ def sample_cost( image_token_id: "int | None" = None, square_merge_size: int = 1, ) -> float: - """Joint per-sample cost: ``encoder_cost_weight·p + language_cost_weight·t``. - - Only the ratio of the two weights matters. Each term is linear in a per-sample count: - - - ``p`` is the vision patch count (mirrors vision-encoder FLOPs). - - ``t`` is the **real (non-pad) token count** of ``input_ids`` (mirrors LM FLOPs), counted - via :func:`real_token_lengths` (from the batch's ``attention_mask``), image placeholder - tokens included. - - Both terms must be **collation-independent** so vision and language derive an identical - assignment with no cross-module communication: the patch count is intrinsic to the image, - and the real token count is intrinsic to the sample (invariant to how each module's collator - pads). The *padded* sequence length is collation-dependent and is deliberately **not** used — - it would differ between the two shards and mispair the ``BridgeCommunicator`` fan-out. - - Patch-count source (must be **module-independent**): + """Return ``encoder_cost_weight * patches + language_cost_weight * real_tokens`` for one sample. - - When ``image_token_id`` is given, ``p`` is the image-placeholder token count in ``input_ids`` - scaled by ``square_merge_size`` (each placeholder token stands for ``spatial_merge_size²`` - patches in Qwen-VL). ``input_ids`` is present and identical on **every** module/PP stage when - reorder is active, so vision and language derive the same cost. This is the correct source: - the rank-aware metadata collate (#4442) nulls ``modality_inputs``/``grid_thw`` on language - shards, so reading ``grid_thw`` there yields ``p = 0`` and mispairs the vision↔language fan-out. - - Otherwise ``p = Σ prod(grid_thw)`` (the ``grid_thw`` fallback, used by mock/tests where no - ``image_token_id`` is wired and ``grid_thw`` is present on the sample). - - With the default ``language_cost_weight=0.0`` the cost is patch-only. + Both terms must be collation-independent so vision and language derive the same cost without + communicating: with ``image_token_id`` the patch count comes from ``input_ids``, which every + module holds (``grid_thw`` is nulled on language shards), and tokens come from ``attention_mask``. Args: flat: One per-sample flat (dotted-key) dict. encoder_cost_weight: Weight per vision patch. language_cost_weight: Weight per real token; ``0.0`` disables the term. - image_token_id: Placeholder token id whose count in ``input_ids`` is the module-independent - patch proxy. ``None`` falls back to the ``grid_thw`` patch sum. - square_merge_size: ``spatial_merge_size²`` — patches represented by one placeholder token. - Recovers the true patch count from the placeholder-token count (only used with - ``image_token_id``). + image_token_id: Placeholder token id counted in ``input_ids``; ``None`` sums ``grid_thw`` instead. + square_merge_size: Patches represented by one placeholder token (``spatial_merge_size**2``). + + Returns: + The weighted per-sample cost. """ patches = 0 if image_token_id is not None: @@ -975,29 +784,23 @@ def reassemble_window( window_size: int, local: int, ) -> List[Dict[str, Any]]: - """Rebuild this rank's ``W`` micro-batches from kept-local + received samples. - - Generalizes the single-micro-batch reassembler to ``W`` slot-buckets: kept-local samples and - deserialized received samples are binned by ``dst_slot``, sorted **within each slot** by - ``canonical_pos`` (the paired-module-identical order), and merged into one micro-batch per slot - via :func:`merge_samples`. ``deserialize_sample`` / ``merge_samples`` are reused verbatim — the - only new step is the per-slot binning. - - The slot invariant asserts are baked in (cheap, always on): they turn a silent slot mispair or - a dropped/duplicated sample into a loud failure at the exact rank. + """Rebuild this rank's ``W`` micro-batches from kept-local and received samples. Args: - local_flats: This rank's input samples (the kept-local ones are indexed by ``local_idx``). - plan: Output of :func:`prepare_sample_exchange` (carries ``local_samples`` / ``recv_schedule`` - with ``dst_slot``). - recv_buf: The ``all_to_all_single`` output buffer (samples this rank received). - all_tensor_meta: All-gathered per-rank tensor metadata (to decode ``recv_buf``). + local_flats: This rank's input samples; kept-local ones are indexed by ``local_idx``. + plan: Output of :func:`prepare_sample_exchange`. + recv_buf: The ``all_to_all_single`` output buffer. + all_tensor_meta: All-gathered per-rank tensor metadata used to decode ``recv_buf``. dp_size: This module's DP size. - window_size: Number of micro-batch slots ``W`` to rebuild. - local: Expected samples per (slot) bucket for this rank (``B/dp`` per slot). + window_size: Number of micro-batch slots ``W``. + local: Expected samples per slot on this rank. Returns: - ``W`` rebuilt micro-batches, slot ``s`` at index ``s``. + ``W`` micro-batches, slot ``s`` at index ``s``, each in canonical order. + + Raises: + RuntimeError: If the received byte count, a slot's sample count, or the canonical positions + do not match ``plan``. """ buckets: List[List[Tuple[int, Dict[str, Any]]]] = [[] for _ in range(window_size)] for dst_slot, canonical_pos, local_idx in plan["local_samples"]: @@ -1010,17 +813,15 @@ def reassemble_window( flat, offset = deserialize_sample(recv_buf, offset, meta, sample_keys(meta)) buckets[dst_slot].append((canonical_pos, flat)) - # Byte accounting: the deserialize cursor must consume exactly the received bytes. if offset != recv_buf.numel(): raise RuntimeError(f"MegatronMIMO reassembly consumed {offset} bytes but received {recv_buf.numel()}.") out: List[Dict[str, Any]] = [] for slot, bucket in enumerate(buckets): - # Slot shape: each slot bucket must hold exactly `local` samples. if len(bucket) != local: raise RuntimeError(f"MegatronMIMO slot {slot} reassembled {len(bucket)} samples, expected {local}.") bucket.sort(key=lambda x: x[0]) - # Order: canonical positions within a slot must be unique (no dropped/duplicated sample). + # Duplicate canonical positions mean a sample was dropped or duplicated in transit. positions = [p for p, _ in bucket] if len(set(positions)) != len(positions): raise RuntimeError(f"MegatronMIMO slot {slot} has duplicate canonical positions: {positions}.") @@ -1041,27 +842,22 @@ def exchange_by_route( dp_group_nccl: "Any", nccl_stream: "Any" = None, ) -> List[Dict[str, Any]]: - """Transport-only per-sample exchange: route -> one ``all_to_all_single`` -> ``W`` micro-batches. - - Given a 3D ``route`` and the all-gathered layout/metadata, build the - plan (:func:`prepare_sample_exchange`), run **one** ragged ``all_to_all_single``, and reassemble - ``W`` micro-batches per slot (:func:`reassemble_window`). Knows nothing about the balancer — an - intra route (``dst_slot == src_slot``, see :func:`intra_route`) and a future inter route flow - through unchanged. + """Run one ragged ``all_to_all_single`` for ``route`` and reassemble ``W`` micro-batches. Args: - local_flats: This rank's input samples (GPU-resident when CUDA is available). + local_flats: This rank's input samples. route: ``(dst_slot, owner_rank, canonical_pos)`` per global sample index. - all_global_indices: ``all_global_indices[r]`` = global indices currently held by rank ``r``. + all_global_indices: Global indices held by each rank, all-gathered. all_tensor_meta: All-gathered per-rank tensor metadata. - dp_rank, dp_size: Module-local DP coordinates. - window_size: Number of micro-batch slots ``W`` in the window. - local: Samples per (slot) bucket for this rank. - dp_group_nccl: NCCL PG for the sample all-to-all. - nccl_stream: Optional CUDA stream to run the all-to-all on (cross-step overlap). + dp_rank: This rank's module-local DP rank. + dp_size: This module's DP size. + window_size: Number of micro-batch slots ``W``. + local: Samples per slot on this rank. + dp_group_nccl: NCCL process group for the sample all-to-all. + nccl_stream: Optional CUDA stream to run the all-to-all on. Returns: - ``W`` rebuilt micro-batches for this rank, slot ``s`` at index ``s``. + ``W`` micro-batches for this rank, slot ``s`` at index ``s``. """ import torch.distributed as dist @@ -1076,7 +872,6 @@ def exchange_by_route( window_size=window_size, ) - # send_buf is already on GPU (built from GPU-resident sample tensors); no H2D copy. send_buf = plan["send_buf"] recv_buf = torch.empty(sum(plan["recv_splits"]), dtype=torch.uint8, device=send_buf.device) if nccl_stream is not None: @@ -1114,36 +909,28 @@ def exchange_window( image_count_of: "Any" = None, probe: bool = False, ) -> List[Dict[str, Any]]: - """Cost-balance a **window** of ``W`` micro-batches across a module DP group in one exchange. - - Collects this rank's ``W`` disjoint shards, all-gathers their per-sample costs + tensor metadata - over Gloo **once**, builds the slot-major window route (:func:`build_window_route` — the intra - balancer run independently per slot), and runs **one** ragged ``all_to_all_single`` over the - whole window, reassembling ``W`` balanced micro-batches per slot (:func:`reassemble_window`). - The single collective lands once per window (≈ once per optimizer step when ``W == GA``). + """Cost-balance a window of ``W`` micro-batches across a module DP group in one exchange. - The window global index of slot ``s``'s slot-local sample ``j`` is ``g = s · B_global + j``; this - rank holds, in ``local_flats`` order, its slot-0 shard then slot-1 shard … (``local = B/dp`` per - slot). Both modules feed identical ``n_groups`` and per-slot costs, so paired replicas end with - the same samples in canonical order in every slot. + Costs and metadata are all-gathered over Gloo once per window and the samples move in one + ``all_to_all_single``; paired modules must pass the same ``n_groups`` and ``cost_of``. Args: - batches: This rank's ``W`` disjoint micro-batch shards (each ``local = B/dp`` samples, same - ``local`` across slots). - dp_rank, dp_size: Module-local DP coordinates. + batches: This rank's ``W`` micro-batch shards, each with the same local sample count. + dp_rank: This rank's module-local DP rank. + dp_size: This module's DP size. n_groups: Canonical group count (LCM of the module DP sizes). - cost_of: ``callable(flat) -> float`` per-sample patch cost (see :func:`sample_cost`). - dp_group_gloo: Gloo PG for the metadata all-gather. - dp_group_nccl: NCCL PG for the sample all-to-all. - nccl_stream: Optional CUDA stream to run the all-to-all on (cross-step overlap). - image_count_of: Optional ``callable(batch) -> LongTensor[local]`` giving a shard's per-sample - image count, enabling a variable number of images per sample (0/1/N). When ``None``, the - one-image-per-sample fallback is used. - probe: When ``True``, log a one-line per-slot cost-spread digest (:func:`window_cost_spread`) - on ``dp_rank == 0`` (the balance probe). Off by default (the caller throttles it). + cost_of: ``callable(flat) -> float`` per-sample cost (see :func:`sample_cost`). + dp_group_gloo: Gloo process group for the metadata all-gather. + dp_group_nccl: NCCL process group for the sample all-to-all. + nccl_stream: Optional CUDA stream to run the all-to-all on. + image_count_of: Optional ``callable(batch) -> LongTensor`` per-sample image counts. + probe: When ``True``, log the per-slot cost spread on ``dp_rank == 0``. Returns: ``W`` rebalanced micro-batches for this rank, slot ``s`` at index ``s``. + + Raises: + ValueError: If ``batches`` is empty or the shards differ in local size. """ import torch.distributed as dist @@ -1151,9 +938,8 @@ def exchange_window( if window_size == 0: raise ValueError("exchange_window requires at least one micro-batch.") - # Cost and metadata are computed on the CPU batches before the optional CUDA move, avoiding - # per-sample .item() synchronizations on the training stream. Per-slot image-count offset tables - # are derived once and reused for both the CPU-meta and GPU-payload splits. + # Costs and metadata come from the CPU batches so the per-sample .item() calls do not + # synchronize the training stream. cu_imgs = [_cu_img_from_counts(image_count_of(b)) if image_count_of is not None else None for b in batches] cpu_flats_per_slot = [split_microbatch(b, cu_img=cu) for b, cu in zip(batches, cu_imgs)] local = len(cpu_flats_per_slot[0]) @@ -1162,15 +948,13 @@ def exchange_window( f"All micro-batches in a window must have the same local size; got {[len(s) for s in cpu_flats_per_slot]}." ) - # Window-local order is slot-major: this rank's slot-0 shard, then slot-1, … (matches the - # g = s·B_global + j layout build_window_route assumes). + # Slot-major local order matches the g = s * B_global + j layout build_window_route assumes. cpu_local_flats = [f for slot in cpu_flats_per_slot for f in slot] local_meta = [tensor_metadata(f) for f in cpu_local_flats] local_costs_per_slot = [[cost_of(f) for f in slot] for slot in cpu_flats_per_slot] - # Keep the exchange GPU-resident after cost calculation: move this rank's (small) shards to the - # device once, then split / serialize / all_to_all / deserialize / merge all run on GPU and the - # returned batches are already on GPU for the forward pass. + # Move the shards to the device once so serialize / all_to_all / merge stay GPU-resident and + # the returned batches need no further copy before the forward pass. if torch.cuda.is_available(): gpu_flats_per_slot = [split_microbatch(_batch_to_cuda(b), cu_img=cu) for b, cu in zip(batches, cu_imgs)] local_flats = [f for slot in gpu_flats_per_slot for f in slot] @@ -1182,21 +966,16 @@ def exchange_window( dist.all_gather_object(gathered, payload, group=dp_group_gloo) b_global = local * dp_size - # Window global index layout: g = slot·b_global + (src_rank·local + i); rank r's local_flats are - # ordered slot-major so its window samples are at all_global_indices[r] (same order). all_global_indices = [ [slot * b_global + r * local + i for slot in range(window_size) for i in range(local)] for r in range(dp_size) ] all_tensor_meta = [g["meta"] for g in gathered] - # Per-slot global costs indexed by slot-local position j = r·local + i (contiguous sharding). costs_per_slot: List[List[float]] = [[0.0] * b_global for _ in range(window_size)] for r, g in enumerate(gathered): for slot in range(window_size): for i, c in enumerate(g["costs_per_slot"][slot]): costs_per_slot[slot][r * local + i] = c - # Balancing is driven by the patch-only cost (see sample_cost): identical across modules, so both - # compute the same per-slot canonical order with no cross-module communication. route = build_window_route(costs_per_slot, n_groups, dp_size) if probe and dp_rank == 0: @@ -1228,15 +1007,13 @@ def reorder_window_local( cost_of: "Any", image_count_of: "Any" = None, ) -> List[Dict[str, Any]]: - """Reorder a window into canonical order on a ``dp_size == 1`` module (no exchange needed). - - Uses the same costs and ``n_groups`` as the paired dp>1 module, so both end in the same order. + """Permute a window into canonical order on a ``dp_size == 1`` module without any exchange. Args: batches: This rank's ``W`` whole micro-batches. - n_groups: Canonical group count (LCM of the module DP sizes). - cost_of: ``callable(flat) -> float`` per-sample cost. - image_count_of: Optional ``callable(batch) -> LongTensor[B]`` per-sample image counts. + n_groups: Canonical group count (LCM of the module DP sizes), same as the paired module. + cost_of: ``callable(flat) -> float`` per-sample cost, same as the paired module. + image_count_of: Optional ``callable(batch) -> LongTensor`` per-sample image counts. Returns: ``W`` micro-batches permuted into canonical order. @@ -1261,14 +1038,14 @@ def reorder_window_local( class _WorkerException: - """Container used to re-raise overlap-worker failures on the consumer thread.""" + """Carry an overlap-worker exception across the queue for re-raising on the consumer thread.""" def __init__(self, exc: BaseException) -> None: self.exc = exc def _record_stream_for_batch(value: Any, stream: torch.cuda.Stream) -> None: - """Recursively tie CUDA tensors in ``value`` to ``stream`` before cross-thread handoff.""" + """Tie every CUDA tensor in ``value`` to ``stream`` before a cross-thread handoff.""" if isinstance(value, torch.Tensor) and value.is_cuda: value.record_stream(stream) elif isinstance(value, dict): @@ -1282,32 +1059,27 @@ def _record_stream_for_batch(value: Any, stream: torch.cuda.Stream) -> None: class ReorderingBuffer: """Wrap a MegatronMIMO data iterator so each yielded micro-batch is cost-balanced. - Each item this rank consumes is its disjoint scalable-data-parallel shard rebalanced by - :func:`exchange_window` (per-sample cost all-gather + ragged all-to-all over the module - DP group); at ``dp_size == 1`` it is permuted locally instead (:func:`reorder_window_local`). - - **Window buffering** (``window_size = W``): collect ``W`` micro-batches, cost-balance them in one - exchange over the window (:func:`exchange_window`), and serve the ``W`` rebalanced micro-batches - one at a time from a cursor. ``W = 1`` (default) exchanges every micro-batch on its own. - ``W == GA`` (the gradient-accumulation count) is the desired setting — the single - window collective then lands once per optimizer step. - - **Cross-window prefetch overlap** (``overlap=True``): a background thread runs the *entire next - window's* exchange — the blocking Gloo metadata all-gather *and* the NCCL all-to-all (on a - dedicated CUDA stream) — while the main thread computes the current window's ``W`` micro-batches. - This is what turns the scalable-data-parallel read win into a net throughput win: one big - all-to-all is hidden behind ``W`` micro-batches of compute instead of sitting on the critical - path (for ``W = 1`` this is the one-step-ahead per-micro-batch overlap). The exchange's GPU work - is synchronized inside the worker before hand-off, so the main thread consumes fully materialized - GPU batches with no further wait. A ``maxsize=1`` queue keeps every rank exactly one *window* - ahead, so all ranks' Gloo/NCCL collectives stay in lockstep — re-synced each window by the - optimizer/DDP all-reduce barrier between consumptions. The worker uses the (separate) - ``dp_group_nccl`` exclusively, so its all-to-all never races the main thread's bridge/DDP - collectives on their own PGs. ``CUDA_DEVICE_MAX_CONNECTIONS=1`` is refused for ``overlap=True``: - it serializes both onto one hardware queue and the rank-dependent issue order deadlocks. Up to three windows - stay resident (consuming + queued + being exchanged) ≈ ``3·W`` micro-batches — the headline memory cost. - - With ``overlap=False`` the window exchange runs synchronously in ``__next__`` (no thread). + ``window_size`` micro-batches are balanced in one :func:`exchange_window` call and served one + at a time; at ``dp_size == 1`` they are permuted locally. With ``overlap=True`` a background + thread exchanges the next window on a dedicated CUDA stream and the side ``dp_group_nccl`` + while the main thread computes the current one, so up to three windows are resident at once. + ``CUDA_DEVICE_MAX_CONNECTIONS=1`` is rejected with ``overlap=True`` because the side + all-to-all and the DDP all-reduce then share one hardware queue and can deadlock. + + Args: + data_iterator: Source iterator yielding this rank's micro-batch shards. + dp_rank: This rank's module-local DP rank. + dp_size: This module's DP size. + n_groups: Canonical group count (LCM of the module DP sizes). + cost_of: ``callable(flat) -> float`` per-sample cost. + dp_group_gloo: Gloo process group for the metadata all-gather. + dp_group_nccl: NCCL process group for the sample all-to-all. + overlap: Exchange the next window on a background thread. + image_count_of: Optional ``callable(batch) -> LongTensor`` per-sample image counts. + window_size: Micro-batches balanced per exchange. + + Raises: + ValueError: If ``window_size < 1`` or ``overlap=True`` under ``CUDA_DEVICE_MAX_CONNECTIONS=1``. """ def __init__( @@ -1338,22 +1110,13 @@ def __init__( self._cuda = torch.cuda.is_available() self._stream = torch.cuda.Stream() if self._cuda else None self._thread = None - # Sync window state (used whenever the prefetch thread is not running): the active window's - # rebalanced micro-batches and the cursor serving them one at a time. self._active: List[Dict[str, Any]] = [] self._cursor = 0 - # Count of window exchanges run, for throttling the balance probe (first few + every 50th). self._exchanges = 0 - # Cross-window prefetch: a thread exchanges window t+1 while the main thread computes - # window t. Started for any window_size (W == 1 is the one-step-ahead per-micro-batch case). if overlap and dp_size > 1 and os.environ.get("CUDA_DEVICE_MAX_CONNECTIONS") == "1": - # With CUDA_DEVICE_MAX_CONNECTIONS=1 every stream maps onto one hardware queue, so the - # worker's side-PG all-to-all and the main thread's DDP all-reduce execute strictly in - # issue order. The two threads issue in rank-dependent order, so one rank can queue - # [a2a, all-reduce] while a peer queues [all-reduce, a2a]; each NCCL kernel then waits - # for a peer kernel that cannot start, and training deadlocks (reproduced at iteration - # 290 on 4 ranks and at iteration 33 with TP=2 on 6 ranks). Refuse rather than degrade; - # the overlap=False path issues both collectives from one thread and is safe. + # With one hardware queue the worker's all-to-all and the main thread's all-reduce run + # in issue order, which differs per rank, so peers can each wait on a kernel that cannot + # start. Refuse rather than degrade; overlap=False issues both from one thread. raise ValueError( "ReorderingBuffer overlap=True deadlocks under CUDA_DEVICE_MAX_CONNECTIONS=1 (single " "hardware queue: the side-PG all-to-all and the DDP all-reduce can wait on each other). " @@ -1377,14 +1140,14 @@ def shutdown(self) -> None: self._thread.join(timeout=1.0) def __del__(self) -> None: - """Best-effort cleanup for runs that drop the iterator without closing it.""" + """Shut down the worker for runs that drop the iterator without closing it.""" try: self.shutdown() except Exception: logger.debug("Ignoring exception while shutting down ReorderingBuffer.", exc_info=True) def _exchange_window(self, batches: List[Dict[str, Any]], stream: "Any") -> List[Dict[str, Any]]: - # Throttle the balance probe: log the first few windows, then every 50th. + # Probe the first few windows, then every 50th, to keep the log quiet. self._exchanges += 1 probe = self._exchanges <= 3 or self._exchanges % 50 == 0 return exchange_window( @@ -1401,7 +1164,7 @@ def _exchange_window(self, batches: List[Dict[str, Any]], stream: "Any") -> List ) def _collect_window(self) -> List[Dict[str, Any]]: - """Pull up to ``window_size`` micro-batches from the source iterator (fewer at the tail).""" + """Pull up to ``window_size`` micro-batches from the source iterator.""" batches: List[Dict[str, Any]] = [] for _ in range(self._window_size): try: @@ -1411,16 +1174,11 @@ def _collect_window(self) -> List[Dict[str, Any]]: return batches def _refill_window(self) -> None: - """Collect up to ``window_size`` micro-batches, exchange the window, reset the cursor. - - A short final window (fewer than ``window_size`` micro-batches before ``StopIteration``) is - exchanged at its actual size. ``self._active`` is left empty when the iterator is exhausted. - """ + """Collect and exchange the next window synchronously; leave ``_active`` empty when exhausted.""" batches = self._collect_window() if not batches: self._active = [] elif self._dp_size <= 1: - # dp_size == 1: nothing to exchange, apply the canonical order locally. self._active = reorder_window_local( batches, n_groups=self._n_groups, cost_of=self._cost_of, image_count_of=self._image_count_of ) @@ -1429,7 +1187,7 @@ def _refill_window(self) -> None: self._cursor = 0 def _put_worker_item(self, item: Any) -> None: - """Put an item unless shutdown was requested while the bounded queue was full.""" + """Put an item on the bounded queue unless shutdown is requested while it is full.""" while not self._stop.is_set(): try: self._queue.put(item, timeout=0.1) @@ -1438,10 +1196,8 @@ def _put_worker_item(self, item: Any) -> None: continue def _worker(self) -> None: - # Cross-window prefetch: exchange window t+1 while the main thread computes window t. - # Each queue item is the LIST of W rebalanced micro-batches for one window. Everything - # (Gloo metadata all-gather + NCCL all-to-all + deserialize/merge) runs here on self._stream; - # we synchronize that stream before queuing so the consumer gets ready GPU batches. + # The whole exchange runs on self._stream; it is synchronized here, on the worker, so the + # consumer receives materialized GPU batches without waiting. torch.cuda.set_device(self._device) try: while not self._stop.is_set(): @@ -1449,8 +1205,8 @@ def _worker(self) -> None: if not batches: break with torch.cuda.stream(self._stream): - out = self._exchange_window(batches, stream=None) # a2a on the current (self._stream) - self._stream.synchronize() # blocks the WORKER (not the main thread) + out = self._exchange_window(batches, stream=None) + self._stream.synchronize() self._put_worker_item(out) except BaseException as exc: self._put_worker_item(_WorkerException(exc)) @@ -1459,8 +1215,6 @@ def _worker(self) -> None: def __next__(self) -> Dict[str, Any]: if self._thread is None: - # Synchronous path (no prefetch thread): serve the active window from the cursor, - # refilling (collect W -> one exchange) when it is exhausted. if self._cursor >= len(self._active): self._refill_window() if not self._active: @@ -1468,8 +1222,6 @@ def __next__(self) -> Dict[str, Any]: out = self._active[self._cursor] self._cursor += 1 return out - # Cross-window prefetch path: serve from the prefetched window; pull the next one (a list of - # W micro-batches) from the worker when the cursor is exhausted. if self._cursor >= len(self._active): item = self._queue.get() if item is _SENTINEL: @@ -1486,18 +1238,14 @@ def __next__(self) -> Dict[str, Any]: def build_module_dp_process_groups(dp_group_nccl_main: "Any", *, overlap: bool) -> "Tuple[int, int, Any, Any]": - """Create the Gloo (metadata) and side NCCL (exchange) PGs for this rank's module DP group. + """Create the Gloo and side NCCL process groups for this rank's module DP group. - ``dist.new_group`` is collective over the world PG, so every rank must create every module's - DP groups in the same order. We all-gather each rank's DP-group global ranks, deduplicate and - sort, then create a Gloo group (and, for overlap, a **separate** NCCL group so the exchange - all-to-all does not serialize against the gradient all-reduce) for each, keeping the ones this - rank belongs to. + ``dist.new_group`` is collective over the world, so every rank creates every module's groups + in the same sorted order and keeps the ones it belongs to. Args: dp_group_nccl_main: This rank's main DP NCCL group (``grid.get_pg(["dp"])``). - overlap: When ``True``, allocate a separate NCCL PG for the side-stream exchange; - otherwise reuse the main DP NCCL group (sync mode). + overlap: Create a separate NCCL group for the side-stream exchange; otherwise reuse the main one. Returns: ``(dp_rank, dp_size, dp_group_gloo, dp_group_nccl)``. diff --git a/tests/functional_tests/test_groups/megatron_mimo/test_reorder_exchange.py b/tests/functional_tests/test_groups/megatron_mimo/test_reorder_exchange.py index c4fe6225b5c..0f090785707 100644 --- a/tests/functional_tests/test_groups/megatron_mimo/test_reorder_exchange.py +++ b/tests/functional_tests/test_groups/megatron_mimo/test_reorder_exchange.py @@ -12,18 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Functional test for the MegatronMIMO intra-microbatch reorder exchange on a real 2-rank -process group. - -The CPU unit tests (``tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py``) emulate the -``all_to_all_single``; this exercises the on-device path that those cannot — the real Gloo cost -all-gather + NCCL ragged ``all_to_all_single`` inside :func:`exchange_window`, plus the side-PG -setup in :func:`build_module_dp_process_groups`. The 2-rank world is treated as a single module's -``dp=2`` group (het-DP pairing needs two modules / 4 ranks and stays unit-emulated). +"""MegatronMIMO intra-microbatch reorder exchange on a real 2-rank process group (NCCL + Gloo side PGs). Run by ``tests/functional_tests/launch_scripts/h100/active/L0_Launch_training_megatron_mimo.sh`` -(``torch.distributed.run --nproc_per_node=2 -m pytest .../test_groups/megatron_mimo``). -""" +(``torch.distributed.run --nproc_per_node=2 -m pytest .../test_groups/megatron_mimo``).""" from __future__ import annotations @@ -43,19 +35,13 @@ from tests.functional_tests.utils import initialize_distributed -_VISION_START = 990 # synthetic vision-start token; one per image, counted by image_count_of +_VISION_START = 990 # one per image; counted by image_count_of _SEQ = 8 _DIM = 4 def _make_shard(samples: List[Tuple[int, List[int]]]) -> dict: - """Build one rank's nested micro-batch from ``[(global_idx, [patches_per_image, ...]), ...]``. - - ``input_ids[i]`` is tagged with ``1 + g*100`` (so the global index is recoverable as - ``(first_token - 1)//100``) and carries one ``_VISION_START`` token per image. Each image - contributes ``p`` ``hidden_states`` rows all filled with ``float(g)`` so vision can be traced to - its owning sample; a text-only sample (no images) contributes empty ``[0, d]`` / ``[0, 3]``. - """ + """One rank's micro-batch from ``[(g, [patches, ...]), ...]``; first token is ``1 + g*100``, vision rows ``g``.""" b = len(samples) rows = [] grids: List[List[int]] = [] @@ -93,9 +79,7 @@ def _expected_owned(global_patches: List[List[int]], dp_rank: int, dp_size: int, def _run_exchange_case(global_samples: List[List[int]], image_count_of: "Any") -> None: - """Shard ``global_samples`` (patch lists per global sample) contiguously across 2 ranks, run the - real exchange, and assert this rank recovers exactly its cost-balanced shard with vision intact. - """ + """Shard contiguously across 2 ranks, run the real exchange, and check this rank's balanced shard.""" dp_size = dist.get_world_size() n_groups = dp_size b = len(global_samples) @@ -125,8 +109,7 @@ def _run_exchange_case(global_samples: List[List[int]], image_count_of: "Any") - got_globals = [int((int(row[0].item()) - 1) // 100) for row in out["input_ids"].cpu()] assert got_globals == expected_globals, (rank, got_globals, expected_globals) - # Vision must travel with its sample through the real all-to-all: reconstruct the expected - # concatenation (per-sample patch blocks tagged float(g)) and compare byte-for-byte. + # Vision must travel with its sample through the real all-to-all exp_hidden = [torch.full((sum(global_samples[g]), _DIM), float(g), dtype=torch.float32) for g in expected_globals] exp_grid = [ torch.tensor([[1, 1, p] for p in global_samples[g]], dtype=torch.int64).reshape(-1, 3) @@ -142,15 +125,13 @@ def _run_exchange_case(global_samples: List[List[int]], image_count_of: "Any") - @pytest.mark.run_only_on("GPU") def test_reorder_exchange_dp2_single_image(): - """One image per sample; per-sample patch counts force a real cross-rank swap.""" initialize_distributed() if not torch.cuda.is_available(): pytest.skip("CUDA is required for this functional test") if dist.get_world_size() != 2: pytest.skip("This functional test requires exactly 2 ranks") - # costs [4,3,1,2]: balanced owner is g0,g2 -> rank0 and g1,g3 -> rank1, but the contiguous - # initial shards are [g0,g1] / [g2,g3], so rank0 must SEND g1 and RECEIVE g2 (a real exchange). + # costs [4,3,1,2]: owners g0,g2 -> rank0, g1,g3 -> rank1; initial shards [g0,g1] / [g2,g3] force a swap global_samples = [[4], [3], [1], [2]] image_count_of = lambda b: (b["input_ids"] == _VISION_START).sum(dim=1).to(torch.long) # noqa: E731 _run_exchange_case(global_samples, image_count_of) @@ -159,16 +140,13 @@ def test_reorder_exchange_dp2_single_image(): @pytest.mark.run_only_on("GPU") def test_reorder_exchange_dp2_variable_images(): - """Mixed image counts per sample (text-only / single / multi) survive the real exchange.""" initialize_distributed() if not torch.cuda.is_available(): pytest.skip("CUDA is required for this functional test") if dist.get_world_size() != 2: pytest.skip("This functional test requires exactly 2 ranks") - # counts [2,1,0,1] (g2 text-only, g0 multi-image), patches=2/image -> costs [4,2,0,2]: owner is - # g0,g2 -> rank0 and g1,g3 -> rank1; initial shards [g0,g1] / [g2,g3], so a multi-image sample - # stays, a single-image (g1) and a TEXT-ONLY (g2) sample swap across ranks. + # counts [2,1,0,1] at 2 patches/image -> costs [4,2,0,2]: g1 (single-image) and g2 (text-only) swap ranks global_samples = [[2, 2], [2], [], [2]] image_count_of = lambda b: (b["input_ids"] == _VISION_START).sum(dim=1).to(torch.long) # noqa: E731 _run_exchange_case(global_samples, image_count_of) diff --git a/tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py b/tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py index 0368ebc9db8..9ed8ea5d938 100644 --- a/tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py +++ b/tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py @@ -12,12 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for the per-sample reorder-buffer building blocks (no process group). - -The all-to-all is emulated on CPU: each rank's ``send_buf`` is split by ``send_splits`` into -per-destination chunks, and each destination's ``recv_buf`` is the concatenation of the chunks -sent to it in source-rank order — exactly the contract of ``dist.all_to_all_single``. -""" +"""Unit tests for the per-sample reorder-buffer building blocks (all-to-all emulated on CPU).""" import pytest import torch @@ -52,11 +47,9 @@ def _make_sample(global_idx: int, n_tokens: int, n_patches: int) -> dict: - """A ragged synthetic MIMO-like sample, content keyed off global_idx for identity checks.""" return { "input_ids": torch.arange(global_idx * 1000, global_idx * 1000 + n_tokens, dtype=torch.int64), "loss_mask": torch.ones(n_tokens, dtype=torch.bool), - # vision patches: [n_patches, 4], 0 patches -> None (text-only) "hidden_states": (torch.full((n_patches, 4), float(global_idx), dtype=torch.float32) if n_patches else None), } @@ -73,7 +66,6 @@ def test_serialize_deserialize_roundtrip(): assert keys == ["hidden_states", "input_ids", "loss_mask"] # sorted, None dropped buf = serialize_sample(flat, keys) - # byte size from metadata alone must match the actual buffer assert buf.numel() == sample_byte_size(meta, keys) assert buf.numel() % 8 == 0 # 8-byte aligned @@ -81,7 +73,6 @@ def test_serialize_deserialize_roundtrip(): assert offset == buf.numel() for k in keys: assert torch.equal(out[k], flat[k]), k - # None keys are restored assert set(out.keys()) == set(flat.keys()) @@ -122,9 +113,7 @@ def test_unknown_metadata_dtype_raises(): def test_tensor_metadata_rejects_unsupported_dtype_early(): - # An out-of-map dtype (float64) must fail at metadata creation on the producing rank, - # before the all-gather — not only later as a receiver-side byte-size/deserialize error. - # Reuses the same _dtype_info check the decode path uses. + # float64 is outside the dtype map: fail on the producing rank, before the all-gather. flat = {"input_ids": torch.tensor([1, 2]), "weird": torch.zeros(2, dtype=torch.float64)} with pytest.raises(ValueError, match="Unsupported tensor dtype"): tensor_metadata(flat) @@ -136,24 +125,22 @@ def test_tensor_metadata_rejects_unsupported_dtype_early(): def test_balanced_assignment_contiguous_blocks_and_balance(): - # 8 samples, dp=2, canonical n_groups=2. costs = [10.0, 1.0, 9.0, 2.0, 8.0, 3.0, 7.0, 4.0] n_groups, dp_size = 2, 2 assignment = balanced_assignment(costs, n_groups, dp_size) n = len(costs) local = n // dp_size - # each rank owns a contiguous canonical block [r*local, (r+1)*local) per_rank = {0: [], 1: []} pos_seen = set() for global_idx, (owner, pos) in enumerate(assignment): per_rank[owner].append(pos) pos_seen.add(pos) - assert pos_seen == set(range(n)) # every canonical position used once + assert pos_seen == set(range(n)) assert sorted(per_rank[0]) == list(range(0, local)) assert sorted(per_rank[1]) == list(range(local, n)) - # cost balance: each rank's total cost within ~max single item of the other + # LPT bound: rank totals differ by at most one item rank_cost = {0: 0.0, 1: 0.0} for global_idx, (owner, _pos) in enumerate(assignment): rank_cost[owner] += costs[global_idx] @@ -190,14 +177,14 @@ class TestBalancedIndexOrder: def test_valid_permutation_and_equal_groups(self): costs = [5.0, 1.0, 4.0, 2.0, 3.0, 6.0] order = balanced_index_order(costs, n_groups=2) - assert sorted(order) == list(range(6)) # permutation + assert sorted(order) == list(range(6)) g0, g1 = order[:3], order[3:] assert len(g0) == 3 and len(g1) == 3 def test_balances_contiguous_shards(self): # Grouped heavies: a naive contiguous split would put all three 10s on one shard. costs = [10.0, 10.0, 10.0, 1.0, 1.0, 1.0] - order = balanced_index_order(costs, n_groups=3) # 2 per group + order = balanced_index_order(costs, n_groups=3) groups = [order[0:2], order[2:4], order[4:6]] totals = [sum(costs[i] for i in g) for g in groups] assert max(totals) - min(totals) == 0.0 # each group one heavy + one light @@ -215,11 +202,7 @@ def test_errors(self): @pytest.mark.parametrize("module_dps", [(1, 2), (2, 4), (2, 3)]) def test_het_dp_pairing_on_canonical_grid(module_dps): - """n_groups = canonical_grid_size(module dps) (the LCM), the geometry the scalable loader reads with. - - Every module must derive the same canonical order, and each module's owner rank must be the - rank that covers that canonical group under covered_canonical_groups. - """ + """Every module derives the same canonical order and its owner covers that canonical group.""" grid = canonical_grid_size(list(module_dps)) n = grid * 3 # 3 samples per canonical group costs = [float((7 * g) % 11 + 1) for g in range(n)] @@ -240,7 +223,6 @@ def test_het_dp_pairing_consistency(): vis = balanced_assignment(costs, n_groups, dp_size=2) lang = balanced_assignment(costs, n_groups, dp_size=4) - # canonical_pos must be identical across modules (same balanced order) for g in range(len(costs)): assert vis[g][1] == lang[g][1], g @@ -258,7 +240,6 @@ def test_het_dp_pairing_consistency(): def _emulate_all_to_all(plans, dp_size): """Emulate dist.all_to_all_single across ``dp_size`` ranks from their plans.""" - # split each source's send_buf into per-destination chunks chunks = {} for s in range(dp_size): buf = plans[s]["send_buf"] @@ -269,7 +250,6 @@ def _emulate_all_to_all(plans, dp_size): off += nbytes assert off == buf.numel(), f"rank {s} send_splits do not cover send_buf" - # each destination receives chunks in source-rank order; verify split symmetry recv_bufs = {} for d in range(dp_size): parts = [] @@ -283,16 +263,11 @@ def _emulate_all_to_all(plans, dp_size): def _run_exchange(costs, dp_size, n_groups, sample_specs): - """End-to-end: read disjoint shards -> assign -> plan -> emulate a2a -> reconstruct. - - sample_specs[global_idx] = (n_tokens, n_patches). Initial (disjoint, contiguous) layout: - rank r holds global indices [r*local, (r+1)*local). - """ + """sample_specs[g] = (n_tokens, n_patches); rank r initially holds globals [r*local, (r+1)*local).""" n = len(costs) local = n // dp_size originals = {g: _make_sample(g, *sample_specs[g]) for g in range(n)} - # initial disjoint layout rank_global_indices = [list(range(r * local, (r + 1) * local)) for r in range(dp_size)] rank_flats = [[originals[g] for g in rank_global_indices[r]] for r in range(dp_size)] all_tensor_meta = [[tensor_metadata(f) for f in rank_flats[r]] for r in range(dp_size)] @@ -317,7 +292,6 @@ def _run_exchange(costs, dp_size, n_groups, sample_specs): recv_bufs = _emulate_all_to_all(plans, dp_size) - # reconstruct each rank's balanced shard for d in range(dp_size): recovered = [] for _dst_slot, canonical_pos, local_idx in plans[d]["local_samples"]: @@ -333,10 +307,8 @@ def _run_exchange(costs, dp_size, n_groups, sample_specs): recovered.sort(key=lambda x: x[0]) positions = [p for p, _ in recovered] - # rank d owns a contiguous canonical block assert positions == list(range(d * local, (d + 1) * local)), (d, positions) - # tensors survived the round-trip intact, matched by global index for canonical_pos, flat in recovered: g = pos2gidx[canonical_pos] for k, v in originals[g].items(): @@ -347,7 +319,6 @@ def _run_exchange(costs, dp_size, n_groups, sample_specs): def test_exchange_dp2_mixed_modality(): - # 8 samples, some text-only (0 patches), varied token/patch loads sample_specs = { 0: (10, 5), 1: (3, 0), @@ -374,17 +345,14 @@ def test_exchange_dp4_canonical(): def _make_mimo_batch(patches_per_sample): - """Synthetic nested MIMO micro-batch: one image per sample, MRoPE position_ids.""" b = len(patches_per_sample) s = 6 - # Start at 1 (not 0) so no token collides with the default pad_token_id=0: every token is real, - # which keeps the real-token count == s for the joint-cost test. + # Start at 1 so no token collides with pad_token_id=0 (the joint-cost test counts real tokens). input_ids = torch.stack([torch.arange(1 + i * 100, 1 + i * 100 + s, dtype=torch.int64) for i in range(b)]) labels = input_ids + 1 loss_mask = torch.ones(b, s, dtype=torch.bool) position_ids = torch.stack([input_ids, input_ids, input_ids], dim=0) # [3, B, S] MRoPE - # one image per sample; grid [1, h, w] with prod(grid) = patches - grids = torch.tensor([[1, 1, p] for p in patches_per_sample], dtype=torch.int64) # [B, 3] + grids = torch.tensor([[1, 1, p] for p in patches_per_sample], dtype=torch.int64) hidden = torch.cat( [torch.full((p, 4), float(i), dtype=torch.float32) for i, p in enumerate(patches_per_sample)], dim=0 ) @@ -414,7 +382,6 @@ def test_split_merge_roundtrip_identity(): batch = _make_mimo_batch([2, 5, 1, 3]) samples = split_microbatch(batch) assert len(samples) == 4 - # per-sample flat dicts carry dotted vision keys assert "modality_inputs.images.enc.hidden_states" in samples[0] merged = merge_samples(samples) _assert_batch_equal(merged, batch) @@ -428,9 +395,7 @@ def test_split_merge_roundtrip_identity(): def _flat_vision_sample(idx: int, n_tokens: int, n_patches: int, *, none_vision: bool = False) -> dict: - """Per-sample flat dict as ``split_microbatch`` produces: ``[1, T]`` input_ids and one image's - ``[p, 4]`` hidden_states + ``[1, 3]`` grid_thw. ``n_patches==0`` is text-only: empty ``[0, 4]`` / - ``[0, 3]`` (or ``None`` when ``none_vision`` to exercise the None-vision merge path).""" + """Flat dict as ``split_microbatch`` produces; ``n_patches == 0`` is text-only (empty or None vision).""" input_ids = torch.arange(idx * 100, idx * 100 + n_tokens, dtype=torch.int64).reshape(1, -1) if n_patches == 0: if none_vision: @@ -445,15 +410,14 @@ def _flat_vision_sample(idx: int, n_tokens: int, n_patches: int, *, none_vision: def test_merge_text_only_first(): - # Sample 0 is text-only with vision keys as None (the case the old ``ref = samples[0]`` - # classification silently dropped); samples 1-2 carry one image each. + # Sample 0 is text-only with None vision keys; samples 1-2 carry one image each. flats = [ _flat_vision_sample(0, n_tokens=4, n_patches=0, none_vision=True), _flat_vision_sample(1, n_tokens=4, n_patches=3), _flat_vision_sample(2, n_tokens=4, n_patches=2), ] enc = merge_samples(flats)["modality_inputs"]["images"]["enc"] - # 0 + 3 + 2 = 5 patches, 0 + 1 + 1 = 2 image rows — no silent loss of sample 1/2 vision. + # 0 + 3 + 2 = 5 patches, 0 + 1 + 1 = 2 image rows assert enc["hidden_states"].shape == (5, 4) assert enc["grid_thw"].shape == (2, 3) assert torch.equal(enc["hidden_states"][:3], torch.full((3, 4), 1.0)) @@ -468,11 +432,11 @@ def test_merge_all_text_only(): def test_serialize_empty_vision_roundtrip(): - flat = _flat_vision_sample(0, n_tokens=4, n_patches=0) # [0,4] / [0,3] empties + flat = _flat_vision_sample(0, n_tokens=4, n_patches=0) meta = tensor_metadata(flat) keys = sample_keys(meta) buf = serialize_sample(flat, keys) - assert buf.numel() == sample_byte_size(meta, keys) # empties contribute 0 bytes + assert buf.numel() == sample_byte_size(meta, keys) out, offset = deserialize_sample(buf, 0, meta, keys) assert offset == buf.numel() assert out[f"{_VKEY}.hidden_states"].shape == (0, 4) @@ -486,21 +450,17 @@ def test_serialize_empty_vision_roundtrip(): def _make_mimo_batch_var(images_per_sample: list[int]) -> dict: - """Nested MIMO micro-batch with a variable number of images per sample (0/1/N). - - ``images_per_sample[s]`` images for sample ``s``; each image gets a distinct patch count and its - ``hidden_states`` rows are tagged with its global image index so reorder/merge can be content-checked. - """ + """Nested micro-batch with ``images_per_sample[s]`` images; hidden rows are tagged with the global image index.""" b = len(images_per_sample) seq = 6 input_ids = torch.stack([torch.arange(i * 100, i * 100 + seq, dtype=torch.int64) for i in range(b)]) - position_ids = torch.stack([input_ids, input_ids, input_ids], dim=0) # [3, B, S] MRoPE + position_ids = torch.stack([input_ids, input_ids, input_ids], dim=0) grids: list[list[int]] = [] hidden_blocks: list[torch.Tensor] = [] img_gid = 0 for n_img in images_per_sample: for _ in range(n_img): - p = img_gid + 2 # distinct, identifiable patch count per image + p = img_gid + 2 grids.append([1, 1, p]) hidden_blocks.append(torch.full((p, 4), float(img_gid), dtype=torch.float32)) img_gid += 1 @@ -565,17 +525,14 @@ def test_soft_validation_mismatch(): def test_split_one_image_per_sample_no_cu_img(): - # one-image-per-sample works with cu_img=None (fallback path) batch = _make_mimo_batch([2, 5, 1, 3]) _assert_batch_equal(merge_samples(split_microbatch(batch)), batch) - # multi-image without cu_img raises a clear error multi = _make_mimo_batch_var([2, 1]) # 3 images, 2 samples with pytest.raises(ValueError, match="one image per sample"): split_microbatch(multi) def _run_exchange_nested(images_per_sample, dp_size, n_groups): - """Full split→assign→exchange→reconstruct on nested variable-image batches (emulated a2a).""" n = len(images_per_sample) local = n // dp_size full = _make_mimo_batch_var(images_per_sample) @@ -626,7 +583,6 @@ def _run_exchange_nested(images_per_sample, dp_size, n_groups): def test_exchange_var_images_dp2(): - # 8 samples: text-only (0), single (1), and multi-image (2/3) mixed so a rank owns a variety. _run_exchange_nested([2, 0, 1, 0, 3, 1, 0, 2], dp_size=2, n_groups=2) @@ -636,8 +592,6 @@ def test_exchange_var_images_dp2(): def _emulate_window_exchange(rank_flats, rank_global_indices, all_tensor_meta, route, *, dp_size, window_size, local): - """Emulate the on-device exchange on CPU and reassemble each rank's window via the production - :func:`reassemble_window`. Returns ``per_rank_batches[d][slot]``.""" plans = [ prepare_sample_exchange( local_flats=rank_flats[r], @@ -667,8 +621,6 @@ def _emulate_window_exchange(rank_flats, rank_global_indices, all_tensor_meta, r def test_reassemble_window_w1_matches_canonical_merge(): - """With the intra route and W=1, ``reassemble_window`` rebuilds exactly the canonically - ordered balanced micro-batch — byte-identical to merging the globally-balanced samples directly.""" images_per_sample = [2, 0, 1, 0, 3, 1, 0, 2] dp_size, n_groups = 2, 2 n = len(images_per_sample) @@ -690,29 +642,25 @@ def test_reassemble_window_w1_matches_canonical_merge(): ) for d in range(dp_size): - assert len(per_rank[d]) == 1 # W=1 -> one micro-batch + assert len(per_rank[d]) == 1 owned_positions = sorted(pos for g, (owner, pos) in enumerate(assignment) if owner == d) expected = merge_samples([all_flats[pos2gidx[p]] for p in owned_positions]) _assert_batch_equal(per_rank[d][0], expected) def test_reassemble_window_w2_3d_route(): - """A hand-built W=2 3D route (no balancer) routes samples into specific slots; assert each - slot reassembles its sample, a text-only sample lands as an empty-vision slot, vision is intact.""" - # 4 samples, dp=2, W=2 slots, local=1 per (slot, rank). Layout: rank0 holds globals [0,1], - # rank1 holds [2,3]. g2 is text-only (0 patches); others carry distinct image patch counts. dp_size, window_size, local = 2, 2, 1 originals = [ _flat_vision_sample(0, n_tokens=4, n_patches=3), _flat_vision_sample(1, n_tokens=4, n_patches=2), - _flat_vision_sample(2, n_tokens=4, n_patches=0), # text-only + _flat_vision_sample(2, n_tokens=4, n_patches=0), _flat_vision_sample(3, n_tokens=4, n_patches=1), ] rank_global_indices = [[0, 1], [2, 3]] rank_flats = [[originals[g] for g in rank_global_indices[r]] for r in range(dp_size)] all_tensor_meta = [[tensor_metadata(f) for f in rank_flats[r]] for r in range(dp_size)] - # route[g] = (dst_slot, owner_rank, canonical_pos). Cross-rank moves into a chosen slot: + # route[g] = (dst_slot, owner_rank, canonical_pos): # g0 -> slot0,rank0 g1 -> slot1,rank1 g2 -> slot0,rank1 g3 -> slot1,rank0 route = [(0, 0, 0), (1, 1, 0), (0, 1, 0), (1, 0, 0)] @@ -723,11 +671,10 @@ def test_reassemble_window_w2_3d_route(): def first_token(batch): return int(batch["input_ids"][0, 0].item()) - # rank0: slot0 <- g0, slot1 <- g3 (received). rank1: slot0 <- g2 (text-only), slot1 <- g1. + # rank0: slot0 <- g0, slot1 <- g3. rank1: slot0 <- g2, slot1 <- g1. assert [first_token(b) for b in per_rank[0]] == [0, 300] assert [first_token(b) for b in per_rank[1]] == [200, 100] - # Vision attribution per slot: g3 (1 patch) on rank0 slot1; g2 (text-only -> [0,4]) on rank1 slot0. r0_slot1_enc = per_rank[0][1]["modality_inputs"]["images"]["enc"] assert r0_slot1_enc["hidden_states"].shape == (1, 4) assert torch.equal(r0_slot1_enc["hidden_states"], torch.full((1, 4), 3.0)) @@ -736,14 +683,12 @@ def first_token(batch): def test_reassemble_window_slot_shape_guard(): - """The slot-shape invariant fires loudly when a slot ends up with the wrong sample count.""" - # Route both rank0's samples into the same slot of rank0 -> slot0 has 2, slot1 has 0 (expected 1). dp_size, window_size, local = 2, 2, 1 originals = [_flat_vision_sample(i, n_tokens=4, n_patches=1) for i in range(4)] rank_global_indices = [[0, 1], [2, 3]] rank_flats = [[originals[g] for g in rank_global_indices[r]] for r in range(dp_size)] all_tensor_meta = [[tensor_metadata(f) for f in rank_flats[r]] for r in range(dp_size)] - # g0,g1 both -> slot0/rank0; g2,g3 -> slot0/rank1. Slot0 over-fills (2) -> guard fires on slot 0. + # g0,g1 both -> slot0/rank0; g2,g3 -> slot0/rank1: slot 0 over-fills (2), slot 1 stays empty. route = [(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1)] with pytest.raises(RuntimeError, match="slot 0 reassembled 2 samples, expected 1"): _emulate_window_exchange( @@ -763,7 +708,6 @@ def test_reassemble_window_slot_shape_guard(): def test_build_window_route_layout(): - """The window route is slot-major: route[s*B + j] has dst_slot == s and a per-slot canonical_pos.""" dp_size, n_groups = 2, 2 costs_per_slot = [[4.0, 1.0, 3.0, 2.0], [1.0, 2.0, 3.0, 4.0], [2.0, 2.0, 2.0, 2.0]] b_global = len(costs_per_slot[0]) @@ -771,20 +715,14 @@ def test_build_window_route_layout(): assert len(route) == len(costs_per_slot) * b_global for s in range(len(costs_per_slot)): slot_entries = route[s * b_global : (s + 1) * b_global] - # intra invariant: every entry in this slot's block routes to slot s. assert all(dst == s for dst, _owner, _pos in slot_entries) - # canonical positions are a permutation of [0, B) within the slot. assert sorted(pos for _dst, _owner, pos in slot_entries) == list(range(b_global)) - # matches an independent per-slot intra_route. assert slot_entries == intra_route(balanced_assignment(costs_per_slot[s], n_groups, dp_size), src_slot=s) def _run_window_exchange(images_per_slot, dp_size, n_groups): - """Full window exchange over W micro-batches (emulated a2a) using the real window indexing - (build_window_route + reassemble_window). Asserts each slot reassembles to exactly the - independent per-slot intra-balanced shard — i.e. slots never mix (intra invariant).""" + """Window exchange over W micro-batches; each slot must equal its independent per-slot balanced shard.""" window_size = len(images_per_slot) - # Per-slot full batches + flats + costs (slot-local global index j in [0, B)). slot_full = [_make_mimo_batch_var(imgs) for imgs in images_per_slot] slot_flats = [ split_microbatch(slot_full[s], cu_img=_cu_img_from_counts(images_per_slot[s])) for s in range(window_size) @@ -807,7 +745,6 @@ def _run_window_exchange(images_per_slot, dp_size, n_groups): rank_flats, rank_global_indices, all_tensor_meta, route, dp_size=dp_size, window_size=window_size, local=local ) - # Expected: each slot independently balanced (W=1 semantics applied per slot). for s in range(window_size): assignment_s = balanced_assignment(costs_per_slot[s], n_groups, dp_size) pos2gidx_s = {pos: g for g, (_o, pos) in enumerate(assignment_s)} @@ -818,30 +755,22 @@ def _run_window_exchange(images_per_slot, dp_size, n_groups): def test_window_exchange_two_slots_dp2(): - # W=2 window, dp=2; distinct image mixes per slot so each slot triggers a real cross-rank swap. _run_window_exchange([[2, 0, 1, 3], [1, 2, 0, 1]], dp_size=2, n_groups=2) def test_window_exchange_w3_dp2_text_only_and_multi(): - # W=3 window with text-only + multi-image samples; vision attribution per slot at W>1. _run_window_exchange([[0, 2, 1, 0], [3, 0, 0, 1], [1, 1, 1, 1]], dp_size=2, n_groups=2) def test_window_exchange_het_dp_paired_slots(): - # Het-DP at the window level: dp2 and dp4 over the same per-slot costs must agree per slot. images_per_slot = [[2, 0, 1, 3, 1, 2, 0, 1], [1, 1, 2, 0, 3, 0, 1, 2]] for dp_size in (2, 4): _run_window_exchange(images_per_slot, dp_size=dp_size, n_groups=4) def test_window_cost_spread_tightens_after_balance(): - """Balance probe: the per-rank cost spread must tighten (and never widen) after balancing — - the quantitative evidence that reordering evens the per-rank load.""" dp_size, n_groups = 4, 4 - # Heavily skewed costs: the natural contiguous shards are very imbalanced (rank0 holds the big - # samples, rank3 the small), so balancing should shrink the per-rank max/min spread. - # Both slots are skewed across the contiguous shards (rank0 heavy, rank1 light) so each must - # move samples and strictly tighten — slot 1 differs from slot 0 to exercise per-slot balancing. + # Both slots are skewed across the contiguous shards (rank0 heavy, rank3 light), so each must move samples. costs_per_slot = [ [100.0, 90.0, 80.0, 70.0, 8.0, 6.0, 4.0, 2.0, 50.0, 40.0, 30.0, 20.0, 9.0, 7.0, 5.0, 3.0], [80.0, 75.0, 70.0, 65.0, 10.0, 9.0, 8.0, 7.0, 45.0, 40.0, 35.0, 30.0, 12.0, 11.0, 10.0, 9.0], @@ -852,22 +781,19 @@ def test_window_cost_spread_tightens_after_balance(): for sp in spreads: before = sp["before_max"] - sp["before_min"] after = sp["after_max"] - sp["after_min"] - assert after <= before, sp # balancing never widens the spread - assert after < before # and on these skewed inputs it strictly tightens - # these skewed inputs require real movement: 0 < remote <= slot size. + assert after <= before, sp + assert after < before assert all(0 < sp["remote"] <= len(costs_per_slot[0]) for sp in spreads) def test_assert_intra_no_cross_slot_passes_for_window_route(): - # build_window_route output (intra) always satisfies dst_slot == src_slot. dp_size, n_groups = 2, 2 costs_per_slot = [[4.0, 1.0, 3.0, 2.0], [1.0, 2.0, 3.0, 4.0]] - route = build_window_route(costs_per_slot, n_groups, dp_size) # calls the guard internally - assert_intra_no_cross_slot(route, b_global=4) # explicit re-check is a no-op (no raise) + route = build_window_route(costs_per_slot, n_groups, dp_size) + assert_intra_no_cross_slot(route, b_global=4) def test_assert_intra_no_cross_slot_trips_on_cross_slot_route(): - # A hand-built route that sends a slot-0 sample (g=1 -> src_slot 0) into slot 1 must trip the guard. b_global = 2 # 2 slots of 2 samples; g//b_global is the source slot route = [(0, 0, 0), (1, 1, 0), (1, 0, 0), (1, 1, 1)] # g=1 leaks slot0 -> slot1 with pytest.raises(RuntimeError, match="cross-slot leak at global index 1"): @@ -875,7 +801,6 @@ def test_assert_intra_no_cross_slot_trips_on_cross_slot_route(): def test_window_cost_spread_already_balanced_no_remote(): - # Uniform costs are already perfectly balanced: spread stays 0 and nothing needs to move. dp_size, n_groups = 2, 2 costs_per_slot = [[5.0, 5.0, 5.0, 5.0]] route = build_window_route(costs_per_slot, n_groups, dp_size) @@ -889,12 +814,12 @@ def test_window_cost_spread_already_balanced_no_remote(): def test_reordering_buffer_window_cursor_single_rank(): - # n_groups=1 + uniform cost -> identity order; exercises windows of [3, 3, 1] over 7 items. + # n_groups=1 + uniform cost -> identity order; windows of [3, 3, 1] over 7 items. items = [{"input_ids": torch.tensor([[i]])} for i in range(7)] buf = ReorderingBuffer( iter(items), dp_rank=0, - dp_size=1, # single rank -> local reorder, no exchange / no PG needed + dp_size=1, n_groups=1, cost_of=lambda f: 0.0, dp_group_gloo=None, @@ -903,11 +828,10 @@ def test_reordering_buffer_window_cursor_single_rank(): window_size=3, ) got = [int(b["input_ids"].item()) for b in buf] - assert got == list(range(7)) # all items served, short final window handled + assert got == list(range(7)) def test_reordering_buffer_overlap_single_rank_no_thread(): - # overlap=True but dp_size=1: no prefetch thread. items = [{"input_ids": torch.tensor([[i]])} for i in range(5)] buf = ReorderingBuffer( iter(items), @@ -920,19 +844,17 @@ def test_reordering_buffer_overlap_single_rank_no_thread(): overlap=True, window_size=2, ) - assert buf._thread is None # dp_size <= 1 never starts the prefetch thread + assert buf._thread is None assert [int(b["input_ids"].item()) for b in buf] == list(range(5)) def test_reorder_window_local_matches_paired_dp2_order(): - # dp=1 local order must equal the paired dp=2 module's post-exchange order (rank 0's, then rank 1's). n_groups = 2 costs = [1.0, 5.0, 2.0, 4.0] # LPT over 2 groups: [1, 0] and [3, 2] -> canonical order [1, 0, 3, 2] expected = [1, 0, 3, 2] batch = { "input_ids": torch.arange(4).view(4, 1), "sample_ids": ["s0", "s1", "s2", "s3"], - # one image / one patch per sample "vision": {"hidden_states": torch.arange(8.0).view(4, 2), "grid_thw": torch.tensor([[1, 1, 1]] * 4)}, } @@ -950,7 +872,6 @@ def cost_of(flat): assert out["vision"]["hidden_states"].tolist() == [[2.0 * i, 2.0 * i + 1] for i in expected] assert out["vision"]["grid_thw"].shape == (4, 3) - # dp=2 order: rank 0's samples by canonical_pos, then rank 1's. assignment = balanced_assignment(costs, n_groups, dp_size=2) dp2_order = [ g @@ -961,7 +882,6 @@ def cost_of(flat): def test_reordering_buffer_single_rank_serves_canonical_order(): - # Each micro-batch of a window is reordered independently; overlap=True runs inline at dp_size=1. costs = [1.0, 5.0, 2.0, 4.0] batches = [{"input_ids": (torch.arange(4) + 10 * s).view(4, 1)} for s in range(2)] @@ -1003,8 +923,7 @@ def _buffer_kwargs(): def test_reordering_buffer_overlap_refuses_single_hardware_queue(monkeypatch): - # CUDA_DEVICE_MAX_CONNECTIONS=1 serializes the worker's all-to-all behind the DDP all-reduce in - # rank-dependent order (reproduced deadlock); overlap=True must be refused before any thread starts. + # CUDA_DEVICE_MAX_CONNECTIONS=1 serializes the worker all-to-all behind the DDP all-reduce (deadlock). monkeypatch.setenv("CUDA_DEVICE_MAX_CONNECTIONS", "1") with pytest.raises(ValueError, match="CUDA_DEVICE_MAX_CONNECTIONS=1"): ReorderingBuffer(iter([]), overlap=True, **_buffer_kwargs()) @@ -1018,7 +937,7 @@ def test_reordering_buffer_no_overlap_allowed_under_single_hardware_queue(monkey def test_reordering_buffer_overlap_allowed_without_single_hardware_queue(monkeypatch): monkeypatch.delenv("CUDA_DEVICE_MAX_CONNECTIONS", raising=False) - ReorderingBuffer(iter([]), overlap=True, **_buffer_kwargs()) # no raise (thread only starts on CUDA) + ReorderingBuffer(iter([]), overlap=True, **_buffer_kwargs()) # --------------------------------------------------------------------------- @@ -1029,21 +948,20 @@ def test_reordering_buffer_overlap_allowed_without_single_hardware_queue(monkeyp def test_sample_cost_patches_only(): from megatron.bridge.data.megatron_mimo.reorder_buffer import sample_cost - batch = _make_mimo_batch([2, 5]) # sample 0: 2 patches, sample 1: 5 patches; 6 tokens each + batch = _make_mimo_batch([2, 5]) samples = split_microbatch(batch) - # default language_cost_weight=0.0 -> patch-only cost: encoder_cost_weight * sum(prod(grid_thw)); token length not counted + # language_cost_weight defaults to 0.0: patch-only cost c0 = sample_cost(samples[0], encoder_cost_weight=1.0) c1 = sample_cost(samples[1], encoder_cost_weight=1.0) assert c0 == 2 assert c1 == 5 - # encoder_cost_weight scales the patch count assert sample_cost(samples[0], encoder_cost_weight=3.0) == 3 * 2 def test_sample_cost_text_only_is_zero(): from megatron.bridge.data.megatron_mimo.reorder_buffer import sample_cost - flat = {"input_ids": torch.tensor([[5, 5, 0, 0]], dtype=torch.int64)} # no vision -> zero cost + flat = {"input_ids": torch.tensor([[5, 5, 0, 0]], dtype=torch.int64)} c = sample_cost(flat, encoder_cost_weight=1.0) assert c == 0 @@ -1051,7 +969,6 @@ def test_sample_cost_text_only_is_zero(): def test_sample_cost_joint_vit_and_lm(): from megatron.bridge.data.megatron_mimo.reorder_buffer import sample_cost - # The LM term reads the real length from attention_mask (all 6 tokens real here). batch = _make_mimo_batch([2, 5]) # patches 2 and 5; 6 real tokens each samples = split_microbatch(batch) for sample in samples: @@ -1059,14 +976,13 @@ def test_sample_cost_joint_vit_and_lm(): # cost = encoder_cost_weight*patches + language_cost_weight*real_tokens assert sample_cost(samples[0], encoder_cost_weight=1.0, language_cost_weight=0.5) == 2 + 0.5 * 6 assert sample_cost(samples[1], encoder_cost_weight=1.0, language_cost_weight=0.5) == 5 + 0.5 * 6 - # encoder_cost_weight=0 -> token-only cost assert sample_cost(samples[0], encoder_cost_weight=0.0, language_cost_weight=2.0) == 2.0 * 6 def test_sample_cost_lm_term_ignores_padding(): from megatron.bridge.data.megatron_mimo.reorder_buffer import sample_cost - # Same 3 real tokens but different padded widths -> identical LM cost (collation-independent). + # Same 3 real tokens at different padded widths -> identical LM cost short = { "input_ids": torch.tensor([[5, 6, 7, 0]], dtype=torch.int64), "attention_mask": torch.tensor([[1, 1, 1, 0]], dtype=torch.int64), @@ -1082,8 +998,7 @@ def test_sample_cost_lm_term_ignores_padding(): def test_sample_cost_lm_term_does_not_compare_pad_id(): from megatron.bridge.data.megatron_mimo.reorder_buffer import sample_cost - # A real token equal to the pad id (pad_token == eos_token tokenizers) still counts: the length - # comes from attention_mask, never from comparing input_ids against a pad id. + # A real token equal to the pad id still counts: the length comes from attention_mask flat = { "input_ids": torch.tensor([[5, 0, 7, 0]], dtype=torch.int64), "attention_mask": torch.tensor([[1, 1, 1, 0]], dtype=torch.int64), @@ -1095,7 +1010,6 @@ def test_sample_cost_lm_term_requires_attention_mask(): from megatron.bridge.data.megatron_mimo.reorder_buffer import sample_cost flat = {"input_ids": torch.tensor([[5, 6, 7, 0]], dtype=torch.int64)} - # Patch-only cost never needs the mask. assert sample_cost(flat, encoder_cost_weight=1.0) == 0 with pytest.raises(ValueError, match="attention_mask"): sample_cost(flat, encoder_cost_weight=1.0, language_cost_weight=1.0) @@ -1104,37 +1018,29 @@ def test_sample_cost_lm_term_requires_attention_mask(): def test_sample_cost_image_token_id_recovers_patches_from_input_ids(): from megatron.bridge.data.megatron_mimo.reorder_buffer import sample_cost - img = 99 # image-placeholder token id - # 3 placeholder tokens * square_merge_size(4) = 12 patches. + img = 99 + # 3 placeholder tokens * square_merge_size(4) = 12 patches flat = {"input_ids": torch.tensor([[1, img, img, img, 2, 0]], dtype=torch.int64)} assert sample_cost(flat, encoder_cost_weight=1.0, image_token_id=img, square_merge_size=4) == 12 - # encoder_cost_weight scales the recovered patch count. assert sample_cost(flat, encoder_cost_weight=2.0, image_token_id=img, square_merge_size=4) == 24 def test_sample_cost_image_token_id_matches_across_modules(): - """Vision (has grid_thw) and language (grid_thw nulled) must derive the SAME cost. - - With image_token_id wired, the patch cost comes from the module-independent input_ids image-token - count, so a vision-style flat and a language-style flat with identical input_ids cost the same even - though the language flat has no grid_thw (nulled by the #4442 rank-aware metadata collate). - """ + """A vision-style flat (with grid_thw) and a language-style flat (grid_thw nulled) must cost the same.""" from megatron.bridge.data.megatron_mimo.reorder_buffer import sample_cost img = 99 - input_ids = torch.tensor([[1, img, img, 2, 0, 0]], dtype=torch.int64) # 2 image tokens - # Vision-style flat: carries grid_thw (would dominate the grid_thw fallback). + input_ids = torch.tensor([[1, img, img, 2, 0, 0]], dtype=torch.int64) vision_flat = { "input_ids": input_ids, f"{_VKEY}.grid_thw": torch.tensor([[1, 4, 2]], dtype=torch.int64), # prod = 8 patches } - # Language-style flat: grid_thw nulled, only input_ids remains. language_flat = {"input_ids": input_ids} vc = sample_cost(vision_flat, encoder_cost_weight=1.0, image_token_id=img, square_merge_size=4) lc = sample_cost(language_flat, encoder_cost_weight=1.0, image_token_id=img, square_merge_size=4) - assert vc == lc == 2 * 4 # 2 image tokens * square_merge_size, independent of grid_thw - # Without image_token_id the two diverge (the bug): vision reads grid_thw, language gets 0. + assert vc == lc == 2 * 4 # 2 image tokens * square_merge_size + # Without image_token_id the two diverge: vision reads grid_thw, language gets 0 assert sample_cost(vision_flat, encoder_cost_weight=1.0) == 8 assert sample_cost(language_flat, encoder_cost_weight=1.0) == 0 @@ -1147,7 +1053,7 @@ def test_sample_cost_image_token_id_matches_across_modules(): @pytest.mark.unit class TestReorderVisionSingleD2H: def test_reorder_vision_single_d2h(self, monkeypatch): - # Count Tensor.item() calls during the reorder; the cumsum should be a single .tolist(). + # The cumsum must be a single .tolist(): no per-image .item() D2H sync in the loop calls = {"item": 0} orig_item = torch.Tensor.item @@ -1159,9 +1065,7 @@ def counting_item(self): hidden = torch.arange(26, dtype=torch.float32).reshape(26, 1) monkeypatch.setattr(torch.Tensor, "item", counting_item) hs, g = _reorder_vision_by_images(hidden, grid, [2, 0, 1]) - # no per-image .item() D2H sync in the loop (cumsum uses one .tolist()). assert calls["item"] == 0 - # output bit-identical to the expected reorder expected = torch.cat([hidden[20:26], hidden[0:4], hidden[4:20]]) torch.testing.assert_close(hs, expected) @@ -1179,13 +1083,9 @@ def __getitem__(self, idx: int) -> int: @pytest.mark.parametrize("dataloader_type", ["single", "cyclic"]) def test_canonical_grid_shards_are_contiguous_blocks_of_one_global_order(dataloader_type): - """exchange_window assumes rank r holds global positions [r*local, (r+1)*local) of one ordered - micro-batch shared by every module. Under the canonical grid that holds for both samplers: a - rank's window is the concatenation of its groups' windows in group order, and the global order - is the concatenation of all groups' windows. - """ + """exchange_window assumes rank r holds positions [r*local, (r+1)*local) of one order shared by all modules.""" module_dps = [1, 2] - grid = canonical_grid_size(module_dps) # 2 + grid = canonical_grid_size(module_dps) micro_batch = 8 dataset = _IndexDataset(64) From 5cb4faa8028e77cb432ff5407e1ce3a77f91c354 Mon Sep 17 00:00:00 2001 From: "kayeon.song" Date: Sat, 3 Oct 2026 04:04:28 +0000 Subject: [PATCH 08/13] fix(data): ship input_ids and attention_mask to MIMO ranks only when the reorder cost needs them Restore the upstream per-rank field selection in _batch_spec_for_rank: input_ids goes to non-first language stages and encoder ranks only when intra-microbatch reorder is on (its cost is computed on every stage), and attention_mask only when packing or a reorder language cost term needs it. Ranks running without reorder receive the same fields as upstream. Signed-off-by: kayeon.song Signed-off-by: Yoonsik Kim Signed-off-by: Chanwoo Park Co-authored-by: Yoonsik Kim Co-authored-by: Chanwoo Park <45866990+sailor1493@users.noreply.github.com> --- .../qwen35_vl/finetune_qwen35_vl.py | 37 +++++--- .../examples/test_qwen35_vl_mimo_finetune.py | 85 +++++++++++++------ 2 files changed, 84 insertions(+), 38 deletions(-) diff --git a/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py b/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py index 10f21d6e518..050140263b6 100644 --- a/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py +++ b/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py @@ -11,8 +11,8 @@ Conversation examples are built with the standard HF VLM provider, then the resulting Qwen batch is adapted into the MIMO forward shape: - - language inputs: ``input_ids``, MRoPE ``position_ids``, labels, loss mask, and the - tokenizer's ``attention_mask`` + - language inputs: ``input_ids``, MRoPE ``position_ids``, labels, loss mask, and (when packing + or with a reorder language cost term) the tokenizer's ``attention_mask`` - image inputs: ``modality_inputs["images"]["qwen_visual"]`` Example 2-GPU smoke: @@ -133,6 +133,7 @@ class MIMOBatchSpec: labels: bool = True loss_mask: bool = True modality_inputs: bool = True + attention_mask: bool = False def describe(self) -> str: enabled = [ @@ -143,6 +144,7 @@ def describe(self) -> str: ("labels", self.labels), ("loss_mask", self.loss_mask), ("modality_inputs", self.modality_inputs), + ("attention_mask", self.attention_mask), ) if value ] @@ -278,24 +280,39 @@ def _batch_spec_for_rank(cfg: Any) -> MIMOBatchSpec: pp_rank = _grid_dim_rank(grid, "pp", 0) pp_size = _grid_dim_size(grid, "pp", 1) + is_first_pp = pp_rank == 0 is_last_pp = pp_rank == pp_size - 1 - # Cheap [B, S] fields are always shipped; forward_step decides what reaches the model. + dataset_cfg = getattr(cfg, "dataset", None) + packing_active = bool(getattr(dataset_cfg, "enable_in_batch_packing", False)) + reorder_active = bool( + getattr(dataset_cfg, "megatron_mimo_scalable_dp", False) + and getattr(dataset_cfg, "megatron_mimo_intra_microbatch_reorder", False) + ) + language_cost_active = reorder_active and ( + getattr(dataset_cfg, "megatron_mimo_reorder_language_cost_weight", 0.0) > 0 + ) + if module_name == MIMO_LANGUAGE_MODULE_KEY: return MIMOBatchSpec( - input_ids=True, + input_ids=is_first_pp or reorder_active, + # Qwen3.5-VL mRoPE needs position_ids on every language PP stage. position_ids=True, labels=is_last_pp, loss_mask=is_last_pp, modality_inputs=False, + # Packing length source; the packer nulls it again before the model. + attention_mask=packing_active or language_cost_active, ) return MIMOBatchSpec( - input_ids=True, + # Encoder first stages need input_ids to attach per-sample split metadata. + input_ids=is_first_pp or reorder_active, position_ids=False, labels=False, loss_mask=False, - modality_inputs=True, + modality_inputs=is_first_pp, + attention_mask=language_cost_active, ) @@ -313,6 +330,8 @@ def _project_adapted_batch( adapted["loss_mask"] = None if not batch_spec.modality_inputs: adapted["modality_inputs"] = None + if not batch_spec.attention_mask: + adapted["attention_mask"] = None return adapted @@ -985,16 +1004,11 @@ def _build_data_iterators(cfg, _megatron_mimo_infra, *, train_state=None): f"(module {_name!r} has PP={_par.pipeline_model_parallel_size}): non-first encoder stages " "skip data loading, so the collective process-group creation below would hang." ) - # n_groups is the canonical grid size (LCM of the module DP sizes), the geometry the canonical - # loader reads with, so vision and language ranks derive one identical per-sample assignment. balance_n_groups = canonical_grid_size(module_dps) _grid, _ = _find_rank_module(cfg.model._grids) dp_rank, dp_size, dp_group_gloo, dp_group_nccl = build_module_dp_process_groups( _grid.get_pg(["dp"]), overlap=cfg.dataset.megatron_mimo_reorder_overlap ) - # Cost = encoder_cost_weight * patches + language_cost_weight * real_tokens. The patch count comes from the - # image-placeholder token count in input_ids (present on every module and PP stage), not - # from grid_thw, which the metadata collate omits on language shards. cost_of = functools.partial( sample_cost, encoder_cost_weight=cfg.dataset.megatron_mimo_reorder_encoder_cost_weight, @@ -1004,7 +1018,6 @@ def _build_data_iterators(cfg, _megatron_mimo_infra, *, train_state=None): ) def image_count_of(b: dict[str, Any]) -> torch.Tensor: - # One vision-start token precedes each image, so this counts images, not patches. return (b["input_ids"] == spec.vision_start_token_id).sum(dim=1).to(torch.long) loader_iter = ReorderingBuffer( diff --git a/tests/unit_tests/examples/test_qwen35_vl_mimo_finetune.py b/tests/unit_tests/examples/test_qwen35_vl_mimo_finetune.py index 086689b1ef6..604ae17a0fa 100644 --- a/tests/unit_tests/examples/test_qwen35_vl_mimo_finetune.py +++ b/tests/unit_tests/examples/test_qwen35_vl_mimo_finetune.py @@ -287,43 +287,76 @@ def test_qwen35_vl_mimo_rejects_truncated_visual_tokens(): sys.modules.pop(name, None) -def test_batch_spec_language_non_first_stage_ships_cheap_fields(monkeypatch): - # Non-first language stage: input_ids/position_ids always shipped (reorder cost, MRoPE); labels and - # loss_mask only on the last stage. forward_step decides what the model receives. +def _spec_for(module, module_name, pp_rank, pp_size, dataset): + class _PG: + def rank(self): + return pp_rank + + grid = SimpleNamespace(dim_names=("pp",), shape=[pp_size], get_pg=lambda dims: _PG()) + module._rank_grid_and_module = lambda grids: (grid, module_name) + cfg = SimpleNamespace(model=SimpleNamespace(_grids={module_name: grid}), dataset=dataset) + return module._batch_spec_for_rank(cfg) + + +def _dataset(reorder=False, language_cost=0.0, packing=False): + return SimpleNamespace( + megatron_mimo_scalable_dp=reorder, + megatron_mimo_intra_microbatch_reorder=reorder, + megatron_mimo_reorder_language_cost_weight=language_cost, + enable_in_batch_packing=packing, + ) + + +def test_batch_spec_language_non_first_stage_input_ids_only_under_reorder(): + # Upstream ships input_ids to the first language stage only; the reorder cost reads it on every + # stage, so reorder (and only reorder) keeps it on later stages. Labels/loss_mask stay last-stage only. name = "qwen35_vl_mimo_finetune_batch_spec_under_test" module = _load_example_module(name) try: - - class _PG: - def rank(self): - return 1 # stage 1 of 3: neither first nor last - - grid = SimpleNamespace(dim_names=("pp",), shape=[3], get_pg=lambda dims: _PG()) - monkeypatch.setattr(module, "_rank_grid_and_module", lambda grids: (grid, module.MIMO_LANGUAGE_MODULE_KEY)) - cfg = SimpleNamespace(model=SimpleNamespace(_grids={module.MIMO_LANGUAGE_MODULE_KEY: grid})) - spec = module._batch_spec_for_rank(cfg) - assert spec.input_ids is True and spec.position_ids is True - assert spec.labels is False and spec.loss_mask is False - assert spec.modality_inputs is False + lang = module.MIMO_LANGUAGE_MODULE_KEY + off = _spec_for(module, lang, pp_rank=1, pp_size=3, dataset=_dataset(reorder=False)) + assert off.input_ids is False and off.position_ids is True + assert off.labels is False and off.loss_mask is False and off.modality_inputs is False + on = _spec_for(module, lang, pp_rank=1, pp_size=3, dataset=_dataset(reorder=True)) + assert on.input_ids is True and on.position_ids is True + assert on.labels is False and on.loss_mask is False and on.modality_inputs is False + first = _spec_for(module, lang, pp_rank=0, pp_size=3, dataset=_dataset(reorder=False)) + assert first.input_ids is True finally: sys.modules.pop(name, None) -def test_batch_spec_encoder_ships_inputs_and_pixels(monkeypatch): - name = "qwen35_vl_mimo_finetune_encoder_spec_under_test" +def test_batch_spec_attention_mask_follows_packing_or_language_cost(): + name = "qwen35_vl_mimo_finetune_attention_mask_spec_under_test" module = _load_example_module(name) try: + lang = module.MIMO_LANGUAGE_MODULE_KEY + neither = _dataset() + assert _spec_for(module, lang, 0, 1, neither).attention_mask is False + assert _spec_for(module, "images", 0, 1, neither).attention_mask is False + packing = _dataset(packing=True) + assert _spec_for(module, lang, 0, 1, packing).attention_mask is True + assert _spec_for(module, "images", 0, 1, packing).attention_mask is False + cost = _dataset(reorder=True, language_cost=0.5) + assert _spec_for(module, lang, 0, 1, cost).attention_mask is True + assert _spec_for(module, "images", 0, 1, cost).attention_mask is True + cost_without_reorder = _dataset(reorder=False, language_cost=0.5) + assert _spec_for(module, "images", 0, 1, cost_without_reorder).attention_mask is False + finally: + sys.modules.pop(name, None) - class _PG: - def rank(self): - return 0 - grid = SimpleNamespace(dim_names=("pp",), shape=[1], get_pg=lambda dims: _PG()) - monkeypatch.setattr(module, "_rank_grid_and_module", lambda grids: (grid, "images")) - cfg = SimpleNamespace(model=SimpleNamespace(_grids={"images": grid})) - spec = module._batch_spec_for_rank(cfg) - assert spec.input_ids is True and spec.modality_inputs is True - assert spec.labels is False and spec.loss_mask is False and spec.position_ids is False +def test_batch_spec_encoder_ships_inputs_and_pixels(): + name = "qwen35_vl_mimo_finetune_encoder_spec_under_test" + module = _load_example_module(name) + try: + first = _spec_for(module, "images", pp_rank=0, pp_size=1, dataset=_dataset()) + assert first.input_ids is True and first.modality_inputs is True + assert first.labels is False and first.loss_mask is False and first.position_ids is False + later_off = _spec_for(module, "images", pp_rank=1, pp_size=2, dataset=_dataset(reorder=False)) + assert later_off.input_ids is False and later_off.modality_inputs is False + later_on = _spec_for(module, "images", pp_rank=1, pp_size=2, dataset=_dataset(reorder=True)) + assert later_on.input_ids is True and later_on.modality_inputs is False finally: sys.modules.pop(name, None) From 329a5e4f47e804805145ce1dabf0b36412969df6 Mon Sep 17 00:00:00 2001 From: "kayeon.song" Date: Sat, 3 Oct 2026 05:16:47 +0000 Subject: [PATCH 09/13] fix(training): read cuda_graph_warmup_steps with getattr in the memory report MegatronMIMOProvider is not a TransformerConfig and has no cuda_graph_warmup_steps; with optimizer_cuda_graph=True the memory report added in #6003 raised AttributeError on the first logged iteration. Same treatment as the cuda_graph_impl read in the sibling line. Signed-off-by: kayeon.song Signed-off-by: Yoonsik Kim Signed-off-by: Chanwoo Park Co-authored-by: Yoonsik Kim Co-authored-by: Chanwoo Park <45866990+sailor1493@users.noreply.github.com> --- src/megatron/bridge/training/utils/train_utils.py | 2 +- tests/unit_tests/training/utils/test_train_utils.py | 8 +++++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/training/utils/train_utils.py b/src/megatron/bridge/training/utils/train_utils.py index db6c96e337e..5097d23ed91 100644 --- a/src/megatron/bridge/training/utils/train_utils.py +++ b/src/megatron/bridge/training/utils/train_utils.py @@ -1224,7 +1224,7 @@ def training_log( # after that step, so warmup_steps + 1 is the post-capture iteration. memory_reporting_iterations = max( memory_reporting_iterations, - config.model.cuda_graph_warmup_steps + 1, + getattr(config.model, "cuda_graph_warmup_steps", 0) + 1, ) if iteration >= loaded_iteration + memory_reporting_iterations: # Always include optimizer state memory and, when enabled, CUDA graph diff --git a/tests/unit_tests/training/utils/test_train_utils.py b/tests/unit_tests/training/utils/test_train_utils.py index 5b8ec41d63c..8485d52a687 100644 --- a/tests/unit_tests/training/utils/test_train_utils.py +++ b/tests/unit_tests/training/utils/test_train_utils.py @@ -974,6 +974,8 @@ def test_memory_reporting( pytest.param("none", True, None, 4, False, id="optimizer-after-capture"), pytest.param("none", False, "transformer_engine", 2, True, id="vision-before-capture"), pytest.param("none", False, "transformer_engine", 4, False, id="vision-after-capture"), + # MegatronMIMOProvider has no cuda_graph_* fields; optimizer graphs must still report. + pytest.param(None, True, None, 2, False, id="mimo-provider-optimizer-graph"), ], ) @mock.patch("megatron.bridge.training.utils.train_utils.get_num_microbatches") @@ -1009,7 +1011,11 @@ def test_memory_reporting_cutoff( mock_get_world_size.return_value = 32 mock_get_rank.return_value = 0 - mock_config.model.cuda_graph_impl = cuda_graph_impl + if cuda_graph_impl is None: + del mock_config.model.cuda_graph_impl + del mock_config.model.cuda_graph_warmup_steps + else: + mock_config.model.cuda_graph_impl = cuda_graph_impl mock_config.model.vision_cuda_graph_impl = vision_cuda_graph_impl mock_config.optimizer.optimizer_cuda_graph = optimizer_cuda_graph mock_global_state.train_state.step = iteration From 94bc3e9a3e925a4dc990c6b23a0b54af4854f56e Mon Sep 17 00:00:00 2001 From: "kayeon.song" Date: Sat, 3 Oct 2026 05:35:12 +0000 Subject: [PATCH 10/13] refactor(training): drop the reorder shutdown hook from the MegatronMIMO train loop The MIMO loop ends only when train_iters is reached or the iterator is exhausted; at that point the ReorderingBuffer worker has already exited on its own (verified on 2 ranks), so the hook never did any work. Dropped iterators are handled by ReorderingBuffer.__del__. The train loop is now unchanged from upstream. Signed-off-by: kayeon.song Signed-off-by: Yoonsik Kim Signed-off-by: Chanwoo Park Co-authored-by: Yoonsik Kim Co-authored-by: Chanwoo Park <45866990+sailor1493@users.noreply.github.com> --- src/megatron/bridge/training/train_megatron_mimo.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/src/megatron/bridge/training/train_megatron_mimo.py b/src/megatron/bridge/training/train_megatron_mimo.py index 9546b7bd10c..600741e5bd8 100644 --- a/src/megatron/bridge/training/train_megatron_mimo.py +++ b/src/megatron/bridge/training/train_megatron_mimo.py @@ -476,9 +476,4 @@ def train_megatron_mimo( timers("interval-time").stop() - # Stop the reorder prefetch thread when the train iterator is a ReorderingBuffer; a plain - # iterator has no shutdown() and is left alone. - if hasattr(train_data_iterator, "shutdown"): - train_data_iterator.shutdown() - logger.info(f"Rank {dist.get_rank()}: MegatronMIMO training completed") From d2d23acd6510dd37e45f3b7e49921e192f60d362 Mon Sep 17 00:00:00 2001 From: "kayeon.song" Date: Sat, 3 Oct 2026 05:40:09 +0000 Subject: [PATCH 11/13] docs(training): add metric expectations and validation patterns to the MIMO reorder page MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - State the Megatron Core limit that the encoder data-parallel size must not exceed the language one, and that a dp-1 module's local permutation runs synchronously. - Add Expected Metric Changes and Representative Validation Patterns in the form used by the other training pages, from the 8×A100 Qwen3.5-0.8B runs. - Mark tensor parallelism as validated; list the page in the docs toctree. Signed-off-by: kayeon.song Signed-off-by: Yoonsik Kim Signed-off-by: Chanwoo Park Co-authored-by: Yoonsik Kim Co-authored-by: Chanwoo Park <45866990+sailor1493@users.noreply.github.com> --- docs/index.md | 1 + .../training/mimo-intra-microbatch-reorder.md | 39 +++++++++++++++++-- 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/docs/index.md b/docs/index.md index 2941a3aa157..5e58ca78ab8 100644 --- a/docs/index.md +++ b/docs/index.md @@ -83,6 +83,7 @@ training/cpu-offloading.md training/moe-optimization.md training/peft.md training/packed-sequences.md +training/mimo-intra-microbatch-reorder.md training/multi-token-prediction.md training/dspark-speculative-decoding.md training/callbacks.md diff --git a/docs/training/mimo-intra-microbatch-reorder.md b/docs/training/mimo-intra-microbatch-reorder.md index d346f690140..970af6e7fc1 100644 --- a/docs/training/mimo-intra-microbatch-reorder.md +++ b/docs/training/mimo-intra-microbatch-reorder.md @@ -39,7 +39,9 @@ fan-out stays aligned. Balancing is done over the canonical grid (the least common multiple of the module data-parallel sizes), so modules with different data-parallel sizes, including a module with data-parallel size 1, agree on the same order. A data-parallel-1 module has nothing to exchange and simply -permutes its whole micro-batch locally. +permutes its whole micro-batch locally; that permutation runs synchronously on +the training thread (about 14 ms per micro-batch of 4 samples on A100) because +there is no exchange to overlap. By default the exchange for the next window runs on a background thread with a dedicated NCCL process group, overlapping with the current window's compute. @@ -99,20 +101,51 @@ with the variable unset. | Configuration | Status | |---|---| -| Heterogeneous module data parallelism (for example vision dp 1 or 2 with language dp 2 or 4) | Supported. Balancing uses the canonical grid; a dp-1 module reorders locally. | +| Heterogeneous module data parallelism (for example vision dp 1 or 2 with language dp 2 or 4) | Supported. Balancing uses the canonical grid; a dp-1 module reorders locally. The encoder data-parallel size must not exceed the language data-parallel size: Megatron Core's bridge fan-in with variable per-sample image-token counts is not implemented and asserts on the first forward step. | | Language pipeline parallelism > 1 | Supported. Every language stage receives `input_ids` so all stages compute the same assignment. | | Encoder pipeline parallelism > 1 | Not supported; the example raises before training starts. | | `--pad-to-seq-length false` (variable batch width) | Not supported with reordering; the example raises before training starts. Exchanged samples are concatenated back into one micro-batch, so all ranks must pad to the same length. | | `micro_batch_size` not divisible by the canonical grid size | Rejected by the canonical loader. | | `drop_last=False`, or a dataloader type other than `single` / `cyclic` | Rejected by scalable data parallelism. | | Expert tensor parallelism > 1 | Rejected by scalable data parallelism. | -| Tensor parallelism > 1, context parallelism > 1 | Untested. | +| Tensor parallelism > 1 | Validated (language tp 2 with and without vision tp 2). | +| Context parallelism > 1 | Untested. | A real exchange needs a module with data-parallel size 2 or more, which exceeds the two-GPU functional-test budget. CI covers the exchange arithmetic with CPU-emulated all-to-all unit tests and a two-GPU smoke test; multi-rank runs are validated manually. +## Expected Metric Changes + +| Metric | Expected Change | Conditions | Evidence | +|---|---|---|---| +| per-rank vision load spread | down | module data-parallel size 2 or more, varied image sizes | measured | +| `step_time` | flat to slightly down | every module has data-parallel size 2 or more, overlap on | measured | +| `step_time` | slightly up | a module has data-parallel size 1 (its local permutation is synchronous) | measured | +| `step_time` | up | overlap off (`--no-overlap-intra-microbatch-reorder`) | measured | +| `peak_memory` | up by up to two extra windows of micro-batches per data rank | overlap on; grows with `megatron_mimo_reorder_window_size` | expected | + +Balancing cannot remove the imbalance entirely: every rank keeps the same +number of samples, so with few samples per rank and a bimodal image-size +distribution the reachable spread is bounded by the largest single image. + +## Representative Validation Patterns + +On one 8×A100 node with Qwen3.5-0.8B on cord-v2 (sequence length 2048, +bf16), language dp 4 with vision dp 4 at micro-batch 8 cut the mean +max-minus-min image-token spread across vision ranks from `643` to `421` and +lowered iteration time by about 2% with overlap on. A 3-rank run with vision +dp 1 was about 5% slower with overlap and 10% slower without, because the dp-1 +module's permutation is synchronous. Across these runs the per-micro-batch +sum of image tokens over all vision ranks matched the unreordered run exactly, +iteration-1 and iteration-2 losses matched to all printed digits, and +checkpoint resume reproduced the continuous run. The same geometry completed +3000 iterations with overlap and 10000 iterations without overlap under +`CUDA_DEVICE_MAX_CONNECTIONS=1`; language tp 2, pp 2, and tp 2 with vision tp 2 +each completed 1000 iterations. Treat these as patterns to reproduce on the +target model and data, not as fixed gains. + ## Feature Interactions - **In-batch sequence packing** (`enable_in_batch_packing`): reordering runs in From da5b14ee6284bdcadb3fb0f804b21fb4ef6cf945 Mon Sep 17 00:00:00 2001 From: "kayeon.song" Date: Sat, 3 Oct 2026 05:49:25 +0000 Subject: [PATCH 12/13] test(data): cover real_token_lengths and the overlapped reorder exchange - Unit: real_token_lengths counts mask ones per sample, keeps pad-valued real tokens, and rejects a missing or mismatched attention_mask. - Functional (2 GPUs): ReorderingBuffer with overlap=True exchanges three windows on its side NCCL group, serves the balanced shards in order, and its prefetch thread exits on its own once the source iterator is exhausted. Signed-off-by: kayeon.song Signed-off-by: Yoonsik Kim Signed-off-by: Chanwoo Park Co-authored-by: Yoonsik Kim Co-authored-by: Chanwoo Park <45866990+sailor1493@users.noreply.github.com> --- .../megatron_mimo/test_reorder_exchange.py | 50 +++++++++++++++++++ .../data/megatron_mimo/test_dp_utils.py | 20 ++++++++ 2 files changed, 70 insertions(+) diff --git a/tests/functional_tests/test_groups/megatron_mimo/test_reorder_exchange.py b/tests/functional_tests/test_groups/megatron_mimo/test_reorder_exchange.py index 0f090785707..13fcb369524 100644 --- a/tests/functional_tests/test_groups/megatron_mimo/test_reorder_exchange.py +++ b/tests/functional_tests/test_groups/megatron_mimo/test_reorder_exchange.py @@ -27,6 +27,7 @@ import torch.distributed as dist from megatron.bridge.data.megatron_mimo.reorder_buffer import ( + ReorderingBuffer, balanced_assignment, build_module_dp_process_groups, exchange_window, @@ -151,3 +152,52 @@ def test_reorder_exchange_dp2_variable_images(): image_count_of = lambda b: (b["input_ids"] == _VISION_START).sum(dim=1).to(torch.long) # noqa: E731 _run_exchange_case(global_samples, image_count_of) dist.barrier() + + +@pytest.mark.run_only_on("GPU") +def test_reorder_buffer_overlap_dp2_exits_on_exhaustion(): + """Overlap path: the prefetch thread exchanges on its own NCCL group and exits once the source is exhausted.""" + initialize_distributed() + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for this functional test") + if dist.get_world_size() != 2: + pytest.skip("This functional test requires exactly 2 ranks") + + dp_size, n_groups, n_windows, window_size = 2, 2, 3, 2 + rank = dist.get_rank() + # Six micro-batches; micro-batch m holds globals [4m + 2*rank, 4m + 2*rank + 1] with costs [4,3,1,2] (a swap every time). + global_samples = [[4], [3], [1], [2]] + main_pg = dist.new_group(ranks=list(range(dp_size)), backend="nccl") + dp_rank, ds, gloo, nccl = build_module_dp_process_groups(main_pg, overlap=True) + assert nccl is not main_pg # overlap=True creates a side NCCL group + + def source(): + for m in range(n_windows * window_size): + yield _make_shard([(4 * m + g, global_samples[g]) for g in range(rank * 2, rank * 2 + 2)]) + + image_count_of = lambda b: (b["input_ids"] == _VISION_START).sum(dim=1).to(torch.long) # noqa: E731 + buf = ReorderingBuffer( + source(), + dp_rank=dp_rank, + dp_size=ds, + n_groups=n_groups, + cost_of=functools.partial(sample_cost, encoder_cost_weight=1.0), + dp_group_gloo=gloo, + dp_group_nccl=nccl, + overlap=True, + image_count_of=image_count_of, + window_size=window_size, + ) + assert buf._thread is not None and buf._thread.is_alive() + + expected = _expected_owned(global_samples, dp_rank, dp_size, n_groups) + served = 0 + for m, out in enumerate(buf): + got = [int((int(row[0].item()) - 1) // 100) for row in out["input_ids"].cpu()] + assert got == [4 * m + g for g in expected], (rank, m, got) + served += 1 + assert served == n_windows * window_size + + buf._thread.join(timeout=5.0) + assert not buf._thread.is_alive(), "worker must exit on its own once the source iterator is exhausted" + dist.barrier() diff --git a/tests/unit_tests/data/megatron_mimo/test_dp_utils.py b/tests/unit_tests/data/megatron_mimo/test_dp_utils.py index 8ab1609b7d3..c401a86f2ad 100644 --- a/tests/unit_tests/data/megatron_mimo/test_dp_utils.py +++ b/tests/unit_tests/data/megatron_mimo/test_dp_utils.py @@ -8,6 +8,7 @@ from megatron.bridge.data.megatron_mimo.dp_utils import ( get_megatron_mimo_dp_info, get_megatron_mimo_sampling_info, + real_token_lengths, slice_batch_for_megatron_mimo, ) from megatron.bridge.models.megatron_mimo.megatron_mimo_config import ( @@ -425,3 +426,22 @@ def test_extra_keys_passed_through(self): sliced = slice_batch_for_megatron_mimo(batch, dp_rank=0, dp_size=2) enc = sliced["modality_inputs"]["images"]["vision_encoder"] assert enc["encoder_meta"] == "fixed-string" + + +class TestRealTokenLengths: + def test_counts_mask_ones_per_sample(self): + input_ids = torch.tensor([[5, 6, 0, 0], [5, 6, 7, 8]]) + mask = torch.tensor([[1, 1, 0, 0], [1, 1, 1, 1]]) + assert real_token_lengths(input_ids, attention_mask=mask).tolist() == [2, 4] + + def test_pad_valued_real_token_still_counts(self): + # pad_token == eos_token tokenizers: a real token equal to the pad id must not be dropped. + input_ids = torch.tensor([[5, 0, 0, 0]]) + mask = torch.tensor([[1, 1, 0, 0]]) + assert real_token_lengths(input_ids, attention_mask=mask).tolist() == [2] + + @pytest.mark.parametrize("mask", [None, torch.ones(3, 4), torch.ones(4), torch.ones(2, 3)]) + def test_rejects_missing_or_mismatched_mask(self, mask): + input_ids = torch.zeros(2, 4, dtype=torch.int64) + with pytest.raises(ValueError, match="attention_mask"): + real_token_lengths(input_ids, attention_mask=mask) From c55ce2c9c59fb4a5298b217d2368cafbb704d733 Mon Sep 17 00:00:00 2001 From: "kayeon.song" Date: Sat, 3 Oct 2026 07:19:41 +0000 Subject: [PATCH 13/13] docs(training): correct test coverage and validation statements on the MIMO reorder page The two-rank functional test does run the real exchange (and now the overlapped prefetch thread); attribute each validation result to the configuration that produced it. Signed-off-by: kayeon.song Signed-off-by: Yoonsik Kim Signed-off-by: Chanwoo Park Co-authored-by: Yoonsik Kim Co-authored-by: Chanwoo Park <45866990+sailor1493@users.noreply.github.com> --- .../training/mimo-intra-microbatch-reorder.md | 22 ++++++++++--------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/docs/training/mimo-intra-microbatch-reorder.md b/docs/training/mimo-intra-microbatch-reorder.md index 970af6e7fc1..5d2c8377a16 100644 --- a/docs/training/mimo-intra-microbatch-reorder.md +++ b/docs/training/mimo-intra-microbatch-reorder.md @@ -111,10 +111,11 @@ with the variable unset. | Tensor parallelism > 1 | Validated (language tp 2 with and without vision tp 2). | | Context parallelism > 1 | Untested. | -A real exchange needs a module with data-parallel size 2 or more, which -exceeds the two-GPU functional-test budget. CI covers the exchange arithmetic -with CPU-emulated all-to-all unit tests and a two-GPU smoke test; multi-rank -runs are validated manually. +A full MIMO run with an exchanging module needs at least three GPUs, which +exceeds the two-GPU functional-test budget. CI covers the balancing and +serialization logic with CPU-emulated unit tests, and the real Gloo/NCCL +exchange, including the overlapped prefetch thread, with a two-rank functional +test; full MIMO runs are validated manually. ## Expected Metric Changes @@ -137,12 +138,13 @@ bf16), language dp 4 with vision dp 4 at micro-batch 8 cut the mean max-minus-min image-token spread across vision ranks from `643` to `421` and lowered iteration time by about 2% with overlap on. A 3-rank run with vision dp 1 was about 5% slower with overlap and 10% slower without, because the dp-1 -module's permutation is synchronous. Across these runs the per-micro-batch -sum of image tokens over all vision ranks matched the unreordered run exactly, -iteration-1 and iteration-2 losses matched to all printed digits, and -checkpoint resume reproduced the continuous run. The same geometry completed -3000 iterations with overlap and 10000 iterations without overlap under -`CUDA_DEVICE_MAX_CONNECTIONS=1`; language tp 2, pp 2, and tp 2 with vision tp 2 +module's permutation is synchronous. In the dp 4 run the per-micro-batch sum +of image tokens over all vision ranks matched the unreordered run for all 40 +micro-batches. On language dp 2 + vision dp 2 the reorder off/on pair had +identical iteration-1 and iteration-2 losses, checkpoint resume reproduced the +continuous run, and an overlapped run completed 3000 iterations. Under +`CUDA_DEVICE_MAX_CONNECTIONS=1` without overlap, language tp 2 dp 2 + vision +dp 2 completed 10000 iterations; language tp 2, pp 2, and tp 2 with vision tp 2 each completed 1000 iterations. Treat these as patterns to reproduce on the target model and data, not as fixed gains.