diff --git a/examples/post_training/modelopt/generation_server.py b/examples/post_training/modelopt/generation_server.py deleted file mode 100644 index b32cca0d73f..00000000000 --- a/examples/post_training/modelopt/generation_server.py +++ /dev/null @@ -1,198 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. - -"""Sample Generate""" -import os -import sys -import warnings -from functools import partial - -sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))) -import os -import sys -from argparse import Namespace -from contextlib import nullcontext - -import torch - -from megatron.core import mpu -from megatron.core.inference.engines import AbstractEngine, StaticInferenceEngine -from megatron.core.inference.engines.abstract_engine import AbstractEngine -from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import ( - InferenceWrapperConfig, -) -from megatron.core.inference.sampling_params import SamplingParams -from megatron.core.inference.text_generation_controllers.text_generation_controller import ( - TextGenerationController, -) -from megatron.core.transformer.module import MegatronModule -from megatron.inference.text_generation import beam_search_and_post_process -from megatron.inference.text_generation.mcore_engine_server import ( - ModelInferenceWrapperServer, - run_mcore_engine, -) -from megatron.inference.text_generation_server import MegatronServer -from megatron.post_training.arguments import add_modelopt_args -from megatron.training import get_args, get_model, get_tokenizer, print_rank_0 -from megatron.training.checkpointing import load_checkpoint -from megatron.training.initialize import initialize_megatron - - -def get_inference_engine(args: Namespace, model: MegatronModule) -> AbstractEngine: - """Get the relevant backend for running inference - - This function will automatically choose the TRTLLMBackend when possible, and default to Mcore - backend if the user does not specify any backends. TRTLLMBackend is not implmented yet. - - Args: - args (Namespace): The user arguments parsed from command line - model (MegatronModule): The megatron model. - - Returns: - AbstractBackend: The chosen backend - """ - tokenizer = get_tokenizer() - - inference_wrapper_config = InferenceWrapperConfig( - hidden_size=args.hidden_size, - inference_batch_times_seqlen_threshold=args.inference_batch_times_seqlen_threshold, - fp32_residual_connection=args.fp32_residual_connection, - params_dtype=args.params_dtype, - padded_vocab_size=args.padded_vocab_size, - inference_max_seq_length=args.inference_max_seq_length, - inference_max_requests=args.inference_max_batch_size, - nccl_all_reduce_for_prefill=args.nccl_all_reduce_for_prefill, - ) - - inference_wrapped_model = ModelInferenceWrapperServer(model, inference_wrapper_config) - text_generation_controller = TextGenerationController( - inference_wrapped_model=inference_wrapped_model, tokenizer=tokenizer - ) - return StaticInferenceEngine( - text_generation_controller=text_generation_controller, - max_batch_size=args.inference_max_batch_size, - ) - - -def add_text_generate_args(parser): - """Adds text generation arguments to parser.""" - group = parser.add_argument_group(title='text generation') - group.add_argument( - "--port", type=int, default=5000, help='port for text generation server to run on' - ) - group.add_argument("--temperature", type=float, default=1.0, help='Sampling temperature.') - group.add_argument("--top_k", type=int, default=1, help='Top k sampling.') - group.add_argument("--top_p", type=float, default=0.0, help='Top p sampling.') - group.add_argument( - "--return-log-probs", - action='store_true', - default=True, - help='Return the log probabilities of the final output tokens', - ) - group.add_argument( - "--num-tokens-to-generate", - type=int, - default=30, - help='Number of tokens to generate for each prompt', - ) - group.add_argument( - "--prompts", - metavar='N', - type=str, - nargs='+', - help='Input prompts with each prompt within quotes and seperated by space', - ) - group.add_argument( - "--max-batch-size", - type=int, - default=None, - help='Deprecated in favor of `--inference-max-batch-size`', - ) - add_modelopt_args(parser) - return parser - - -@torch.inference_mode() -def main(model_provider: str = "gpt"): - """Runs the text generation server with the specified model provider.""" - initialize_megatron( - extra_args_provider=add_text_generate_args, - args_defaults={ - 'no_load_rng': True, - 'no_load_optim': True, - 'exit_on_missing_checkpoint': True, - }, - ) - 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.") - exit() - print_rank_0("WARNING: Forcing exit_on_missing_checkpoint to True for text " "generation.") - args.exit_on_missing_checkpoint = True - - # Set up model and load checkpoint - load_context = nullcontext() - if args.fp8: - from transformer_engine.pytorch.fp8 import fp8_model_init - - load_context = fp8_model_init() - with load_context: - - from megatron.post_training.model_builder import modelopt_gpt_mamba_builder - from model_provider import model_provider as root_model_provider - if model_provider == "gpt": - model = get_model(partial(root_model_provider, modelopt_gpt_mamba_builder), wrap_with_ddp=False) - elif model_provider == "mamba": - pass - else: - raise ValueError(f"Invalid model provider {model_provider}") - - if args.load is not None: - _ = load_checkpoint(model, None, None, strict=False) - - assert len(model) == 1, "Above condition should have caught this" - model = model[0] - model.eval() - - if args.max_batch_size is not None: - assert args.inference_max_batch_size is not None - args.inference_max_batch_size = max(args.inference_max_batch_size, args.max_batch_size) - warnings.warn( - "`--max-batch-size` has been deprecated in favor of `--inference-max-requests`, " - f"setting maximum batch size to {args.inference_max_batch_size}" - ) - - inference_engine = get_inference_engine(args, model) - - if args.cuda_graph_impl == "local": - print(f"Running warmup for CUDA graphs...") - inference_engine.generate( - prompts=["Test prompt"], sampling_params=SamplingParams(num_tokens_to_generate=10) - ) - - if ( - mpu.is_pipeline_first_stage() - and mpu.get_tensor_model_parallel_rank() == 0 - and mpu.get_expert_model_parallel_rank() == 0 - ): - server = MegatronServer(inference_engine, args) - server.run("0.0.0.0", port=args.port) - - while True: - choice = torch.tensor(1, dtype=torch.long, device='cuda') - torch.distributed.broadcast(choice, 0) - if choice.item() == 0: - try: - run_mcore_engine(inference_engine) - except ValueError as ve: - pass - elif choice.item() == 1: - try: - beam_search_and_post_process( - inference_engine.text_generation_controller.inference_wrapped_model.model - ) - except ValueError as ve: - pass - - -if __name__ == "__main__": - main(model_provider="gpt") diff --git a/examples/post_training/modelopt/generation_server.sh b/examples/post_training/modelopt/generation_server.sh old mode 100644 new mode 100755 index a4d7ff2dada..9acd61f3d04 --- a/examples/post_training/modelopt/generation_server.sh +++ b/examples/post_training/modelopt/generation_server.sh @@ -14,11 +14,9 @@ if [ -z ${MLM_MODEL_CKPT} ]; then exit 1 fi -if [ -z ${DRAFT_LEN} ]; then - DRAFT_LEN=0 -fi +TOOLS_DIR="$(realpath ${SCRIPT_DIR}/../../../tools)" -${LAUNCH_SCRIPT} ${SCRIPT_DIR}/generation_server.py \ +${LAUNCH_SCRIPT} ${TOOLS_DIR}/run_text_generation_server.py \ ${MODEL_ARGS} \ --tensor-model-parallel-size ${TP} \ --expert-tensor-parallel-size ${ETP} \ diff --git a/megatron/training/training.py b/megatron/training/training.py index 55cc6f00e41..d8b57326f67 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -855,6 +855,18 @@ def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap args = get_args() args.model_type = model_type + if has_nvidia_modelopt: + from megatron.post_training.checkpointing import has_modelopt_state + # [ModelOpt]: Check if the checkpoint is a ModelOpt checkpoint and + # set a flag to use our model provider if so. + if args.load is not None and has_modelopt_state(args.load): + print_rank_0(f'ModelOpt checkpoint detected') + args.modelopt_enabled = True + elif getattr(args, "export_kd_teacher_load", None): + # For distillation ckpts without ModelOpt state + args.modelopt_enabled = True + + # Build model. def build_model(): if ( @@ -1081,17 +1093,6 @@ def setup_model_and_optimizer( timers = get_timers() one_logger = get_one_logger() - if has_nvidia_modelopt: - from megatron.post_training.checkpointing import has_modelopt_state - # [ModelOpt]: Check if the checkpoint is a ModelOpt checkpoint and - # set a flag to use our model provider if so. - if args.load is not None and has_modelopt_state(args.load): - print_rank_0(f'ModelOpt checkpoint detected') - args.modelopt_enabled = True - elif getattr(args, "export_kd_teacher_load", None): - # For distillation ckpts without ModelOpt state - args.modelopt_enabled = True - model = get_model(model_provider_func, model_type) unwrapped_model = unwrap_model(model) diff --git a/tools/run_text_generation_server.py b/tools/run_text_generation_server.py index fb5212f7649..350173dc16f 100644 --- a/tools/run_text_generation_server.py +++ b/tools/run_text_generation_server.py @@ -32,6 +32,7 @@ from megatron.core.inference.text_generation_server import MegatronServer from megatron.core.inference.text_generation_server.run_mcore_engine import run_mcore_engine from megatron.core.transformer.module import MegatronModule +from megatron.post_training.arguments import add_modelopt_args from megatron.training import get_model, print_rank_0 from model_provider import model_provider @@ -120,6 +121,7 @@ def add_text_generate_args(parser): default=None, help='Deprecated in favor of `--inference-max-batch-size`', ) + add_modelopt_args(parser) return parser