feat(data): scalable canonical-grid reads for MegatronMIMO data loading

[going-song]

What

Adds opt-in scalable data-parallel reads to MegatronMIMO. By default every data-loading
rank reads the full global micro-batch and forward_step slices out its module-DP shard, so
per-rank read and preprocessing cost scales with the DP degree. With
dataset.megatron_mimo_scalable_dp each rank's sampler reads only its shard and the slice is
skipped:

default            scalable
read    [4, S]     read    [2, S]      (per language rank, dp=2)
slice   [4]->[2]   slice   skipped
forward [2, S]     forward [2, S]      (same shard geometry, identical routing)

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

Why a canonical grid (mechanism deviation from #4608)

#4608 sharded each module's loader by its own module-DP (rank, size). A sampler's shard
assignment is a function of (micro_batch_size, data_parallel_size), so under "cyclic" —
the direct-HF-SFT default since #5048 — modules with different DP sizes read different
sample sets, while the BridgeCommunicator routes modality embeddings positionally along
the batch dim. We hit this live as an embedding-count crash in
align_embeddings_by_token_positions; with uniform image sizes it would silently train on
mispaired image-text.

Here all loaders shard on one canonical grid — the LCM of the module DP sizes. Every
group sampler uses the identical (micro_batch // grid, grid) geometry and global
consumed_samples; a rank whose module has DP size d concatenates the windows of its
grid // d consecutive groups (CanonicalGroupBatchSampler). Group streams are
module-independent, so every module materializes the identical ordered micro-batch under
"single" and "cyclic"; for "single" the stream is bit-identical to the default
path's slice. Geometries are truncated to whole micro-batches so ragged dataset sizes cannot
desynchronize the groups. ("Canonical group" is the vocabulary of #4608's reordering
machinery, which ③ will port.)

Configuration

cfg.dataset.megatron_mimo_scalable_dp = True   # DirectHFSFTDatasetConfig; example: --scalable-dp

"single" and "cyclic" are supported; "batch"/"external" and drop_last=False are
rejected in validate() and again at loader build. Off by default — with the switch unset, the
sampler parameters, loader arguments and slicing behavior are identical to main.

Notes for reviewers

• Scalable reads reject expert_tensor_parallel_size > 1 (unvalidated combination). The
  default slicing path shares the same assumption — happy to hoist the guard into
  MegatronMIMOParallelismConfig validation if preferred.
• No train-loop changes: the provider forces variable_seq_lengths=True, so the schedule's
  micro_batch_size is shape-inert and the LR/consumed accounting already uses the full
  micro-batch in both modes (feat(data): intra-microbatch reordering for MegatronMIMO (+ sequence packing, scalable DP) #4608's schedule/LR adjustment is unnecessary here).
• Under "cyclic", scalable on/off are two different valid epoch shuffles (per-iteration
  losses not comparable); "single" is order-identical.

Testing

258 unit tests passing locally, including the cross-geometry alignment invariant (a rank
covering groups {0,1} reads exactly the concatenation of the single-group ranks' windows, for
"cyclic" in both data_sharding modes) and a library-path end-to-end test with DP {2,3}
(grid = 6). Live PP=2 runs (5 GPUs of an 8xA100 node, language=tp1,pp2,dp2 + images=dp1, Qwen3.5-0.8B random
init, cord_v2, mbs 4, 10 iters): cyclic+scalable completes 10/10 where module-DP sharding
crashes; single scalable matches the default losses within bf16 noise (max |diff| 0.0043);
the default path matches prior losses within run-to-run noise (max |diff| 0.0024); packing on/off parity under scalable ~0.001.
(One drive-by repair: the example-config test's arg namespace was already failing on main —
#5131 added pack_sequences_in_batch to _build_dataset_config without updating it — so this
PR adds both missing kwargs there.)

Known limitations

ETP > 1 and "batch"/"external" loaders are rejected for scalable reads; MegatronMIMO still
cannot run CP (pre-existing); checkpoint resume is unit-tested but not yet exercised in a live
run, and multi-node validation is pending, as with #4608.

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]

The default slicing path relies on the same assumption — should this live in MegatronMIMOParallelismConfig validation instead?

[going-song]

Prevents unequal per-group window counts on ragged dataset sizes (cyclic + data_sharding=False); regression-tested.

[going-song]

This test has been failing on main since #5131 (missing pack_sequences_in_batch); adding both kwargs.

[liding-nv]

/ok to test 958575f