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

examples/megatron_mimo/qwen35_vl/finetune_qwen35_vl.py

@@ -58,6 +58,7 @@
     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.samplers import build_pretraining_data_loader
 from megatron.bridge.data.sources.hf import hf_dataset_supports_split
@@ -416,6 +417,7 @@ def _build_dataset_config(args: argparse.Namespace) -> DirectHFSFTDatasetConfig:
         # MegatronMIMO packs in the step, after the module-DP slice (deferred packing).
         enable_in_batch_packing=args.pack_sequences_in_batch,
         defer_in_batch_packing_to_step=True,
+        megatron_mimo_scalable_dp=args.scalable_dp,
         do_validation=do_validation,
         do_test=False,
         trust_remote_code=args.trust_remote_code,
@@ -857,9 +859,11 @@ def _build_data_iterators(cfg, _megatron_mimo_infra, *, train_state=None):
         if cfg.model._grids is None:
             raise ValueError("MegatronMIMOProvider._grids is None. Model must be built before data iterators.")
 
+        scalable_dp = bool(getattr(cfg.dataset, "megatron_mimo_scalable_dp", False))
         sampler_dp_rank, sampler_dp_size, needs_data = get_megatron_mimo_sampling_info(
             cfg.model.megatron_mimo_parallelism_config,
             cfg.model._grids,
+            scalable_dp=scalable_dp,
         )
         if not needs_data:
             return None, None
@@ -909,20 +913,42 @@ def _build_data_iterators(cfg, _megatron_mimo_infra, *, train_state=None):
                 batch_spec=batch_spec,
             )
 
-        train_loader = build_pretraining_data_loader(
-            dataset=train_ds,
-            consumed_samples=train_state.consumed_train_samples,
-            dataloader_type=cfg.dataset.dataloader_type,
-            micro_batch_size=cfg.train.micro_batch_size,
-            num_workers=cfg.dataset.num_workers,
-            data_sharding=cfg.dataset.data_sharding,
-            collate_fn=collate_fn,
-            pin_memory=cfg.dataset.pin_memory,
-            persistent_workers=cfg.dataset.persistent_workers,
-            data_parallel_rank=sampler_dp_rank,
-            data_parallel_size=sampler_dp_size,
-            drop_last=cfg.dataset.drop_last,
-        )
+        if scalable_dp:
+            # Shard reads on the canonical grid (LCM of the module DP sizes) so every
+            # module materializes the same ordered global micro-batch under any sampler.
+            module_dps = [
+                p.data_parallel_size for p in cfg.model.megatron_mimo_parallelism_config.module_parallelisms.values()
+            ]
+            train_loader = build_canonical_mimo_data_loader(
+                train_ds,
+                consumed_samples=train_state.consumed_train_samples,
+                dataloader_type=cfg.dataset.dataloader_type,
+                micro_batch_size=cfg.train.micro_batch_size,
+                module_dp_sizes=module_dps,
+                dp_rank=sampler_dp_rank,
+                dp_size=sampler_dp_size,
+                data_sharding=cfg.dataset.data_sharding,
+                drop_last=cfg.dataset.drop_last,
+                num_workers=cfg.dataset.num_workers,
+                pin_memory=cfg.dataset.pin_memory,
+                collate_fn=collate_fn,
+                persistent_workers=cfg.dataset.persistent_workers,
+            )
+        else:
+            train_loader = build_pretraining_data_loader(
+                dataset=train_ds,
+                consumed_samples=train_state.consumed_train_samples,
+                dataloader_type=cfg.dataset.dataloader_type,
+                micro_batch_size=cfg.train.micro_batch_size,
+                num_workers=cfg.dataset.num_workers,
+                data_sharding=cfg.dataset.data_sharding,
+                collate_fn=collate_fn,
+                pin_memory=cfg.dataset.pin_memory,
+                persistent_workers=cfg.dataset.persistent_workers,
+                data_parallel_rank=sampler_dp_rank,
+                data_parallel_size=sampler_dp_size,
+                drop_last=cfg.dataset.drop_last,
+            )
 
         # `pretrain_megatron_mimo` calls `next(data_iterator)` per microbatch, so
         # return an iterator (DataLoader is iterable but not itself an iterator).
@@ -1230,6 +1256,14 @@ def _parse_args() -> argparse.Namespace:
         "tokens into one [1, T] THD sequence so the language model skips padding compute "
         "(block-diagonal attention comes from cu_seqlens).",
     )
+    parser.add_argument(
+        "--scalable-dp",
+        action="store_true",
+        help="Scalable data parallelism: each rank reads only its disjoint 1/dp shard of the global "
+        "micro-batch instead of every rank reading the full batch and slicing locally (IO scales with "
+        "DP). Each rank processes its natural, unbalanced shard. Uses the same DP loss reduction as "
+        "non-scalable runs.",
+    )
     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)

src/megatron/bridge/data/builders/direct_hf_sft.py

@@ -93,6 +93,7 @@ class DirectHFSFTDatasetConfig(DataloaderConfig):
     pad_to_max_length: bool = False
     pad_to_multiple_of: int = 128
     in_batch_packing_pad_to_multiple_of: int = 1
+    megatron_mimo_scalable_dp: bool = False
 
     def validate(self) -> None:
         """Validate declarative source and dataset settings."""
@@ -112,6 +113,18 @@ def validate(self) -> None:
             self.test_source.validate()
         if self.hf_processor_path is not None and not self.hf_processor_path.strip():
             raise ValueError("hf_processor_path must be a non-empty string when set.")
+        if self.megatron_mimo_scalable_dp:
+            if self.dataloader_type not in ("single", "cyclic"):
+                raise ValueError(
+                    "megatron_mimo_scalable_dp requires dataloader_type 'single' or 'cyclic' "
+                    f"(got {self.dataloader_type!r}); other samplers have no cross-module-consistent "
+                    "shard assignment."
+                )
+            if not self.drop_last:
+                raise ValueError(
+                    "megatron_mimo_scalable_dp requires drop_last=True; a partial final micro-batch "
+                    "gives modules unequal shares and misaligns the modality routing."
+                )
         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(

src/megatron/bridge/data/megatron_mimo/canonical_sampler.py

@@ -0,0 +1,213 @@
+# 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.
+
+"""Canonical-grid batch sampling for MegatronMIMO scalable data-parallel reads.
+
+With ``megatron_mimo_scalable_dp`` every module's loaders must materialize the same
+ordered global micro-batch, because the ``BridgeCommunicator`` routes modality
+embeddings to language ranks by contiguous position along the batch dim. A sampler's
+shard assignment depends on its ``(data_parallel_rank, data_parallel_size,
+micro_batch_size)``, so modules with different DP sizes cannot shard by their own DP:
+a shuffling sampler would hand them different sample sets. Instead, all loaders shard
+on one shared **canonical grid** — the least common multiple of the module DP sizes.
+A rank whose module has DP size ``d`` covers ``grid // d`` consecutive canonical
+groups and concatenates their windows, so concatenating any module's rank batches
+reproduces the identical ordered global micro-batch under any deterministic sampler
+(``single`` and ``cyclic`` included; the shuffle seed derives from the epoch alone).
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Callable, Iterator
+
+from torch.utils.data import DataLoader, Dataset
+
+from megatron.bridge.data.samplers import MegatronPretrainingRandomSampler, MegatronPretrainingSampler
+
+
+def canonical_grid_size(module_dp_sizes: list[int]) -> int:
+    """Return the shared sampler grid size: the LCM of every module's DP size."""
+    if not module_dp_sizes or any(dp is None or dp < 1 for dp in module_dp_sizes):
+        raise ValueError(f"module DP sizes must be positive integers (got {module_dp_sizes}).")
+    return math.lcm(*module_dp_sizes)
+
+
+def covered_canonical_groups(dp_rank: int, dp_size: int, grid_size: int) -> list[int]:
+    """Return the canonical groups this rank reads.
+
+    The groups are the ``grid_size // dp_size`` consecutive slots matching the
+    contiguous batch-dim chunk the ``BridgeCommunicator`` routes to this DP rank.
+    """
+    if grid_size % dp_size != 0:
+        raise ValueError(
+            f"canonical grid size ({grid_size}) is not divisible by the module DP size ({dp_size}); "
+            "module DP sizes must divide their least common multiple."
+        )
+    span = grid_size // dp_size
+    return list(range(dp_rank * span, (dp_rank + 1) * span))
+
+
+class CanonicalGroupBatchSampler:
+    """Concatenate per-canonical-group Megatron samplers into one batch sampler.
+
+    Holds one flat sampler per covered canonical group, all built with the same
+    ``(micro_batch // grid, grid)`` geometry and the same global ``consumed_samples``.
+    Each yield emits the groups' current windows concatenated in group order. Group
+    streams are functions of ``(group, grid, micro_batch, consumed, epoch seed)`` only,
+    so every rank covering group ``g`` — in any module — sees the identical stream.
+    """
+
+    def __init__(self, samplers: list) -> None:
+        if not samplers:
+            raise ValueError("CanonicalGroupBatchSampler needs at least one group sampler.")
+        self.samplers = samplers
+
+    def __len__(self) -> int:
+        """Match the flat Megatron samplers' convention of reporting total samples."""
+        return min(len(sampler) for sampler in self.samplers)
+
+    def __iter__(self) -> Iterator[list[int]]:
+        """Yield one concatenated index window per micro-batch."""
+        iterators = [iter(sampler) for sampler in self.samplers]
+        while True:
+            window: list[int] = []
+            for iterator in iterators:
+                group_batch = next(iterator, None)
+                if group_batch is None:
+                    # Groups share one truncated geometry, so they exhaust on the same window.
+                    return
+                window.extend(group_batch)
+            yield window
+
+
+def build_canonical_group_batch_sampler(
+    *,
+    dataloader_type: str,
+    dataset: Dataset,
+    consumed_samples: int,
+    micro_batch_size: int,
+    grid_size: int,
+    groups: list[int],
+    data_sharding: bool,
+    drop_last: bool = True,
+) -> CanonicalGroupBatchSampler:
+    """Build this rank's canonical-group batch sampler for scalable MIMO reads.
+
+    Args:
+        dataloader_type: ``"single"`` or ``"cyclic"``; other types have no shard
+            assignment that is consistent across modules and are rejected.
+        dataset: The dataset the loader reads.
+        consumed_samples: Global consumed-sample count, exactly as the flat samplers
+            expect (used for resume / epoch derivation).
+        micro_batch_size: The global micro-batch size (not the per-rank share).
+        grid_size: Canonical grid size from :func:`canonical_grid_size`.
+        groups: This rank's groups from :func:`covered_canonical_groups`.
+        data_sharding: Passed through to the cyclic sampler.
+        drop_last: Must stay ``True``: a partial final window gives the groups
+            unequal shares and breaks the positional routing.
+
+    Returns:
+        The merged batch sampler for ``torch.utils.data.DataLoader(batch_sampler=...)``.
+    """
+    if not drop_last:
+        raise ValueError("megatron_mimo_scalable_dp requires drop_last=True (partial windows misalign modules).")
+    if micro_batch_size % grid_size != 0:
+        raise ValueError(
+            f"micro_batch_size ({micro_batch_size}) must be divisible by the canonical grid size ({grid_size})."
+        )
+    group_micro_batch_size = micro_batch_size // grid_size
+    # Truncate to whole global micro-batches: the flat samplers round their active range at
+    # per-group granularity, which diverges per group for a non-multiple dataset size (the
+    # cyclic data_sharding=False stride would give groups unequal window counts).
+    total_samples = (len(dataset) // micro_batch_size) * micro_batch_size
+    if total_samples <= 0:
+        raise ValueError(f"dataset ({len(dataset)} samples) is smaller than one micro-batch ({micro_batch_size}).")
+
+    samplers = []
+    for group in groups:
+        if dataloader_type == "single":
+            samplers.append(
+                MegatronPretrainingSampler(
+                    total_samples=total_samples,
+                    consumed_samples=consumed_samples,
+                    micro_batch_size=group_micro_batch_size,
+                    data_parallel_rank=group,
+                    data_parallel_size=grid_size,
+                    drop_last=True,
+                )
+            )
+        elif dataloader_type == "cyclic":
+            samplers.append(
+                MegatronPretrainingRandomSampler(
+                    dataset,
+                    total_samples=total_samples,
+                    consumed_samples=consumed_samples,
+                    micro_batch_size=group_micro_batch_size,
+                    data_parallel_rank=group,
+                    data_parallel_size=grid_size,
+                    data_sharding=data_sharding,
+                )
+            )
+        else:
+            raise ValueError(
+                f"megatron_mimo_scalable_dp supports dataloader_type 'single' or 'cyclic' (got {dataloader_type!r})."
+            )
+    return CanonicalGroupBatchSampler(samplers)
+
+
+def build_canonical_mimo_data_loader(
+    dataset: Dataset | None,
+    *,
+    consumed_samples: int,
+    dataloader_type: str,
+    micro_batch_size: int,
+    module_dp_sizes: list[int],
+    dp_rank: int,
+    dp_size: int,
+    data_sharding: bool,
+    drop_last: bool,
+    num_workers: int,
+    pin_memory: bool,
+    collate_fn: Callable | None,
+    persistent_workers: bool,
+) -> DataLoader | None:
+    """Build this rank's read-sharded DataLoader for scalable MegatronMIMO reads.
+
+    Computes the canonical grid from ``module_dp_sizes``, derives this rank's covered
+    groups from its module-local ``(dp_rank, dp_size)``, and wraps the merged batch
+    sampler in a ``DataLoader``. Returns ``None`` when ``dataset`` is ``None``.
+    """
+    if dataset is None:
+        return None
+    grid = canonical_grid_size(module_dp_sizes)
+    groups = covered_canonical_groups(dp_rank, dp_size, grid)
+    batch_sampler = build_canonical_group_batch_sampler(
+        dataloader_type=dataloader_type,
+        dataset=dataset,
+        consumed_samples=consumed_samples,
+        micro_batch_size=micro_batch_size,
+        grid_size=grid,
+        groups=groups,
+        data_sharding=data_sharding,
+        drop_last=drop_last,
+    )
+    return DataLoader(
+        dataset,
+        batch_sampler=batch_sampler,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        collate_fn=collate_fn,
+        persistent_workers=persistent_workers,
+    )