Extract args init to launch scripts

megatron/training/arguments.py

@@ -31,6 +31,7 @@
 )
 from megatron.core.activations import squared_relu
 from megatron.core.fusions.fused_bias_geglu import quick_gelu
+from megatron.training.global_vars import set_global_variables
 from megatron.training.utils import (
     get_device_arch_version,
     update_use_dist_ckpt,
@@ -85,6 +86,35 @@ def add_megatron_arguments(parser: argparse.ArgumentParser):
 
     return parser
 
+def parse_and_validate_args(extra_args_provider=None, ignore_unknown_args=False, args_defaults={}):
+    args = parse_args(extra_args_provider, ignore_unknown_args)
+
+    if args.use_checkpoint_args or args_defaults.get("use_checkpoint_args", False):
+        from megatron.training.checkpointing import load_args_from_checkpoint
+
+        assert args.load is not None or args.pretrained_checkpoint is not None, "--use-checkpoint-args requires --load or --pretrained-checkpoint argument"
+        assert args.non_persistent_ckpt_type != "local", (
+            "--use-checkpoint-args is not supported with --non_persistent_ckpt_type=local. "
+            "Two-stage checkpoint loading is not implemented, and all arguments must be defined "
+            "before initializing LocalCheckpointManager."
+        )
+        load_args_from_checkpoint(args, load_arg='pretrained_checkpoint')
+        load_args_from_checkpoint(args)
+
+    if args.yaml_cfg is not None:
+        from megatron.training.yaml_arguments import validate_yaml
+
+        args = validate_yaml(args, args_defaults)
+    else:
+        validate_args(args, args_defaults)
+
+    # set global args, build tokenizer, and set adlr-autoresume,
+    # tensorboard-writer, and timers.
+    set_global_variables(args)
+
+    return args
+
+
 def parse_args(extra_args_provider=None, ignore_unknown_args=False):
     """Parse all arguments."""
     parser = argparse.ArgumentParser(description='Megatron-LM Arguments',
@@ -293,6 +323,12 @@ def tuple_type(x):
 
 def validate_args(args, defaults={}):
 
+    # Prep for checkpoint conversion.
+    if args.ckpt_convert_format is not None:
+        assert args.ckpt_convert_save is not None
+        assert args.load is not None
+        args.exit_on_missing_checkpoint = True
+
     # Temporary
     assert args.non_persistent_ckpt_type in ['global', 'local', None], \
         'Currently only global and local checkpoints are supported'

megatron/training/initialize.py

@@ -32,25 +32,17 @@
     get_tensorboard_writer,
     inprocess_restart,
 )
-from megatron.training.arguments import parse_args, validate_args
 from megatron.training.async_utils import init_persistent_async_worker
-from megatron.training.checkpointing import load_args_from_checkpoint
-from megatron.training.global_vars import set_global_variables
 from megatron.training.utils import is_rank0, print_rank_0, warn_rank_0
-from megatron.training.yaml_arguments import validate_yaml
 
 logger = logging.getLogger(__name__)
 
 
 def initialize_megatron(
-    extra_args_provider=None,
-    args_defaults={},
-    ignore_unknown_args=False,
     allow_no_cuda=False,
     skip_mpu_initialization=False,
     get_embedding_ranks=None,
     get_position_embedding_ranks=None,
-    parsed_args=None,
     store=None,
 ):
     """Set global variables, initialize distributed, and
@@ -65,38 +57,7 @@ def initialize_megatron(
         # Make sure cuda is available.
         assert torch.cuda.is_available(), "Megatron requires CUDA."
 
-    # Parse arguments
-    if parsed_args is None:
-        args = parse_args(extra_args_provider, ignore_unknown_args)
-    else:
-        args = parsed_args
-
-    # Prep for checkpoint conversion.
-    if args.ckpt_convert_format is not None:
-        assert args.ckpt_convert_save is not None
-        assert args.load is not None
-        args.exit_on_missing_checkpoint = True
-
-    if args.use_checkpoint_args or args_defaults.get("use_checkpoint_args", False):
-        assert (
-            args.load is not None or args.pretrained_checkpoint is not None
-        ), "--use-checkpoint-args requires --load or --pretrained-checkpoint argument"
-        assert args.non_persistent_ckpt_type != "local", (
-            "--use-checkpoint-args is not supported with --non_persistent_ckpt_type=local. "
-            "Two-stage checkpoint loading is not implemented, and all arguments must be defined "
-            "before initializing LocalCheckpointManager."
-        )
-        load_args_from_checkpoint(args, load_arg='pretrained_checkpoint')
-        load_args_from_checkpoint(args)
-
-    if args.yaml_cfg is not None:
-        args = validate_yaml(args, args_defaults)
-    else:
-        validate_args(args, args_defaults)
-
-    # set global args, build tokenizer, and set adlr-autoresume,
-    # tensorboard-writer, and timers.
-    set_global_variables(args)
+    args = get_args()
 
     # set logging level
     setup_logging()

megatron/training/training.py

@@ -826,8 +826,6 @@ def pretrain(
     model_type,
     forward_step_func,
     process_non_loss_data_func=None,
-    extra_args_provider=None,
-    args_defaults={},
     get_embedding_ranks=None,
     get_position_embedding_ranks=None,
     non_loss_data_func=None,
@@ -893,8 +891,6 @@ def pretrain(
 
     # Initalize and get arguments, timers, and Tensorboard writer.
     initialize_megatron(
-        extra_args_provider=extra_args_provider,
-        args_defaults=args_defaults,
         get_embedding_ranks=get_embedding_ranks,
         get_position_embedding_ranks=get_position_embedding_ranks,
         store=store,

pretrain_bert.py

@@ -16,7 +16,7 @@
 from megatron.core.models.bert.bert_model import BertModel
 from megatron.training import pretrain
 from megatron.training.utils import average_losses_across_data_parallel_group
-from megatron.training.arguments import core_transformer_config_from_args
+from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args
 from megatron.core.transformer.spec_utils import import_module
 from megatron.core.models.bert.bert_layer_specs import bert_layer_with_transformer_engine_spec, bert_layer_local_spec
 from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer
@@ -191,6 +191,7 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None
     # Temporary for transition to core datasets
     train_valid_test_datasets_provider.is_distributed = True
 
+    parse_and_validate_args(args_defaults={'tokenizer_type': 'BertWordPieceLowerCase'})
     pretrain(train_valid_test_datasets_provider, model_provider,
              ModelType.encoder_or_decoder,
-             forward_step, args_defaults={'tokenizer_type': 'BertWordPieceLowerCase'})
+             forward_step)

pretrain_gpt.py

@@ -41,7 +41,7 @@
 )
 from megatron.training.datasets.sft_dataset import SFTDataset
 from megatron.core.transformer.multi_token_prediction import mtp_on_this_rank, get_mtp_ranks
-from megatron.training.arguments import core_transformer_config_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.utils import (
     get_batch_on_this_cp_rank,
@@ -409,13 +409,15 @@ def get_embedding_ranks(pp_ranks: List[int]):
     # Optionally enable inprocess restart on pretrain
     pretrain, store = inprocess_restart.maybe_wrap_for_inprocess_restart(pretrain)
 
+    args = parse_and_validate_args(
+        extra_args_provider=add_modelopt_args if has_nvidia_modelopt else None,
+        args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
+    )
     pretrain(
         train_valid_test_datasets_provider,
         partial(model_provider, gpt_builder),
         ModelType.encoder_or_decoder,
         forward_step,
-        args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
-        extra_args_provider=add_modelopt_args if has_nvidia_modelopt else None,
         store=store,
         get_embedding_ranks=get_embedding_ranks,
     )

pretrain_mamba.py

@@ -42,6 +42,7 @@
     print_rank_0,
     set_startup_timestamps,
 )
+from megatron.training.arguments import parse_and_validate_args
 from megatron.training.datasets.sft_dataset import SFTDataset
 from megatron.training.utils import (
     get_batch_on_this_cp_rank,
@@ -356,11 +357,13 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None
     # Optionally enable inprocess restart on pretrain
     pretrain, store = inprocess_restart.maybe_wrap_for_inprocess_restart(pretrain)
 
+    args = parse_and_validate_args(
+        extra_args_provider=add_modelopt_args if has_nvidia_modelopt else None,
+        args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
+    )
     pretrain(train_valid_test_datasets_provider,
              partial(model_provider, mamba_builder),
              ModelType.encoder_or_decoder,
              forward_step,
-             args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
              store=store,
-             extra_args_provider=add_modelopt_args if has_nvidia_modelopt else None,
              )

pretrain_t5.py

@@ -26,7 +26,7 @@
     get_t5_encoder_with_transformer_engine_block_spec,
 )
 from megatron.training import get_args, get_timers, pretrain, print_rank_0
-from megatron.training.arguments import core_transformer_config_from_args
+from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args
 from pretrain_gpt import loss_func
 
 """
@@ -279,12 +279,12 @@ def t5_position_embedding_ranks(pp_ranks):
     # Temporary for transition to core datasets
     train_valid_test_datasets_provider.is_distributed = True
 
+    parse_and_validate_args(args_defaults={'tokenizer_type': 'BertWordPieceLowerCase'})
     pretrain(
         train_valid_test_datasets_provider,
         model_provider,
         ModelType.encoder_or_decoder,
         forward_step,
-        args_defaults={'tokenizer_type': 'BertWordPieceLowerCase'},
         get_embedding_ranks=t5_embedding_ranks,
         get_position_embedding_ranks=t5_position_embedding_ranks,
     )

pretrain_vlm.py

@@ -38,7 +38,7 @@
     pretrain,
     print_rank_0,
 )
-from megatron.training.arguments import core_transformer_config_from_args
+from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args
 from pretrain_gpt import loss_func
 
 
@@ -439,13 +439,15 @@ def llava_position_embedding_ranks(pp_ranks):
 if __name__ == "__main__":
     train_valid_test_datasets_provider.is_distributed = True
 
+    parse_and_validate_args(
+        extra_args_provider=add_vlm_extra_args,
+        args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
+    )
     pretrain(
         train_valid_test_datasets_provider,
         model_provider,
         ModelType.encoder_or_decoder,
         forward_step,
-        args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
-        extra_args_provider=add_vlm_extra_args,
         get_embedding_ranks=llava_embedding_ranks,
         get_position_embedding_ranks=llava_position_embedding_ranks,
     )

train_rl.py

@@ -23,7 +23,7 @@
 )
 from megatron.training import get_args, get_timers, pretrain, print_rank_0
 from megatron.training.utils import is_hybrid_model
-from megatron.training.arguments import core_transformer_config_from_args
+from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args
 from model_provider import model_provider
 
 from megatron.core.packed_seq_params import PackedSeqParams
@@ -410,11 +410,13 @@ def _model_builder(
                 pg_collection=pg_collection,
             )
 
+    parse_and_validate_args(
+        extra_args_provider=add_inference_args,
+        args_defaults={},
+    )
     pretrain(
         None,  # we don't need to build any datasets for RL training
         partial(model_provider, _model_builder),
         ModelType.encoder_or_decoder,
         forward_step,
-        args_defaults={},
-        extra_args_provider=add_inference_args,
     )