feat: Enhance dataset loading efficiency with tensor parallelism

src/megatron/bridge/data/utils.py

@@ -15,6 +15,7 @@
 from dataclasses import fields
 from typing import Any, Callable, Dict, Optional, Type, Union
 
+from megatron.core import parallel_state
 from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
 from megatron.core.datasets.blended_megatron_dataset_config import BlendedMegatronDatasetConfig
 from megatron.core.datasets.gpt_dataset import GPTDataset, MockGPTDataset
@@ -37,17 +38,27 @@
 from megatron.bridge.utils.common_utils import print_rank_0
 
 
-def is_dataset_built_on_rank(pg_collection: ProcessGroupCollection) -> bool:
+def is_dataset_built_on_rank(pg_collection: Optional[ProcessGroupCollection] = None) -> bool:
     """Determines whether the dataset should be built on the current rank.
 
     Datasets are typically built only on the first and last pipeline stages
     and the first tensor parallel rank to save memory and avoid redundancy.
 
+    Args:
+        pg_collection: Process group collection. When provided, uses the
+            explicit process groups. When ``None``, falls back to the global
+            parallel state, which allows this function to be passed directly
+            as a zero-argument callable to ``BlendedMegatronDatasetBuilder``.
+
     Returns:
         True if the dataset should be built on the current rank, False otherwise.
     """
-    return (is_pp_first_stage(pg_collection.pp) or is_pp_last_stage(pg_collection.pp)) and (
-        pg_collection.tp.rank() == 0
+    if pg_collection is not None:
+        return (is_pp_first_stage(pg_collection.pp) or is_pp_last_stage(pg_collection.pp)) and (
+            pg_collection.tp.rank() == 0
+        )
+    return (parallel_state.is_pipeline_first_stage() or parallel_state.is_pipeline_last_stage()) and (
+        parallel_state.get_tensor_model_parallel_rank() == 0
     )
 
 
@@ -76,9 +87,11 @@ def pretrain_train_valid_test_datasets_provider(
 
     print_rank_0("> building train, validation, and test datasets for GPT ...")
 
-    # Build the dataset on all ranks for TP-replicated loading
+    broadcast_data = getattr(dataset_config, "broadcast_data_across_tp", False)
+    is_built_on_rank = is_dataset_built_on_rank if broadcast_data else (lambda: True)
+
     train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder(
-        dataset_type, train_val_test_num_samples, lambda: True, dataset_config
+        dataset_type, train_val_test_num_samples, is_built_on_rank, dataset_config
     ).build()
 
     print_rank_0("> finished creating GPT datasets ...")

src/megatron/bridge/training/config.py

@@ -255,6 +255,15 @@ class DataloaderConfig:
     trust_remote_code: Optional[bool] = None
     """Whether remote code execution should be trusted for a given HF path."""
 
+    broadcast_data_across_tp: bool = False
+    """When True, only TP-rank-0 loads data and broadcasts to other TP ranks.
+
+    This eliminates redundant I/O across tensor-parallel ranks and is critical
+    for storage backends with high contention costs (e.g. VAST / network-attached
+    storage).  When False (default), every rank loads data independently, which
+    works well on low-latency parallel file-systems like Lustre.
+    """
+
 
 @dataclass(frozen=True)
 class DatasetBuildContext:

src/megatron/bridge/training/gpt_step.py

@@ -32,6 +32,7 @@
 from megatron.bridge.training.losses import masked_next_token_loss
 from megatron.bridge.training.post_training.distillation import loss_func_kd
 from megatron.bridge.training.state import GlobalState
+from megatron.bridge.training.utils.batch_utils import get_batch_on_this_tp_rank
 from megatron.bridge.training.utils.packed_seq_utils import get_packed_seq_params
 from megatron.bridge.training.utils.pg_utils import get_pg_collection
 
@@ -169,13 +170,19 @@ def get_batch(
     if (not is_first) and (not is_last):
         return None, None, None, None, None, None, None, None, None, None
 
-    batch = get_batch_from_iterator(
-        data_iterator,
-        use_mtp,
-        getattr(cfg.dataset, "skip_getting_attention_mask_from_dataset", True),
-        is_first_pp_stage=is_first,
-        is_last_pp_stage=is_last,
-    )
+    broadcast_data = getattr(cfg.dataset, "broadcast_data_across_tp", False)
+    if broadcast_data:
+        # TP-rank-0 loads data and broadcasts to other TP ranks.
+        # Reduces I/O by a factor of TP on high-latency storage.
+        batch = get_batch_on_this_tp_rank(data_iterator, cfg, use_mtp, pg_collection=pg_collection)
+    else:
+        batch = get_batch_from_iterator(
+            data_iterator,
+            use_mtp,
+            getattr(cfg.dataset, "skip_getting_attention_mask_from_dataset", True),
+            is_first_pp_stage=is_first,
+            is_last_pp_stage=is_last,
+        )
 
     cp_size = pg_collection.cp.size()
     has_packed = batch.get("cu_seqlens") is not None

src/megatron/bridge/training/utils/batch_utils.py

@@ -149,4 +149,16 @@ def _broadcast(item):
             "position_ids": position_ids,
         }
 
+    # Broadcast any extra keys (e.g. packed-sequence metadata) that are not
+    # covered by the fixed-shape direct broadcasts above.
+    extra = {k: v for k, v in data.items() if k not in batch} if is_tp_rank0 else None
+    obj_list = [extra]
+    torch.distributed.broadcast_object_list(obj_list, src=tp_ranks[0], group=tp_group)
+    extra = obj_list[0]
+    for key, val in extra.items():
+        if isinstance(val, torch.Tensor):
+            batch[key] = val.cuda(non_blocking=True)
+        else:
+            batch[key] = val
+
     return batch