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/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..5d2c8377a16 --- /dev/null +++ b/docs/training/mimo-intra-microbatch-reorder.md @@ -0,0 +1,168 @@ +# 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; 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. + +## 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 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 + +| 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. 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 | Validated (language tp 2 with and without vision tp 2). | +| Context parallelism > 1 | Untested. | + +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 + +| 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. 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. + +## 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` diff --git a/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py b/examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py index 338ca008225..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 - (when packing) 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: @@ -29,6 +29,7 @@ from __future__ import annotations import argparse +import functools import logging import math import os @@ -58,8 +59,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 @@ -284,26 +285,34 @@ def _batch_spec_for_rank(cfg: Any) -> MIMOBatchSpec: 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=is_first_pp, + 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, + attention_mask=packing_active or language_cost_active, ) return MIMOBatchSpec( # Encoder first stages need input_ids to attach per-sample split metadata. - input_ids=is_first_pp, + input_ids=is_first_pp or reorder_active, position_ids=False, labels=False, loss_mask=False, modality_inputs=is_first_pp, + attention_mask=language_cost_active, ) @@ -340,6 +349,22 @@ 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 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(): dp = parallelism.data_parallel_size @@ -418,6 +443,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 +983,55 @@ 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 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)) + 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." + ) + 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_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: + 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 +1343,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 per image patch.", + ) + parser.add_argument( + "--reorder-language-cost-weight", + type=float, + default=0.0, + help="Reorder cost weight per real 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..cb7c81b9568 100644 --- a/src/megatron/bridge/data/builders/direct_hf_sft.py +++ b/src/megatron/bridge/data/builders/direct_hf_sft.py @@ -94,6 +94,12 @@ 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 * 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 + megatron_mimo_reorder_window_size: int = 1 def validate(self) -> None: """Validate declarative source and dataset settings.""" @@ -125,6 +131,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..f1d33717bcb --- /dev/null +++ b/src/megatron/bridge/data/megatron_mimo/reorder_buffer.py @@ -0,0 +1,1286 @@ +# 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. + +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 +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__) + +# 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: + """Return the zero-padding bytes that round ``nbytes`` up to the serialization alignment.""" + return (-nbytes) % _SERIALIZED_TENSOR_BYTE_ALIGNMENT + + +# 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 { + "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 name from sample metadata.""" + 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. + + 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; must divide ``len(costs)``. + + Returns: + A permutation of ``range(len(costs))`` laid out group-by-group. + + 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)): + 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]]: + """Map each sample of one micro-batch to its owner rank and canonical position. + + ``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 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: + ``(owner_rank, canonical_pos)`` per global sample index. + + Raises: + 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: + 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) + 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 assignment to the ``(dst_slot, owner_rank, canonical_pos)`` window route. + + Intra routing never moves a sample across slots, so ``dst_slot == src_slot`` for every sample. + + Args: + assignment: Output of :func:`balanced_assignment` for one slot. + src_slot: Slot the samples were read from. + + Returns: + ``(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]]: + """Build the slot-major window route by balancing each slot independently. + + The window global index of slot ``s``, slot-local position ``j`` is ``g = s * B_global + j``. + + Args: + 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: + ``(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)) + 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: + """Raise if any sample's ``dst_slot`` differs from its source slot ``g // b_global``. + + Args: + 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 a 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]]: + """Compute per-slot per-rank cost max/min before and after balancing and the remote-sample count. + + Args: + 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 ``{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): + 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]: + """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]: + """Build per-field ``{shape, dtype}`` / ``{non_tensor: value}`` / ``None`` metadata for one sample. + + Args: + flat: One flattened sample dict. + + Returns: + Metadata keyed like ``flat``, sufficient to size and decode the sample without its data. + + Raises: + 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) # 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} + 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 aligned ``uint8`` buffer in ``keys`` order. + + Args: + flat: Flattened sample dict; ``None`` values are skipped. + keys: Sorted tensor keys from :func:`sample_keys`. + + Returns: + 1D ``uint8`` tensor with each field padded to the serialization alignment. + """ + parts = [] + for key in keys: + t = flat.get(key) + if t is None: + continue + raw = t.contiguous().view(torch.uint8).reshape(-1) + pad = _alignment_padding_bytes(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: + """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) + if info is None: + continue + _, elem = _dtype_info(info["dtype"]) + nbytes = math.prod(info["shape"]) * elem + nbytes += _alignment_padding_bytes(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]: + """Reconstruct one sample from a ``uint8`` buffer (inverse of :func:`serialize_sample`). + + Args: + 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 rebuilt sample and the byte offset just past it in ``buf``. + """ + 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 + _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 + 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]: + """Build the kept-local partition, packed send buffer, and splits for one ``all_to_all_single``. + + 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 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; every ``dst_slot`` must be in ``[0, window_size)``. + + Returns: + 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)``. + """ + location: Dict[int, Tuple[int, int]] = {} + for r, gidxs in enumerate(all_global_indices): + for i, g in enumerate(gidxs): + location[g] = (r, i) + + 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: + 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)) + + # 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() + + # 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) + if t is not None: + dev = t.device + break + + 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_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 +# --------------------------------------------------------------------------- + + +def _flatten_nested(d: Dict[str, Any], prefix: str = "") -> Dict[str, Any]: + """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: + flat[path] = val + return flat + + +def _unflatten_nested(flat: Dict[str, Any]) -> Dict[str, Any]: + """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(".") + 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]: + """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: + 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]: + """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. + + 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])) + return empty_hidden_states, empty_grid_thw + + +def patches_per_image(grid_thw: torch.Tensor) -> torch.Tensor: + """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. + + Returns: + ``[n_images]`` patch counts. + """ + 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 ``hidden_states`` (patch dim) and ``grid_thw`` (image dim) by whole image.""" + if not image_perm: + # 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) + # 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) + 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]: + """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: + 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]: + """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(): + 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): + # 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: + 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 with size-1 batch dims. + + Args: + batch: Nested micro-batch. + cu_img: Cumulative per-sample image offsets (length ``B + 1``); ``None`` assumes one image per sample. + + 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)] + + +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`). + + 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.") + nested = [_unflatten_nested(f) for f in flats] + return _merge_nested(nested) + + +def _merge_nested(samples: List[Dict[str, Any]]) -> Dict[str, Any]: + # 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: + 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: + 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 + return out + + +def _batch_to_cuda(batch: Dict[str, Any]) -> Dict[str, Any]: + """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): + 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: + """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): + 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: + """Return ``encoder_cost_weight * patches + language_cost_weight * real_tokens`` for one sample. + + 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 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: + 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 and received samples. + + Args: + 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``. + local: Expected samples per slot on this rank. + + Returns: + ``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"]: + 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)) + + 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): + if len(bucket) != local: + raise RuntimeError(f"MegatronMIMO slot {slot} reassembled {len(bucket)} samples, expected {local}.") + bucket.sort(key=lambda x: x[0]) + # 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}.") + 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]]: + """Run one ragged ``all_to_all_single`` for ``route`` and reassemble ``W`` micro-batches. + + Args: + local_flats: This rank's input samples. + 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: All-gathered per-rank 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``. + 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`` 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 = 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. + + 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`` 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 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 + + window_size = len(batches) + if window_size == 0: + raise ValueError("exchange_window requires at least one micro-batch.") + + # 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]) + 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]}." + ) + + # 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] + + # 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] + 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 + 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] + 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 + + 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]]: + """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), 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. + """ + 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: + """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: + """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): + 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. + + ``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__( + 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 + self._active: List[Dict[str, Any]] = [] + self._cursor = 0 + self._exchanges = 0 + if overlap and dp_size > 1 and os.environ.get("CUDA_DEVICE_MAX_CONNECTIONS") == "1": + # 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). " + "Pass overlap=False (--no-overlap-intra-microbatch-reorder) or unset the variable." + ) + if overlap and self._cuda and dp_size > 1: + 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: + """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]]: + # 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( + 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.""" + 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 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: + 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 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) + return + except queue.Full: + continue + + def _worker(self) -> None: + # 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(): + batches = self._collect_window() + if not batches: + break + with torch.cuda.stream(self._stream): + 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)) + finally: + self._put_worker_item(_SENTINEL) + + def __next__(self) -> Dict[str, Any]: + if self._thread is None: + 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 + 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 and side NCCL process groups for this rank's module DP group. + + ``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: 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)``. + """ + 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/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/megatron_mimo_step.py b/src/megatron/bridge/training/megatron_mimo_step.py index 03173d5426d..5e2622a7a20 100644 --- a/src/megatron/bridge/training/megatron_mimo_step.py +++ b/src/megatron/bridge/training/megatron_mimo_step.py @@ -280,6 +280,8 @@ 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 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. data_batch = { diff --git a/src/megatron/bridge/training/utils/train_utils.py b/src/megatron/bridge/training/utils/train_utils.py index 800a6da97f4..5097d23ed91 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" ) @@ -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/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..13fcb369524 --- /dev/null +++ b/tests/functional_tests/test_groups/megatron_mimo/test_reorder_exchange.py @@ -0,0 +1,203 @@ +# 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. + +"""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``).""" + +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 ( + ReorderingBuffer, + balanced_assignment, + build_module_dp_process_groups, + exchange_window, + sample_cost, +) +from tests.functional_tests.utils import initialize_distributed + + +_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: + """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]] = [] + 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 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) + 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 + 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(): + 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]: 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) + dist.barrier() + + +@pytest.mark.run_only_on("GPU") +def test_reorder_exchange_dp2_variable_images(): + 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] 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) + 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/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_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) 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..9ed8ea5d938 --- /dev/null +++ b/tests/unit_tests/data/megatron_mimo/test_reorder_buffer.py @@ -0,0 +1,1109 @@ +# 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 (all-to-all emulated on CPU).""" + +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: + 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), + "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) + 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 + 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(): + # 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) + + +# --------------------------------------------------------------------------- +# Assignment +# --------------------------------------------------------------------------- + + +def test_balanced_assignment_contiguous_blocks_and_balance(): + 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 + 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)) + assert sorted(per_rank[0]) == list(range(0, local)) + assert sorted(per_rank[1]) == list(range(local, n)) + + # 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] + 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)) + 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) + 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): + """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)] + 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) + + 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.""" + 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" + + 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): + """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)} + + 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) + + 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] + assert positions == list(range(d * local, (d + 1) * local)), (d, positions) + + 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(): + 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): + b = len(patches_per_sample) + s = 6 + # 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 + 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 + ) + 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 + 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 +# --------------------------------------------------------------------------- + +_VKEY = "modality_inputs.images.enc" + + +def _flat_vision_sample(idx: int, n_tokens: int, n_patches: int, *, none_vision: bool = False) -> dict: + """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: + 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 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 + 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) + meta = tensor_metadata(flat) + keys = sample_keys(meta) + buf = serialize_sample(flat, keys) + 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) + 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 +# --------------------------------------------------------------------------- + + +def _make_mimo_batch_var(images_per_sample: list[int]) -> dict: + """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) + 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 + 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_one_image_per_sample_no_cu_img(): + batch = _make_mimo_batch([2, 5, 1, 3]) + _assert_batch_equal(merge_samples(split_microbatch(batch)), batch) + 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): + 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(): + _run_exchange_nested([2, 0, 1, 0, 3, 1, 0, 2], dp_size=2, n_groups=2) + + +# --------------------------------------------------------------------------- +# 3D routing + per-slot reassembly (reassemble_window) +# --------------------------------------------------------------------------- + + +def _emulate_window_exchange(rank_flats, rank_global_indices, all_tensor_meta, route, *, dp_size, window_size, local): + 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(): + 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 + 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(): + 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), + _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): + # 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. 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] + + 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(): + 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: 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( + 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 +# --------------------------------------------------------------------------- + + +def test_build_window_route_layout(): + 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] + assert all(dst == s for dst, _owner, _pos in slot_entries) + assert sorted(pos for _dst, _owner, pos in slot_entries) == list(range(b_global)) + 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): + """Window exchange over W micro-batches; each slot must equal its independent per-slot balanced shard.""" + window_size = len(images_per_slot) + 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 + ) + + 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(): + _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(): + _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(): + 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(): + dp_size, n_groups = 4, 4 + # 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], + ] + 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 + 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(): + 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) + assert_intra_no_cross_slot(route, b_global=4) + + +def test_assert_intra_no_cross_slot_trips_on_cross_slot_route(): + 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(): + 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; 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, + 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)) + + +def test_reordering_buffer_overlap_single_rank_no_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 + assert [int(b["input_ids"].item()) for b in buf] == list(range(5)) + + +def test_reorder_window_local_matches_paired_dp2_order(): + 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"], + "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) + + 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(): + 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, + ) + + +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 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()) + + +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()) + + +# --------------------------------------------------------------------------- +# 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]) + samples = split_microbatch(batch) + # 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 + 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)} + 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 + + 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 + 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 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), + } + 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 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), + } + 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)} + 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 + # 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 + 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(): + """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) + vision_flat = { + "input_ids": input_ids, + f"{_VKEY}.grid_thw": torch.tensor([[1, 4, 2]], dtype=torch.int64), # prod = 8 patches + } + 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 + # 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 + + +# --------------------------------------------------------------------------- +# Single D2H on the vision reorder path +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestReorderVisionSingleD2H: + def test_reorder_vision_single_d2h(self, monkeypatch): + # The cumsum must be a single .tolist(): no per-image .item() D2H sync in the loop + 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]) + assert calls["item"] == 0 + 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 positions [r*local, (r+1)*local) of one order shared by all modules.""" + module_dps = [1, 2] + grid = canonical_grid_size(module_dps) + 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..604ae17a0fa 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,113 @@ 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 _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: + 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_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) + + +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) + + +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) 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") 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