From bee7ca8f1e8b06eb4a608873749172fd63ff7c7f Mon Sep 17 00:00:00 2001 From: ilml Date: Mon, 15 Jun 2026 17:14:03 -0700 Subject: [PATCH 1/2] Clean up pretrain_gpt.py formatting and remove module globals Run black/isort and apply minor readability fixes to pretrain_gpt.py: - Hoist the get_batch BATCH_KEYS list to a module-level constant. - Replace the mutable loss_func_cached_logits global and manual lazy-init with an lru_cache-backed builder. - Drop the unnecessary `global stimer` declaration in forward_step (it is only read, never rebound). - Fix a recurring "determinisic" typo and add docstrings to two helpers. No behavior change. Signed-off-by: ilml --- pretrain_gpt.py | 157 ++++++++++++++++++++++++++++++++++++------------ 1 file changed, 120 insertions(+), 37 deletions(-) diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 503aa31503d..1884b728e05 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -18,7 +18,7 @@ warnings.filterwarnings("ignore", category=UserWarning) warnings.filterwarnings("ignore", category=FutureWarning) -from functools import partial +from functools import lru_cache, partial from typing import Any, List, Optional, Tuple import torch @@ -54,7 +54,7 @@ print_rank_0, set_startup_timestamps, ) -from megatron.training.argument_utils import pretrain_cfg_container_from_args, gpt_config_from_args +from megatron.training.argument_utils import gpt_config_from_args, pretrain_cfg_container_from_args from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args from megatron.training.datasets.fim_dataset import GPTFIMDataset, GPTFIMDatasetConfig from megatron.training.datasets.sft_dataset import SFTDataset @@ -72,12 +72,25 @@ stimer = StragglerDetector() +# Canonical, ordered schema of the fields ``get_batch`` returns. Kept alphabetical +# to match the historical ``sorted(batch.keys())`` order that callers unpack into. +BATCH_KEYS = [ + "attention_mask", + "cu_seqlens", + "cu_seqlens_padded", + "hybrid_cp_group", + "labels", + "local_cp_size", + "loss_mask", + "max_seqlen", + "position_ids", + "tokens", +] + def get_batch(data_iterator, vp_stage: Optional[int] = None): """Generate a batch.""" - BATCH_KEYS = ["attention_mask", "cu_seqlens", "cu_seqlens_padded", "hybrid_cp_group", "labels", "local_cp_size", "loss_mask", "max_seqlen", "position_ids", "tokens"] - args = get_args() config = core_transformer_config_from_args(args) @@ -85,7 +98,12 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): tp_rank = mpu.get_tensor_model_parallel_rank() is_sft = args.sft create_attention_mask_in_dataloader = args.create_attention_mask_in_dataloader - mtp_on_this_rank = mtp_on_this_rank_func(layout=config.pipeline_model_parallel_layout, mtp_num_layers=config.mtp_num_layers, ignore_virtual=False, vp_stage=vp_stage) + mtp_on_this_rank = mtp_on_this_rank_func( + layout=config.pipeline_model_parallel_layout, + mtp_num_layers=config.mtp_num_layers, + ignore_virtual=False, + vp_stage=vp_stage, + ) is_hybrid_cp = args.hybrid_context_parallel if not is_first_or_last_pipeline_stage(vp_stage) and not mtp_on_this_rank and not is_sft: @@ -95,15 +113,50 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): if tp_rank == 0: batch = next(data_iterator) for key in BATCH_KEYS: - batch[key] = batch[key].cuda(non_blocking=True) if key in batch and batch[key] is not None else None + batch[key] = ( + batch[key].cuda(non_blocking=True) + if key in batch and batch[key] is not None + else None + ) + + batch = get_batch_on_this_tp_rank( + batch, + broadcast_src_rank=mpu.get_tensor_model_parallel_src_rank(), + broadcast_group=mpu.get_tensor_model_parallel_group(), + is_sft=is_sft, + is_hybrid_cp=is_hybrid_cp, + create_attention_mask_in_dataloader=create_attention_mask_in_dataloader, + cp_size=cp_size, + tp_rank=tp_rank, + micro_batch_size=args.micro_batch_size, + seq_length=args.seq_length, + mtp_on_this_rank=mtp_on_this_rank, + pipeline_model_parallel_size=args.pipeline_model_parallel_size, + is_pipeline_first_stage=mpu.is_pipeline_first_stage(), + is_pipeline_last_stage=mpu.is_pipeline_last_stage(), + ) - batch = get_batch_on_this_tp_rank(batch, broadcast_src_rank=mpu.get_tensor_model_parallel_src_rank(), broadcast_group=mpu.get_tensor_model_parallel_group(), is_sft=is_sft, is_hybrid_cp=is_hybrid_cp, create_attention_mask_in_dataloader=create_attention_mask_in_dataloader, cp_size=cp_size, tp_rank=tp_rank, micro_batch_size=args.micro_batch_size, seq_length=args.seq_length, mtp_on_this_rank=mtp_on_this_rank, pipeline_model_parallel_size=args.pipeline_model_parallel_size, is_pipeline_first_stage=mpu.is_pipeline_first_stage(), is_pipeline_last_stage=mpu.is_pipeline_last_stage()) - if not is_first_or_last_pipeline_stage(vp_stage) and not mtp_on_this_rank: assert is_sft - return None, batch['cu_seqlens'], batch['cu_seqlens_padded'], None, None, None, None, batch['max_seqlen'], None, None - - batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=is_hybrid_cp, cp_group=get_context_parallel_group(), hybrid_cp_group_func=get_hybrid_data_context_parallel_groups) + return ( + None, + batch['cu_seqlens'], + batch['cu_seqlens_padded'], + None, + None, + None, + None, + batch['max_seqlen'], + None, + None, + ) + + batch = get_batch_on_this_cp_rank( + batch, + is_hybrid_cp=is_hybrid_cp, + cp_group=get_context_parallel_group(), + hybrid_cp_group_func=get_hybrid_data_context_parallel_groups, + ) # Return values in BATCH_KEYS order so callers can unpack into the fixed # names regardless of any provenance fields wrappers like BlendedDataset @@ -117,7 +170,27 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): # define spiky loss as a loss that's 10x the max loss observed SPIKY_LOSS_FACTOR = 10 -loss_func_cached_logits = None + +@lru_cache(maxsize=1) +def _build_cached_logits_loss_func( + logprobs_dir, decode_threads, prefetch_factor, msc_prefetch_depth, kd_loss_alpha, ignore_errors +): + """Build (once) the offline knowledge-distillation loss callable for cached logits. + + Memoized so the teacher log-probability reader is constructed a single time per + process, replacing the previous module-level mutable global. + """ + from megatron.training.distillation import LossFuncCallable + + return LossFuncCallable( + logprobs_dir=logprobs_dir, + decode_threads=decode_threads, + prefetch_factor=prefetch_factor, + msc_prefetch_depth=msc_prefetch_depth, + kd_loss_alpha=kd_loss_alpha, + ignore_errors=ignore_errors, + ) + def loss_func( loss_mask: torch.Tensor, output_tensor: torch.Tensor, model: Optional[GPTModel] = None @@ -139,17 +212,14 @@ def loss_func( if args.logits_load_dir is not None: # Offline knowledge distillation loss using cached teacher log-probabilities. - global loss_func_cached_logits - if loss_func_cached_logits is None: - from megatron.training.distillation import LossFuncCallable - loss_func_cached_logits = LossFuncCallable( - logprobs_dir=args.logits_load_dir, - decode_threads=args.logits_load_decode_threads, - prefetch_factor=args.logits_load_prefetch_factor, - msc_prefetch_depth=args.logits_load_msc_prefetch_depth, - kd_loss_alpha=args.logits_load_kd_loss_alpha, - ignore_errors=args.logits_load_ignore_errors, - ) + loss_func_cached_logits = _build_cached_logits_loss_func( + logprobs_dir=args.logits_load_dir, + decode_threads=args.logits_load_decode_threads, + prefetch_factor=args.logits_load_prefetch_factor, + msc_prefetch_depth=args.logits_load_msc_prefetch_depth, + kd_loss_alpha=args.logits_load_kd_loss_alpha, + ignore_errors=args.logits_load_ignore_errors, + ) loss, num_tokens, report = loss_func_cached_logits(loss_mask, output_tensor, model=model) elif has_nvidia_modelopt and getattr(args, 'modelopt_enabled', False): # [ModelOpt] loss, num_tokens, report = loss_func_modelopt(loss_mask, output_tensor, model=model) @@ -168,14 +238,14 @@ def loss_func( result=loss, rejection_func=torch.isnan, message="found NaN in local forward loss calculation", - tolerance=0.0, # forward pass calculations are determinisic + tolerance=0.0, # forward pass calculations are deterministic fatal=True, ) rerun_state_machine.validate_result( result=loss, rejection_func=torch.isinf, message="found Inf in local forward loss calculation", - tolerance=0.0, # forward pass calculations are determinisic + tolerance=0.0, # forward pass calculations are deterministic fatal=True, ) # Check for spiky loss @@ -188,7 +258,7 @@ def loss_func( context="loss", ), message="Spiky loss", - tolerance=0.0, # forward pass calculations are determinisic + tolerance=0.0, # forward pass calculations are deterministic fatal=False, ) @@ -208,7 +278,6 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa # Get the batch. timers('batch-generator', log_level=2).start() - global stimer with stimer(bdata=True): vp_stage = get_attr_wrapped_model(model, "vp_stage") ( @@ -234,7 +303,9 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa # Use real (unpadded) cu_seqlens to feed the FLOPs accounting: varlen # attention only computes work for real tokens within each chunk. update_seqlen_stats_from_cu_seqlens(cu_seqlens) - cu_seqlens_for_params = cu_seqlens_padded if cu_seqlens_padded is not None else cu_seqlens # TODO(asolergi-nv): Currently there is a bug forcing cu_seqlens to be cu_seqlens_padded + cu_seqlens_for_params = ( + cu_seqlens_padded if cu_seqlens_padded is not None else cu_seqlens + ) # TODO(asolergi-nv): Currently there is a bug forcing cu_seqlens to be cu_seqlens_padded packed_seq_params = PackedSeqParams( qkv_format="thd", cu_seqlens_q=cu_seqlens_for_params, @@ -251,15 +322,21 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa with stimer: if return_schedule_plan: - assert args.overlap_moe_expert_parallel_comm, \ - "overlap_moe_expert_parallel_comm must be enabled to return the schedule plan" + assert ( + args.overlap_moe_expert_parallel_comm + ), "overlap_moe_expert_parallel_comm must be enabled to return the schedule plan" schedule_plan = model.build_schedule_plan( tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask ) return schedule_plan, partial(loss_func, loss_mask, model=model) else: output_tensor = model( - tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask, packed_seq_params=packed_seq_params + tokens, + position_ids, + attention_mask, + labels=labels, + loss_mask=loss_mask, + packed_seq_params=packed_seq_params, ) # [ModelOpt]: model is needed to access ModelOpt distillation losses @@ -267,19 +344,23 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa def is_dataset_built_on_rank(vp_stage=None, is_packed_sequence=False): + """Whether the dataset should be built on the current rank.""" args = get_args() config = core_transformer_config_from_args(args) if mpu.get_tensor_model_parallel_rank() != 0: return False elif is_packed_sequence: return True - return ( - is_first_or_last_pipeline_stage(vp_stage) - or mtp_on_this_rank_func(layout=config.pipeline_model_parallel_layout, mtp_num_layers=config.mtp_num_layers, ignore_virtual=False, vp_stage=vp_stage) + return is_first_or_last_pipeline_stage(vp_stage) or mtp_on_this_rank_func( + layout=config.pipeline_model_parallel_layout, + mtp_num_layers=config.mtp_num_layers, + ignore_virtual=False, + vp_stage=vp_stage, ) def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig: + """Build the GPT (or FIM) dataset config from parsed CLI args.""" tokenizer = build_tokenizer(args) # Sometimes --data-path is too long, instead we parse it from a file. @@ -354,7 +435,6 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None config = core_gpt_dataset_config_from_args(args) - is_packed_sequence = False if args.sft: dataset_type = SFTDataset @@ -369,7 +449,9 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None print_rank_0("> building train, validation, and test datasets for GPT ...") - is_dataset_built = partial(is_dataset_built_on_rank, vp_stage=vp_stage, is_packed_sequence=is_packed_sequence) + is_dataset_built = partial( + is_dataset_built_on_rank, vp_stage=vp_stage, is_packed_sequence=is_packed_sequence + ) train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder( dataset_type, train_val_test_num_samples, is_dataset_built, config ).build() @@ -413,7 +495,8 @@ def get_embedding_ranks(pp_ranks: List[int]): ) model_cfg = gpt_config_from_args(args) full_config = pretrain_cfg_container_from_args(args, model_cfg) - pretrain(full_config, + pretrain( + full_config, train_valid_test_datasets_provider, partial(model_provider, gpt_builder), ModelType.encoder_or_decoder, From 0dc4317e231f5f7041997cd548aadef50eafd909 Mon Sep 17 00:00:00 2001 From: ilml Date: Mon, 15 Jun 2026 18:57:35 -0700 Subject: [PATCH 2/2] Apply the same cleanup to pretrain_hybrid.py Mirror the pretrain_gpt.py housekeeping in pretrain_hybrid.py: - Hoist the get_batch BATCH_KEYS list to a module-level constant. - Replace the mutable loss_func_cached_logits global and manual lazy-init with an lru_cache-backed builder. - Drop the unnecessary `global stimer` declaration in forward_step. - Add docstrings to two helpers and run black/isort. No behavior change. Signed-off-by: ilml --- pretrain_hybrid.py | 152 +++++++++++++++++++++++++++++++++------------ 1 file changed, 114 insertions(+), 38 deletions(-) diff --git a/pretrain_hybrid.py b/pretrain_hybrid.py index 6530aa6531d..02448b2f755 100644 --- a/pretrain_hybrid.py +++ b/pretrain_hybrid.py @@ -17,7 +17,7 @@ warnings.filterwarnings("ignore", category=UserWarning) warnings.filterwarnings("ignore", category=FutureWarning) -from functools import partial +from functools import lru_cache, partial from typing import Any, List, Optional, Tuple import torch @@ -65,18 +65,32 @@ try: from megatron.post_training.arguments import add_modelopt_args from megatron.post_training.loss_func import loss_func as loss_func_modelopt + has_nvidia_modelopt = True except ImportError: has_nvidia_modelopt = False stimer = StragglerDetector() +# Canonical, ordered schema of the fields ``get_batch`` returns. Kept alphabetical +# to match the historical ``sorted(batch.keys())`` order that callers unpack into. +BATCH_KEYS = [ + "attention_mask", + "cu_seqlens", + "cu_seqlens_padded", + "hybrid_cp_group", + "labels", + "local_cp_size", + "loss_mask", + "max_seqlen", + "position_ids", + "tokens", +] + def get_batch(data_iterator, vp_stage=None): """Generate a batch.""" - BATCH_KEYS = ["attention_mask", "cu_seqlens", "cu_seqlens_padded", "hybrid_cp_group", "labels", "local_cp_size", "loss_mask", "max_seqlen", "position_ids", "tokens"] - args = get_args() config = core_transformer_config_from_args(args) @@ -84,7 +98,12 @@ def get_batch(data_iterator, vp_stage=None): tp_rank = mpu.get_tensor_model_parallel_rank() is_sft = args.sft create_attention_mask_in_dataloader = args.create_attention_mask_in_dataloader - mtp_on_this_rank = mtp_on_this_rank_func(layout=config.pipeline_model_parallel_layout, mtp_num_layers=config.mtp_num_layers, ignore_virtual=False, vp_stage=vp_stage) + mtp_on_this_rank = mtp_on_this_rank_func( + layout=config.pipeline_model_parallel_layout, + mtp_num_layers=config.mtp_num_layers, + ignore_virtual=False, + vp_stage=vp_stage, + ) is_hybrid_cp = args.hybrid_context_parallel if not is_first_or_last_pipeline_stage(vp_stage) and not mtp_on_this_rank and not is_sft: @@ -94,15 +113,50 @@ def get_batch(data_iterator, vp_stage=None): if tp_rank == 0: batch = next(data_iterator) for key in BATCH_KEYS: - batch[key] = batch[key].cuda(non_blocking=True) if key in batch and batch[key] is not None else None + batch[key] = ( + batch[key].cuda(non_blocking=True) + if key in batch and batch[key] is not None + else None + ) - batch = get_batch_on_this_tp_rank(batch, broadcast_src_rank=mpu.get_tensor_model_parallel_src_rank(), broadcast_group=mpu.get_tensor_model_parallel_group(), is_sft=is_sft, is_hybrid_cp=is_hybrid_cp, create_attention_mask_in_dataloader=create_attention_mask_in_dataloader, cp_size=cp_size, tp_rank=tp_rank, micro_batch_size=args.micro_batch_size, seq_length=args.seq_length, mtp_on_this_rank=mtp_on_this_rank, pipeline_model_parallel_size=args.pipeline_model_parallel_size, is_pipeline_first_stage=mpu.is_pipeline_first_stage(), is_pipeline_last_stage=mpu.is_pipeline_last_stage()) + batch = get_batch_on_this_tp_rank( + batch, + broadcast_src_rank=mpu.get_tensor_model_parallel_src_rank(), + broadcast_group=mpu.get_tensor_model_parallel_group(), + is_sft=is_sft, + is_hybrid_cp=is_hybrid_cp, + create_attention_mask_in_dataloader=create_attention_mask_in_dataloader, + cp_size=cp_size, + tp_rank=tp_rank, + micro_batch_size=args.micro_batch_size, + seq_length=args.seq_length, + mtp_on_this_rank=mtp_on_this_rank, + pipeline_model_parallel_size=args.pipeline_model_parallel_size, + is_pipeline_first_stage=mpu.is_pipeline_first_stage(), + is_pipeline_last_stage=mpu.is_pipeline_last_stage(), + ) if not is_first_or_last_pipeline_stage(vp_stage) and not mtp_on_this_rank: assert is_sft - return None, batch['cu_seqlens'], batch['cu_seqlens_padded'], None, None, None, None, batch['max_seqlen'], None, None + return ( + None, + batch['cu_seqlens'], + batch['cu_seqlens_padded'], + None, + None, + None, + None, + batch['max_seqlen'], + None, + None, + ) - batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=is_hybrid_cp, cp_group=get_context_parallel_group(), hybrid_cp_group_func=get_hybrid_data_context_parallel_groups) + batch = get_batch_on_this_cp_rank( + batch, + is_hybrid_cp=is_hybrid_cp, + cp_group=get_context_parallel_group(), + hybrid_cp_group_func=get_hybrid_data_context_parallel_groups, + ) # Return values in BATCH_KEYS order so callers can unpack into the fixed # names regardless of any provenance fields wrappers like BlendedDataset @@ -116,9 +170,31 @@ def get_batch(data_iterator, vp_stage=None): # define spiky loss as a loss that's 10x the max loss observed SPIKY_LOSS_FACTOR = 10 -loss_func_cached_logits = None -def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor, model: Optional[HybridModel] = None): +@lru_cache(maxsize=1) +def _build_cached_logits_loss_func( + logprobs_dir, decode_threads, prefetch_factor, msc_prefetch_depth, kd_loss_alpha, ignore_errors +): + """Build (once) the offline knowledge-distillation loss callable for cached logits. + + Memoized so the teacher log-probability reader is constructed a single time per + process, replacing the previous module-level mutable global. + """ + from megatron.training.distillation import LossFuncCallable + + return LossFuncCallable( + logprobs_dir=logprobs_dir, + decode_threads=decode_threads, + prefetch_factor=prefetch_factor, + msc_prefetch_depth=msc_prefetch_depth, + kd_loss_alpha=kd_loss_alpha, + ignore_errors=ignore_errors, + ) + + +def loss_func( + loss_mask: torch.Tensor, output_tensor: torch.Tensor, model: Optional[HybridModel] = None +): """Loss function. Args: @@ -134,17 +210,14 @@ def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor, model: Optio args = get_args() if args.logits_load_dir is not None: # Offline knowledge distillation loss using cached teacher log-probabilities. - global loss_func_cached_logits - if loss_func_cached_logits is None: - from megatron.training.distillation import LossFuncCallable - loss_func_cached_logits = LossFuncCallable( - logprobs_dir=args.logits_load_dir, - decode_threads=args.logits_load_decode_threads, - prefetch_factor=args.logits_load_prefetch_factor, - msc_prefetch_depth=args.logits_load_msc_prefetch_depth, - kd_loss_alpha=args.logits_load_kd_loss_alpha, - ignore_errors=args.logits_load_ignore_errors, - ) + loss_func_cached_logits = _build_cached_logits_loss_func( + logprobs_dir=args.logits_load_dir, + decode_threads=args.logits_load_decode_threads, + prefetch_factor=args.logits_load_prefetch_factor, + msc_prefetch_depth=args.logits_load_msc_prefetch_depth, + kd_loss_alpha=args.logits_load_kd_loss_alpha, + ignore_errors=args.logits_load_ignore_errors, + ) loss, num_tokens, report = loss_func_cached_logits(loss_mask, output_tensor, model=model) elif has_nvidia_modelopt and getattr(args, 'modelopt_enabled', False): # [ModelOpt] loss, num_tokens, report = loss_func_modelopt(loss_mask, output_tensor, model=model) @@ -163,14 +236,14 @@ def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor, model: Optio result=loss, rejection_func=torch.isnan, message="found NaN in local forward loss calculation", - tolerance=0.0, # forward pass calculations are deterministic + tolerance=0.0, # forward pass calculations are deterministic fatal=True, ) rerun_state_machine.validate_result( result=loss, rejection_func=torch.isinf, message="found Inf in local forward loss calculation", - tolerance=0.0, # forward pass calculations are deterministic + tolerance=0.0, # forward pass calculations are deterministic fatal=True, ) # Check for spiky loss @@ -183,7 +256,7 @@ def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor, model: Optio context="loss", ), message="Spiky loss", - tolerance=0.0, # forward pass calculations are deterministic + tolerance=0.0, # forward pass calculations are deterministic fatal=False, ) @@ -202,8 +275,6 @@ def forward_step(data_iterator, model: HybridModel): # Get the batch. timers('batch-generator', log_level=2).start() - global stimer - with stimer(bdata=True): vp_stage = get_attr_wrapped_model(model, "vp_stage") ( @@ -252,7 +323,7 @@ def forward_step(data_iterator, model: HybridModel): attention_mask, labels=labels, packed_seq_params=packed_seq_params, - loss_mask=loss_mask + loss_mask=loss_mask, ) # [ModelOpt]: model is needed to access ModelOpt distillation losses @@ -260,19 +331,23 @@ def forward_step(data_iterator, model: HybridModel): def is_dataset_built_on_rank(vp_stage=None, is_packed_sequence=False): + """Whether the dataset should be built on the current rank.""" args = get_args() config = core_transformer_config_from_args(args) if mpu.get_tensor_model_parallel_rank() != 0: return False elif is_packed_sequence: return True - return ( - is_first_or_last_pipeline_stage(vp_stage) - or mtp_on_this_rank_func(layout=config.pipeline_model_parallel_layout, mtp_num_layers=config.mtp_num_layers, ignore_virtual=False, vp_stage=vp_stage) + return is_first_or_last_pipeline_stage(vp_stage) or mtp_on_this_rank_func( + layout=config.pipeline_model_parallel_layout, + mtp_num_layers=config.mtp_num_layers, + ignore_virtual=False, + vp_stage=vp_stage, ) def core_gpt_dataset_config_from_args(args: Any) -> GPTDatasetConfig: + """Build the GPT dataset config from parsed CLI args.""" tokenizer = build_tokenizer(args) # Sometimes --data-path is too long, instead we parse it from a file. @@ -339,7 +414,7 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None dataset_type, train_val_test_num_samples, partial(is_dataset_built_on_rank, vp_stage=vp_stage, is_packed_sequence=is_packed_sequence), - config + config, ).build() print_rank_0("> finished creating GPT datasets ...") @@ -366,10 +441,11 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None ) model_cfg = hybrid_config_from_args(args) full_config = pretrain_cfg_container_from_args(args, model_cfg) - pretrain(full_config, - train_valid_test_datasets_provider, - partial(model_provider, hybrid_builder), - ModelType.encoder_or_decoder, - forward_step, - store=store, - ) + pretrain( + full_config, + train_valid_test_datasets_provider, + partial(model_provider, hybrid_builder), + ModelType.encoder_or_decoder, + forward_step, + store=store, + )