From 8d15161981e087864153ea50ebab3b278988be5b Mon Sep 17 00:00:00 2001 From: Keshav Santhanam Date: Wed, 15 Apr 2026 12:35:57 -0700 Subject: [PATCH] Fix args parsing Signed-off-by: Keshav Santhanam --- .../detxoify_lm/generate_samples_gpt.py | 12 +++++++----- .../model_converter/vision_model_tester.py | 4 +++- examples/multimodal/run_text_generation.py | 4 +++- tools/run_dynamic_text_generation_server.py | 4 +++- tools/run_inference_performance_test.py | 4 +++- tools/run_text_generation_server.py | 4 +++- tools/run_vlm_text_generation.py | 4 +++- 7 files changed, 25 insertions(+), 11 deletions(-) diff --git a/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py b/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py index 895a45d0242..183c6c695bd 100644 --- a/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py +++ b/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py @@ -14,6 +14,7 @@ from megatron.training import print_rank_0 from megatron.training.checkpointing import load_checkpoint from megatron.core import mpu +from megatron.training.arguments import parse_and_validate_args from megatron.training.initialize import initialize_megatron from megatron.legacy.model import GPTModel from megatron.training import get_model @@ -232,11 +233,12 @@ def generate_and_write_samples_conditional(model): def main(): """Main program.""" - initialize_megatron(extra_args_provider=add_text_generate_args, - args_defaults={'tokenizer_type': 'GPT2BPETokenizer', - 'no_load_rng': True, - 'no_load_optim': True, - 'seq_length': 2048}) + parse_and_validate_args(extra_args_provider=add_text_generate_args, + args_defaults={'tokenizer_type': 'GPT2BPETokenizer', + 'no_load_rng': True, + 'no_load_optim': True, + 'seq_length': 2048}) + initialize_megatron() # Set up model and load checkpoint model = get_model(model_provider, wrap_with_ddp=False) diff --git a/examples/multimodal/model_converter/vision_model_tester.py b/examples/multimodal/model_converter/vision_model_tester.py index ef36dd5f9e0..36e2122b555 100644 --- a/examples/multimodal/model_converter/vision_model_tester.py +++ b/examples/multimodal/model_converter/vision_model_tester.py @@ -17,6 +17,7 @@ from examples.multimodal.model import model_provider from examples.multimodal.multimodal_args import add_multimodal_extra_args from megatron.training import get_model +from megatron.training.arguments import parse_and_validate_args from megatron.training.checkpointing import load_checkpoint from megatron.training.initialize import initialize_megatron @@ -50,7 +51,8 @@ def run_mcore_vision(model_path): f"--pretrained-checkpoint={model_path}", ] - initialize_megatron(extra_args_provider=add_multimodal_extra_args) + parse_and_validate_args(extra_args_provider=add_multimodal_extra_args) + initialize_megatron() def wrapped_model_provider(pre_process, post_process): return model_provider(pre_process, post_process, parallel_output=False) diff --git a/examples/multimodal/run_text_generation.py b/examples/multimodal/run_text_generation.py index e55679c1b2e..532d0771a94 100644 --- a/examples/multimodal/run_text_generation.py +++ b/examples/multimodal/run_text_generation.py @@ -39,6 +39,7 @@ VLMInferenceWrapper, ) from megatron.training import get_args, get_model, get_tokenizer, print_rank_0, is_last_rank +from megatron.training.arguments import parse_and_validate_args from megatron.training.checkpointing import load_checkpoint from megatron.training.initialize import initialize_megatron @@ -842,7 +843,8 @@ def run_evaluation_loop(model, configs, output_dir_override=None, iteration=None def eval_tasks(): """Vision language model text generation for single or batch tasks.""" - initialize_megatron(extra_args_provider=add_text_generation_args) + parse_and_validate_args(extra_args_provider=add_text_generation_args) + initialize_megatron() args = get_args() diff --git a/tools/run_dynamic_text_generation_server.py b/tools/run_dynamic_text_generation_server.py index 74fbfd8e3d3..ed5ab473d63 100644 --- a/tools/run_dynamic_text_generation_server.py +++ b/tools/run_dynamic_text_generation_server.py @@ -14,6 +14,7 @@ from megatron.inference.utils import add_inference_args, get_dynamic_inference_engine from megatron.post_training.arguments import add_modelopt_args from megatron.training import get_args +from megatron.training.arguments import parse_and_validate_args from megatron.training.initialize import initialize_megatron @@ -76,10 +77,11 @@ async def run_text_generation_server( if __name__ == "__main__": with torch.inference_mode(): - initialize_megatron( + parse_and_validate_args( extra_args_provider=add_text_generation_server_args, args_defaults={'no_load_rng': True, 'no_load_optim': True}, ) + initialize_megatron() # Enable return_log_probs to allow prompt logprobs computation for echo=True requests # This sets materialize_only_last_token_logits=False in the inference context, diff --git a/tools/run_inference_performance_test.py b/tools/run_inference_performance_test.py index a1feef380ef..ac9e92d3639 100644 --- a/tools/run_inference_performance_test.py +++ b/tools/run_inference_performance_test.py @@ -39,6 +39,7 @@ from megatron.core import mpu from megatron.training import get_args, get_model, get_tokenizer from megatron.training.checkpointing import load_checkpoint +from megatron.training.arguments import parse_and_validate_args from megatron.training.initialize import initialize_megatron REQUEST_ID = 0 @@ -151,7 +152,7 @@ def generate_dynamic( def main(): """Main program.""" - initialize_megatron( + parse_and_validate_args( extra_args_provider=add_inference_benchmarking_args, args_defaults={ 'no_load_rng': True, @@ -160,6 +161,7 @@ def main(): 'exit_on_missing_checkpoint': True, }, ) + initialize_megatron() args = get_args() diff --git a/tools/run_text_generation_server.py b/tools/run_text_generation_server.py index 83a0bbc2369..5a2940f1a4c 100644 --- a/tools/run_text_generation_server.py +++ b/tools/run_text_generation_server.py @@ -40,6 +40,7 @@ from megatron.core import mpu from megatron.training import get_args, get_model, get_tokenizer from megatron.training.checkpointing import load_checkpoint +from megatron.training.arguments import parse_and_validate_args from megatron.training.initialize import initialize_megatron @@ -113,7 +114,7 @@ def add_text_generate_args(parser): @torch.inference_mode() def main(model_type: str = "gpt"): """Runs the text generation server with the specified model type.""" - initialize_megatron( + parse_and_validate_args( extra_args_provider=add_text_generate_args, args_defaults={ 'no_load_rng': True, @@ -121,6 +122,7 @@ def main(model_type: str = "gpt"): 'exit_on_missing_checkpoint': True, }, ) + initialize_megatron() args = get_args() if args.num_layers_per_virtual_pipeline_stage is not None: print("Interleaved pipeline schedule is not yet supported for text generation.") diff --git a/tools/run_vlm_text_generation.py b/tools/run_vlm_text_generation.py index 3e9a6f785de..17c3b1030dc 100644 --- a/tools/run_vlm_text_generation.py +++ b/tools/run_vlm_text_generation.py @@ -19,6 +19,7 @@ from megatron.inference.text_generation.api import generate_and_post_process from megatron.inference.text_generation.forward_step import ForwardStep from megatron.training import get_args, get_model, print_rank_0 +from megatron.training.arguments import parse_and_validate_args from megatron.training.checkpointing import load_checkpoint from megatron.training.initialize import initialize_megatron from pretrain_vlm import model_provider @@ -199,7 +200,8 @@ def main(): logging.getLogger(__name__).warning("Models using pipeline parallelism are not supported yet.") - initialize_megatron(extra_args_provider=add_text_generation_args) + parse_and_validate_args(extra_args_provider=add_text_generation_args) + initialize_megatron() # Set up model and load checkpoint. model = get_model(model_provider, wrap_with_ddp=False)