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

[going-song]

What

Adds opt-in intra-microbatch reordering to MegatronMIMO scalable data-parallel reads
(#5525). With scalable reads each data rank holds a fixed 1/dp shard of every micro-batch,
so the per-rank vision load is whatever the sampler assigned. Image sizes vary widely, and the
slowest vision rank gates every step. Reordering rebalances each micro-batch across a module's
DP group before the forward step:

                      vision rank 0        vision rank 1
read (scalable)       [big, big]           [small, small]      straggler: rank 0
reorder               [big, small]         [big, small]        one ragged all_to_all

1. Every rank computes a per-sample cost from fields identical on both modules: the
   image-placeholder token count in input_ids (= vision patches / spatial_merge_size²),
   optionally plus the real token count from attention_mask.
2. Costs are all-gathered over the module DP group (Gloo) and each rank runs the same
   deterministic equal-count LPT assignment over the canonical grid (LCM of the module DP
   sizes), so the vision and language modules derive one identical assignment without
   talking to each other; a dp-1 module permutes locally.
3. Samples move with one ragged all_to_all_single (NCCL) and each rank reassembles its
   micro-batch in canonical order, so the BridgeCommunicator positional fan-out stays
   aligned.

By default the next window's exchange runs on a background thread with its own NCCL group,
overlapping the current window's compute. megatron_mimo_reorder_window_size balances W
micro-batches in one exchange.

Third of three PRs splitting #4608 (original work by Chanwoo Park @sailor1493, credited as
co-author and signer), per the maintainers' split request in #4609. Parts 1 and 2 merged as
#5131 and #5525.

Differences from #4608

• Balancing runs on the canonical grid of feat(data): scalable canonical-grid reads for MegatronMIMO data loading #5525 instead of per-module DP, so modules with
  different DP sizes (including dp 1) agree on one assignment.
• attention_mask is also shipped to every module when the language cost term is on;
  input_ids follows feat(data): intra-microbatch reordering for MegatronMIMO (+ sequence packing, scalable DP) #4608 (every language PP stage when reordering). Both are decided in
  _batch_spec_for_rank; otherwise the per-rank field selection of improve Qwen MegatronMIMO data loading #4442 is unchanged.
• Overlap is refused under CUDA_DEVICE_MAX_CONNECTIONS=1. feat(data): intra-microbatch reordering for MegatronMIMO (+ sequence packing, scalable DP) #4608 only warned ("correct but no
  faster"); with a single hardware queue the worker's all-to-all and the DDP all-reduce execute
  in issue order, which differs per rank, and training deadlocks (reproduced at iteration 290
  on 4 ranks and iteration 33 on 6 ranks with TP 2); the non-overlapped path is unaffected.
• --pad-to-seq-length false is rejected with reordering (ranks must pad to one width to
  concatenate exchanged samples).
• Cost weights renamed encoder_cost_weight / language_cost_weight; negative weights rejected.
• No training-loop change: the worker exits on its own when the data iterator is exhausted.

Configuration

cfg.dataset.megatron_mimo_scalable_dp = True                 # required
cfg.dataset.megatron_mimo_intra_microbatch_reorder = True
cfg.dataset.megatron_mimo_reorder_encoder_cost_weight = 1.0  # per vision patch
cfg.dataset.megatron_mimo_reorder_language_cost_weight = 0.0 # per real token; 0 = vision only
cfg.dataset.megatron_mimo_reorder_overlap = True
cfg.dataset.megatron_mimo_reorder_window_size = 1

All on DirectHFSFTDatasetConfig, off by default; the Qwen3.5-VL example exposes them as
--intra-microbatch-reorder, --reorder-encoder-cost-weight, --reorder-language-cost-weight,
--no-overlap-intra-microbatch-reorder, --reorder-window-size. The library loader does not
wire the buffer (the cost needs model token ids); the example does.

Notes for reviewers

• The equal-count constraint bounds the reachable balance: with two samples per rank and a
  bimodal image-size distribution the measured spread (max−min 643 → 421 image tokens, 8 ranks)
  matches the best reachable value.
• Up to three windows are resident per data rank with overlap (consuming, queued, exchanging).
• The dp-1 module's local permutation runs synchronously (about 14 ms per micro-batch of 4 on
  A100); a 3-rank run with vision dp 1 is about 5% slower with reordering, an 8-rank run with
  vision dp 4 about 2% faster.

Testing

• Unit (CPU, all-to-all emulated): balanced assignment and canonical-grid pairing for module
  DP {(1,2), (2,4), (2,3)}, serialization round-trip for every supported dtype and MRoPE
  [3,B,S] position ids, per-slot reassembly for windows of 1–3, cost module-independence
  (vision-style vs language-style sample), the CUDA_DEVICE_MAX_CONNECTIONS=1 refusal, the
  example's batch-spec and validate paths.
• Functional (2 GPUs): real Gloo all-gather + NCCL ragged all-to-all over a dp-2 group with
  single- and variable-image samples, vision rows travelling with their sample byte-for-byte;
  and ReorderingBuffer(overlap=True, window_size=2) serving three windows on its side NCCL
  group, with the prefetch thread exiting on its own once the source is exhausted.
• Live, 8×A100, Qwen3.5-0.8B on cord-v2, bf16: reorder off/on pairs with identical
  iteration-1/2 losses (language dp 2 + vision dp 2, 30 iterations; pretrained weights
  iteration-1 difference 1.3e-4, so image–text pairing is intact); per-micro-batch image-token
  sums over all vision ranks identical to the unreordered run (dp 4 + dp 4, 40 micro-batches);
  checkpoint resume reproduces the continuous run; language pp 2, tp 2, tp 2 + vision tp 2 and
  pp 2 + dp 2 each 1000 iterations; window {1,2,4} × overlap on/off × language cost {0, 0.5}
  13 × 1000 iterations; 3000 iterations with overlap; 10000 iterations without overlap under
  CUDA_DEVICE_MAX_CONNECTIONS=1.

Known limitations

Encoder PP > 1 is rejected with reordering; context parallelism is untested; encoder dp must
not exceed language dp, because Megatron Core's bridge fan-in with variable image-token counts
asserts (MimoModel._attach_modality_split_sizes) — this is independent of reordering.

Also included

Two small fixes unrelated to reordering: Qwen3VLGPTModel.forward accepts the
mtp_input_mask kwarg that MimoModel passes since NVIDIA/Megatron-LM#6770 (MIMO training
fails on the first step without it), and the one-shot memory report from #6003 reads
cuda_graph_impl / cuda_graph_warmup_steps with getattr because MegatronMIMOProvider is
not a TransformerConfig (fails after iteration 1, and with optimizer_cuda_graph=True).

Signed-off-by: kayeon.song kayeon.song@navercorp.com
Signed-off-by: Yoonsik Kim yoonsik.kim90@navercorp.com
Signed-off-by: Chanwoo Park chanwoo.park98@navercorp.com

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[going-song]

Was a warning in #4608; reproduced as a hang (iteration 290 on 4 ranks, 33 on 6 ranks with TP 2).

[going-song]

attention_mask now reaches non-packing stages when the language cost term is on; the model must not receive it.

[going-song]

Unrelated to reordering: MegatronMIMOProvider has no cuda_graph_* fields, so MIMO crashed here after iteration 1.

[going-song]

Unrelated to reordering: MimoModel passes this kwarg since NVIDIA/Megatron-LM#6770; MIMO failed on the first step without it.

[going-song]

Gated on the canonical grid, not this rank's DP size, so dp-1 ranks also enter the collective process-group setup below.

[cuichenx]

The two compatibility fixes included here correspond to NVBug 6864016: forwarding mtp_input_mask through Qwen3VLGPTModel.forward and handling absent CUDA graph fields in MegatronMIMO training logging. Both failure sites described by the report are covered in the diff. Tracking this as pending QA verification; this comment does not claim a completed release validation.

[cuichenx]

/ok to test c55ce2c

[github-actions]

Automatic Claude reviews have been retired. To request a pull-request review, post a comment containing:

/review

Add model=claude to use a Claude reviewer (the default is model=codex). mode=light|strict selects the review depth; for example, /review model=claude mode=strict. Comment /review help for all options.

[github-actions]

🌿 Preview your docs: https://nvidia-preview-pr-6306-c55ce2c9c59f.docs.buildwithfern.com/nemo/megatron-bridge