feat(data): intra-microbatch reordering for MegatronMIMO scalable reads

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

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 |

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`

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)