From 74cc8d18e994760aca5383e87a78b6ea43159b66 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Tue, 26 Aug 2025 14:14:27 -0700 Subject: [PATCH] switch pretrain to replicated load from broadcast (#491) Signed-off-by: Ananth Subramaniam --- src/megatron/bridge/data/utils.py | 3 ++- src/megatron/bridge/training/gpt_step.py | 9 +-------- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/src/megatron/bridge/data/utils.py b/src/megatron/bridge/data/utils.py index 2973553c02..fb5d2b8383 100644 --- a/src/megatron/bridge/data/utils.py +++ b/src/megatron/bridge/data/utils.py @@ -69,8 +69,9 @@ 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 train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder( - dataset_type, train_val_test_num_samples, is_dataset_built_on_rank, dataset_config + dataset_type, train_val_test_num_samples, lambda: True, dataset_config ).build() print_rank_0("> finished creating GPT datasets ...") diff --git a/src/megatron/bridge/training/gpt_step.py b/src/megatron/bridge/training/gpt_step.py index 2cdacdee5c..5d0db05e48 100644 --- a/src/megatron/bridge/training/gpt_step.py +++ b/src/megatron/bridge/training/gpt_step.py @@ -269,14 +269,7 @@ def get_batch( if (not parallel_state.is_pipeline_first_stage()) and (not parallel_state.is_pipeline_last_stage()): return None, None, None, None, None, None, None, None - if isinstance(cfg.dataset, FinetuningDatasetConfig): - batch = get_batch_from_iterator(data_iterator, use_mtp) - else: - # get batches based on the TP rank you are on - batch = get_batch_on_this_tp_rank(data_iterator, cfg, use_mtp) - batch["cu_seqlens"] = None - batch["cu_seqlens_argmin"] = None - batch["max_seqlen"] = None + batch = get_batch_from_iterator(data_iterator, use_mtp) # slice batch along sequence dimension for context parallelism batch = get_batch_on_this_cp_rank(batch)