From 7d35d0caed31cb445ee364d96358d8f267492540 Mon Sep 17 00:00:00 2001 From: shanmugamr1992 Date: Thu, 4 Jun 2026 15:15:15 -0700 Subject: [PATCH 1/5] Doing first set of changes for refactor --- megatron/inference/utils.py | 98 ++++++++++++++++++++----- tools/run_inference_performance_test.py | 11 +-- tools/run_text_generation_server.py | 26 ++----- 3 files changed, 94 insertions(+), 41 deletions(-) diff --git a/megatron/inference/utils.py b/megatron/inference/utils.py index 91a9d954617..20e6d66e283 100644 --- a/megatron/inference/utils.py +++ b/megatron/inference/utils.py @@ -1,13 +1,13 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging -from argparse import ArgumentParser +import warnings +from argparse import ArgumentParser, Namespace from functools import partial -from typing import Optional +from typing import Callable, Optional + import torch -from gpt_builders import gpt_builder -from hybrid_builders import hybrid_builder from megatron.core.inference.config import ( CudaGraphSizingDistribution, InferenceConfig, @@ -32,34 +32,98 @@ from megatron.training import get_args from megatron.training import get_model as _get_model from megatron.training import get_tokenizer, get_wandb_writer +from megatron.training.argument_utils import gpt_config_from_args, hybrid_config_from_args from megatron.training.checkpointing import load_checkpoint +from megatron.training.models import ( + GPTModelBuilder, + HybridModelBuilder, + ModelBuilder, +) from model_provider import model_provider logger = logging.getLogger(__name__) -def get_model_for_inference() -> MegatronModule: - """Initialize model and load checkpoint for inference.""" +def get_model_builder(args: Namespace, provider: Optional[str] = None) -> ModelBuilder: + """Construct a :class:`ModelBuilder` for the requested model provider. - args = get_args() + Replaces the legacy ``gpt_builder`` / ``hybrid_builder`` function selector with + a config-driven dispatch that returns a fully-configured :class:`ModelBuilder` + instance whose ``build_model()`` and ``build_distributed_models()`` methods can + be used to materialize the model. - if args.model_provider == "gpt": - model_builder = gpt_builder - elif args.model_provider in ("hybrid", "mamba"): - if args.model_provider == "mamba": - import warnings + Args: + args: The parsed argparse namespace, used to populate the model config via + ``gpt_config_from_args`` / ``hybrid_config_from_args``. + provider: Optional override for the model provider name. Must be one of + ``"gpt"``, ``"hybrid"``, or the deprecated ``"mamba"``. When omitted, + falls back to ``args.model_provider`` (set by ``add_inference_args``). + Returns: + A :class:`ModelBuilder` instance bound to a config derived from ``args``. + """ + if provider is None: + provider = args.model_provider + if provider == "gpt": + return GPTModelBuilder(gpt_config_from_args(args)) + if provider in ("hybrid", "mamba"): + if provider == "mamba": warnings.warn( - '--model-provider "mamba" is deprecated. Use --model-provider "hybrid" instead.', + '"mamba" model provider is deprecated. Use "hybrid" instead.', DeprecationWarning, stacklevel=2, ) - model_builder = hybrid_builder - else: - raise ValueError(f"Invalid model provider {args.model_provider}") + return HybridModelBuilder(hybrid_config_from_args(args)) + raise ValueError(f"Invalid model provider {provider}") + + +def builder_to_legacy_callable(builder: ModelBuilder) -> Callable: + """Adapt a :class:`ModelBuilder` to the legacy ``model_provider`` callable signature. + + Routing through ``model_provider`` (from the top-level ``model_provider.py``) + preserves legacy hooks that are still on the function-style builder API + (e.g. ``args.modelopt_enabled`` swaps in ``modelopt_gpt_hybrid_builder`` and + ``args.record_memory_history`` registers the OOM observer) until those code + paths are migrated to the new builder API. + + NOTE : If the args.record_history_memory and has_nvidia_modelopt are not needed, then can do a direct pass through. + + + + Intended usage:: + + builder = get_model_builder(args) + model = get_model( + partial(model_provider, builder_to_legacy_callable(builder)), + wrap_with_ddp=False, + ) + """ + + def legacy_builder( + args, pre_process, post_process, vp_stage=None, config=None, pg_collection=None + ): + return builder.build_model( + pg_collection=pg_collection, + pre_process=pre_process, + post_process=post_process, + vp_stage=vp_stage, + ) + + return legacy_builder + + +def get_model_for_inference() -> MegatronModule: + """Initialize model and load checkpoint for inference.""" + + args = get_args() + + builder = get_model_builder(args) # Build model. - model = _get_model(partial(model_provider, model_builder), wrap_with_ddp=False) + model = _get_model( + partial(model_provider, builder_to_legacy_callable(builder)), + wrap_with_ddp=False, + ) # Load checkpoint. assert args.load is not None diff --git a/tools/run_inference_performance_test.py b/tools/run_inference_performance_test.py index d42453c62ed..934bcd1878f 100644 --- a/tools/run_inference_performance_test.py +++ b/tools/run_inference_performance_test.py @@ -8,8 +8,6 @@ import torch -from gpt_builders import gpt_builder -from hybrid_builders import hybrid_builder from megatron.core.inference.contexts import StaticInferenceContext from megatron.core.inference.engines import DynamicInferenceEngine, StaticInferenceEngine from megatron.core.inference.engines.abstract_engine import AbstractEngine @@ -26,8 +24,11 @@ ) from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer from megatron.core.transformer.module import MegatronModule -from megatron.inference.utils import add_inference_args, get_dynamic_inference_engine, get_model_for_inference -from model_provider import model_provider +from megatron.inference.utils import ( + add_inference_args, + get_dynamic_inference_engine, + get_model_for_inference, +) sys.path.append( os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir)) @@ -38,8 +39,8 @@ 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.checkpointing import load_checkpoint from megatron.training.initialize import initialize_megatron REQUEST_ID = 0 diff --git a/tools/run_text_generation_server.py b/tools/run_text_generation_server.py index e871214e739..abab95479c7 100644 --- a/tools/run_text_generation_server.py +++ b/tools/run_text_generation_server.py @@ -14,8 +14,6 @@ import torch -from gpt_builders import gpt_builder -from hybrid_builders import hybrid_builder from megatron.core.inference.contexts import StaticInferenceContext from megatron.core.inference.engines import AbstractEngine, StaticInferenceEngine from megatron.core.inference.engines.abstract_engine import AbstractEngine @@ -29,6 +27,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.inference.utils import builder_to_legacy_callable, get_model_builder from megatron.post_training.arguments import add_modelopt_args from megatron.training import get_model, print_rank_0 from model_provider import model_provider @@ -39,8 +38,8 @@ 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.checkpointing import load_checkpoint from megatron.training.initialize import initialize_megatron @@ -137,22 +136,11 @@ def main(model_type: str = "gpt"): load_context = fp8_model_init() with load_context: - # Set up model and load checkpoint - if model_type == "gpt": - model_builder = gpt_builder - elif model_type in ("hybrid", "mamba"): - if model_type == "mamba": - import warnings - - warnings.warn( - 'model_type="mamba" is deprecated. Use model_type="hybrid" instead.', - DeprecationWarning, - stacklevel=2, - ) - model_builder = hybrid_builder - else: - raise ValueError(f"Invalid model provider {model_type}") - model = get_model(partial(model_provider, model_builder), wrap_with_ddp=False) + builder = get_model_builder(args, provider=model_type) + model = get_model( + partial(model_provider, builder_to_legacy_callable(builder)), + wrap_with_ddp=False, + ) if args.load is not None: _ = load_checkpoint(model, None, None, strict=False) From f51bd1ff67c390d18dc10532ad02f7a0b1da9815 Mon Sep 17 00:00:00 2001 From: shanmugamr1992 Date: Tue, 9 Jun 2026 12:41:22 -0700 Subject: [PATCH 2/5] Address review: drop legacy callable adapter, build models directly - Tighten `provider` typehint to `Optional[Literal['gpt','hybrid','mamba']]`. - Drop `builder_to_legacy_callable`: with `model_provider` no longer in the call path, the adapter is dead weight. - `get_model_for_inference` / `run_text_generation_server.main` now call `builder.build_distributed_models(pg_collection, wrap_with_ddp=False)` directly in the common case, so we don't have to keep `get_model()` alive for the inference path. The modelopt branch still routes through `_get_model(modelopt_gpt_hybrid_builder, ...)` because the modelopt hooks haven't been ported to the new `ModelBuilder` API yet. Co-Authored-By: Claude Opus 4.7 (1M context) --- megatron/inference/utils.py | 76 ++++++++++------------------- tools/run_text_generation_server.py | 29 ++++++++--- 2 files changed, 46 insertions(+), 59 deletions(-) diff --git a/megatron/inference/utils.py b/megatron/inference/utils.py index 20e6d66e283..40e909d11dc 100644 --- a/megatron/inference/utils.py +++ b/megatron/inference/utils.py @@ -3,8 +3,7 @@ import logging import warnings from argparse import ArgumentParser, Namespace -from functools import partial -from typing import Callable, Optional +from typing import Literal, Optional import torch @@ -25,6 +24,7 @@ from megatron.core.inference.text_generation_controllers.text_generation_controller import ( TextGenerationController, ) +from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer from megatron.core.transformer.enums import InferenceCudaGraphScope from megatron.core.transformer.module import MegatronModule @@ -34,17 +34,21 @@ from megatron.training import get_tokenizer, get_wandb_writer from megatron.training.argument_utils import gpt_config_from_args, hybrid_config_from_args from megatron.training.checkpointing import load_checkpoint -from megatron.training.models import ( - GPTModelBuilder, - HybridModelBuilder, - ModelBuilder, -) -from model_provider import model_provider +from megatron.training.models import GPTModelBuilder, HybridModelBuilder, ModelBuilder + +try: + from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder + + HAS_NVIDIA_MODELOPT = True +except ImportError: + HAS_NVIDIA_MODELOPT = False logger = logging.getLogger(__name__) -def get_model_builder(args: Namespace, provider: Optional[str] = None) -> ModelBuilder: +def get_model_builder( + args: Namespace, provider: Optional[Literal["gpt", "hybrid", "mamba"]] = None +) -> ModelBuilder: """Construct a :class:`ModelBuilder` for the requested model provider. Replaces the legacy ``gpt_builder`` / ``hybrid_builder`` function selector with @@ -77,53 +81,23 @@ def get_model_builder(args: Namespace, provider: Optional[str] = None) -> ModelB raise ValueError(f"Invalid model provider {provider}") -def builder_to_legacy_callable(builder: ModelBuilder) -> Callable: - """Adapt a :class:`ModelBuilder` to the legacy ``model_provider`` callable signature. - - Routing through ``model_provider`` (from the top-level ``model_provider.py``) - preserves legacy hooks that are still on the function-style builder API - (e.g. ``args.modelopt_enabled`` swaps in ``modelopt_gpt_hybrid_builder`` and - ``args.record_memory_history`` registers the OOM observer) until those code - paths are migrated to the new builder API. - - NOTE : If the args.record_history_memory and has_nvidia_modelopt are not needed, then can do a direct pass through. - - - - Intended usage:: - - builder = get_model_builder(args) - model = get_model( - partial(model_provider, builder_to_legacy_callable(builder)), - wrap_with_ddp=False, - ) - """ - - def legacy_builder( - args, pre_process, post_process, vp_stage=None, config=None, pg_collection=None - ): - return builder.build_model( - pg_collection=pg_collection, - pre_process=pre_process, - post_process=post_process, - vp_stage=vp_stage, - ) - - return legacy_builder - - def get_model_for_inference() -> MegatronModule: """Initialize model and load checkpoint for inference.""" args = get_args() - builder = get_model_builder(args) - - # Build model. - model = _get_model( - partial(model_provider, builder_to_legacy_callable(builder)), - wrap_with_ddp=False, - ) + if HAS_NVIDIA_MODELOPT and getattr(args, "modelopt_enabled", False): + # ModelOpt path keeps the legacy callable-based builder because the + # modelopt hooks (custom layer specs, calibration, etc.) have not been + # ported to the new ``ModelBuilder`` API yet. ``_get_model`` also takes + # care of running the modelopt-checkpoint auto-detection side effect. + model = _get_model(modelopt_gpt_hybrid_builder, wrap_with_ddp=False) + else: + builder = get_model_builder(args) + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + model = builder.build_distributed_models( + pg_collection=pg_collection, wrap_with_ddp=False + ) # Load checkpoint. assert args.load is not None diff --git a/tools/run_text_generation_server.py b/tools/run_text_generation_server.py index abab95479c7..967c1668943 100644 --- a/tools/run_text_generation_server.py +++ b/tools/run_text_generation_server.py @@ -4,7 +4,6 @@ import os import sys import warnings -from functools import partial sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir))) import os @@ -26,11 +25,18 @@ ) 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.process_groups_config import ProcessGroupCollection from megatron.core.transformer.module import MegatronModule -from megatron.inference.utils import builder_to_legacy_callable, get_model_builder +from megatron.inference.utils import get_model_builder from megatron.post_training.arguments import add_modelopt_args from megatron.training import get_model, print_rank_0 -from model_provider import model_provider + +try: + from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder + + HAS_NVIDIA_MODELOPT = True +except ImportError: + HAS_NVIDIA_MODELOPT = False sys.path.append( os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir)) @@ -136,11 +142,18 @@ def main(model_type: str = "gpt"): load_context = fp8_model_init() with load_context: - builder = get_model_builder(args, provider=model_type) - model = get_model( - partial(model_provider, builder_to_legacy_callable(builder)), - wrap_with_ddp=False, - ) + if HAS_NVIDIA_MODELOPT and getattr(args, "modelopt_enabled", False): + # ModelOpt path keeps the legacy callable-based builder because the + # modelopt hooks have not been ported to the new ``ModelBuilder`` + # API yet. ``get_model`` also handles the modelopt-checkpoint + # auto-detection side effect. + model = get_model(modelopt_gpt_hybrid_builder, wrap_with_ddp=False) + else: + builder = get_model_builder(args, provider=model_type) + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + model = builder.build_distributed_models( + pg_collection=pg_collection, wrap_with_ddp=False + ) if args.load is not None: _ = load_checkpoint(model, None, None, strict=False) From c6893aa473080bd6a43aeec5b63774ddf5a51f05 Mon Sep 17 00:00:00 2001 From: shanmugamr1992 Date: Thu, 11 Jun 2026 15:38:26 -0700 Subject: [PATCH 3/5] Add InferenceConfigContainer and InferenceSetupConfig for arg-free init Introduce a declarative, serializable InferenceSetupConfig (mirroring the inference argparse group) and an InferenceConfigContainer so inference entry points can drive the refactored initialize_megatron() without constructing a training-shaped PretrainConfigContainer. - Add InferenceSetupConfig (megatron/training/config/inference_config.py) with a to_inference_config() mapper that produces the runtime megatron.core.inference.config.InferenceConfig from the built model. - Add InferenceConfigContainer holding only inference-relevant configs. - Add inference_cfg_from_args() / inference_cfg_container_from_args() helpers. - Make get_inference_config_from_model_and_args() delegate to the mapper so the declarative config is the single source of truth. --- megatron/inference/utils.py | 102 +----- megatron/training/argument_utils.py | 61 ++++ megatron/training/config/__init__.py | 3 +- megatron/training/config/container.py | 33 ++ megatron/training/config/inference_config.py | 350 +++++++++++++++++++ 5 files changed, 464 insertions(+), 85 deletions(-) create mode 100644 megatron/training/config/inference_config.py diff --git a/megatron/inference/utils.py b/megatron/inference/utils.py index 40e909d11dc..d6d27df1a38 100644 --- a/megatron/inference/utils.py +++ b/megatron/inference/utils.py @@ -7,14 +7,6 @@ import torch -from megatron.core.inference.config import ( - CudaGraphSizingDistribution, - InferenceConfig, - KVCacheManagementMode, - MambaInferenceStateConfig, - PrefixCachingCoordinatorPolicy, - PrefixCachingEvictionPolicy, -) from megatron.core.inference.contexts import DynamicInferenceContext from megatron.core.inference.engines import DynamicInferenceEngine from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import ( @@ -28,7 +20,7 @@ from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer from megatron.core.transformer.enums import InferenceCudaGraphScope from megatron.core.transformer.module import MegatronModule -from megatron.core.utils import get_attr_wrapped_model, log_single_rank, unwrap_model +from megatron.core.utils import log_single_rank, unwrap_model from megatron.training import get_args from megatron.training import get_model as _get_model from megatron.training import get_tokenizer, get_wandb_writer @@ -327,47 +319,20 @@ def add_inference_args(parser: ArgumentParser) -> ArgumentParser: def get_inference_config_from_model_and_args(model: MegatronModule, args): - """Returns a `InferenceConfig` constructed from the model and command line arguments.""" - - # Max sequence length. - position_embedding_type = get_attr_wrapped_model(model, "position_embedding_type") - model_max_seq_len = get_attr_wrapped_model(model, "max_sequence_length") - inf_max_seq_len = args.inference_max_seq_length - max_batch_size = args.inference_dynamic_batching_max_requests - - if position_embedding_type == "learned_absolute": - # When using absolute position embeddings, it is critical that the - # context's `max_sequence_length` is less than or equal to the model's - # `max_sequence_length`. Otherwise, the context's `position_ids` will - # contain ids greater than the dimension of the position embedding - # tensor, which will result in an index error. - if inf_max_seq_len: - max_sequence_length = min(model_max_seq_len, inf_max_seq_len) - else: - max_sequence_length = model_max_seq_len - assert max_batch_size is None or max_batch_size <= model_max_seq_len - else: - max_sequence_length = inf_max_seq_len - if args.inference_dynamic_batching_max_requests is not None: - max_sequence_length = max(max_sequence_length, max_batch_size) + """Returns an `InferenceConfig` constructed from the model and command line arguments. - mamba_inference_state_config = MambaInferenceStateConfig.from_model( - model, - conv_states_dtype=args.mamba_inference_conv_states_dtype, - ssm_states_dtype=args.mamba_inference_ssm_states_dtype, - ) - pg_collection = get_attr_wrapped_model(model, "pg_collection") - - # Get inference logging configuration from args - log_inference_wandb = args.inference_wandb_logging - inference_logging_step_interval = args.inference_logging_step_interval + Delegates to ``InferenceSetupConfig.to_inference_config`` so the declarative + ``InferenceSetupConfig`` (built from args) is the single source of truth for translating + inference args into the runtime engine ``InferenceConfig``. + """ + from megatron.training.argument_utils import inference_cfg_from_args - # Get metrics writer if logging is enabled and on the logging rank - # Use the same rank convention as training (last rank logs) + # Get metrics writer if logging is enabled and on the logging rank. + # Use the same rank convention as training (last rank logs). metrics_writer = None if ( - inference_logging_step_interval > 0 - and log_inference_wandb + args.inference_logging_step_interval > 0 + and args.inference_wandb_logging and args.rank == (args.world_size - 1) ): metrics_writer = get_wandb_writer() @@ -379,47 +344,16 @@ def get_inference_config_from_model_and_args(model: MegatronModule, args): "wandb module is available. Inference logging will be disabled.", ) - return InferenceConfig( - verbose=True, - block_size_tokens=args.inference_dynamic_batching_block_size, - buffer_size_gb=args.inference_dynamic_batching_buffer_size_gb, - paused_buffer_size_gb=args.inference_dynamic_batching_paused_buffer_size_gb, - mamba_memory_ratio=args.inference_dynamic_batching_mamba_memory_ratio, - num_cuda_graphs=( - args.inference_dynamic_batching_num_cuda_graphs - if args.inference_cuda_graph_scope != InferenceCudaGraphScope.none - else None - ), - max_requests=args.inference_dynamic_batching_max_requests, - max_tokens=args.inference_dynamic_batching_max_tokens, - unified_memory_level=args.inference_dynamic_batching_unified_memory_level, - kv_cache_management_mode=KVCacheManagementMode(args.rl_kv_cache_management_mode), - cuda_graph_mixed_prefill_count=args.inference_dynamic_batching_cuda_graph_mixed_prefill_count, # pylint: disable=line-too-long - cuda_graph_sizing_distribution=CudaGraphSizingDistribution( - args.inference_dynamic_batching_cuda_graph_sizing_distribution - ), - use_cuda_graphs_for_non_decode_steps=not args.decode_only_cuda_graphs, - cuda_graph_all_prefills=args.inference_cuda_graph_all_prefills, + setup_cfg = inference_cfg_from_args(args) + return setup_cfg.to_inference_config( + model, + return_log_probs=args.return_log_probs, + kv_cache_management_mode=args.rl_kv_cache_management_mode, static_kv_memory_pointers=args.rl_persist_cuda_graphs, - max_sequence_length=max_sequence_length, - mamba_inference_state_config=mamba_inference_state_config, - pg_collection=pg_collection, use_flashinfer_fused_rope=args.use_flashinfer_fused_rope, - materialize_only_last_token_logits=not (args.return_log_probs and not args.skip_prompt_log_probs), - track_generated_token_events=args.inference_dynamic_batching_track_generated_token_events, - track_paused_request_events=args.inference_dynamic_batching_track_paused_request_events, - enable_chunked_prefill=args.enable_chunked_prefill, - enable_prefix_caching=args.inference_dynamic_batching_enable_prefix_caching, - prefix_caching_eviction_policy=PrefixCachingEvictionPolicy(args.inference_dynamic_batching_prefix_caching_eviction_policy), - prefix_caching_coordinator_policy=PrefixCachingCoordinatorPolicy(args.inference_dynamic_batching_prefix_caching_coordinator_policy), - prefix_caching_routing_alpha=getattr(args, 'inference_dynamic_batching_prefix_caching_routing_alpha', 0.5), - prefix_caching_mamba_gb=getattr(args, 'inference_dynamic_batching_prefix_caching_mamba_gb', None), + enable_cuda_graphs=(args.inference_cuda_graph_scope != InferenceCudaGraphScope.none), + skip_prompt_log_probs=args.skip_prompt_log_probs, metrics_writer=metrics_writer, - logging_step_interval=args.inference_logging_step_interval, - num_speculative_tokens=args.num_speculative_tokens, - use_synchronous_zmq_collectives=args.inference_use_synchronous_zmq_collectives, - disable_ep_consensus=args.inference_disable_ep_consensus, - sampling_backend=args.inference_dynamic_batching_sampling_backend, ) diff --git a/megatron/training/argument_utils.py b/megatron/training/argument_utils.py index 2cfb3f0f17b..abe437e2ee7 100644 --- a/megatron/training/argument_utils.py +++ b/megatron/training/argument_utils.py @@ -20,6 +20,8 @@ from megatron.training.config import ( DistributedInitConfig, + InferenceSetupConfig, + InferenceConfigContainer, PretrainConfigContainer, SchedulerConfig, TokenizerConfig, @@ -522,3 +524,62 @@ def pretrain_cfg_container_from_args(args: Namespace, model_cfg=None) -> Pretrai ) return cfg + + +def inference_cfg_from_args(args: Namespace) -> InferenceSetupConfig: + """Build an InferenceSetupConfig from the argparse arguments. + + InferenceSetupConfig field names map one-to-one onto the argparse ``dest`` names produced + by ``_add_inference_args``, so this is a direct copy of the relevant values from ``args``. + + This builds the declarative/serializable inference config. To obtain the runtime engine + config (``megatron.core.inference.config.InferenceConfig``), call + ``inference_cfg_from_args(args).to_inference_config(model, ...)``. + """ + return _default_config_from_args(InferenceSetupConfig, args) + + +def inference_cfg_container_from_args( + args: Namespace, model_cfg=None +) -> InferenceConfigContainer: + """Build an InferenceConfigContainer from the argparse arguments. + + This mirrors ``pretrain_cfg_container_from_args`` but assembles only the configs that + inference needs (no optimizer, scheduler, training, validation, DDP, rerun, or straggler + configs). It is intended to be passed to ``initialize_megatron`` from inference entry points. + + Args: + args: Parsed and validated argparse namespace (e.g. from ``parse_and_validate_args``). + model_cfg: Optional pre-built model config. If None, a model config is constructed from + ``args`` (a HybridModelConfig when ``--hybrid-layer-pattern`` is set, otherwise a + GPTModelConfig). + """ + if model_cfg is None: + if getattr(args, "hybrid_layer_pattern", None) is not None: + model_cfg = hybrid_config_from_args(args) + else: + model_cfg = gpt_config_from_args(args) + + ckpt_kwargs = _default_config_from_args(CheckpointConfig, args, return_instance=False) + ckpt_kwargs["save_optim"] = not args.no_save_optim + ckpt_kwargs["save_rng"] = not args.no_save_rng + ckpt_kwargs["load_optim"] = not args.no_load_optim + ckpt_kwargs["load_rng"] = not args.no_load_rng + ckpt_kwargs["fully_parallel_save"] = args.ckpt_fully_parallel_save + ckpt_kwargs["fully_parallel_load"] = args.ckpt_fully_parallel_load + + prof_kwargs = _default_config_from_args(ProfilingConfig, args, return_instance=False) + prof_kwargs["use_nsys_profiler"] = args.profile + + cfg = InferenceConfigContainer( + model=model_cfg, + checkpoint=CheckpointConfig(**ckpt_kwargs), + inference=inference_cfg_from_args(args), + dist=_default_config_from_args(DistributedInitConfig, args), + rng=_default_config_from_args(RNGConfig, args), + tokenizer=_default_config_from_args(TokenizerConfig, args), + logger=_default_config_from_args(LoggerConfig, args), + profiling=ProfilingConfig(**prof_kwargs), + ) + + return cfg diff --git a/megatron/training/config/__init__.py b/megatron/training/config/__init__.py index 4b8b67109e4..63e2c6ceaeb 100644 --- a/megatron/training/config/__init__.py +++ b/megatron/training/config/__init__.py @@ -18,6 +18,7 @@ RerunStateMachineConfig, StragglerDetectionConfig, ) +from megatron.training.config.inference_config import InferenceSetupConfig -from megatron.training.config.container import PretrainConfigContainer +from megatron.training.config.container import InferenceConfigContainer, PretrainConfigContainer from megatron.training.config.instantiate_utils import TargetAllowlist, target_allowlist diff --git a/megatron/training/config/container.py b/megatron/training/config/container.py index c13f73f52e9..7f4c882695e 100644 --- a/megatron/training/config/container.py +++ b/megatron/training/config/container.py @@ -12,6 +12,7 @@ from megatron.core.msc_utils import MultiStorageClientFeature from megatron.core.optimizer import OptimizerConfig from megatron.training.config.common_config import DistributedInitConfig, ProfilingConfig, RNGConfig +from megatron.training.config.inference_config import InferenceSetupConfig from megatron.training.config.instantiate_utils import InstantiationMode, instantiate from megatron.training.config.resilience_config import ( RerunStateMachineConfig, @@ -247,3 +248,35 @@ class PretrainConfigContainer(ConfigContainerBase): rerun_state_machine: RerunStateMachineConfig = field(default_factory=RerunStateMachineConfig) straggler: StragglerDetectionConfig | None = None + + +@dataclass(kw_only=True) +class InferenceConfigContainer(ConfigContainerBase): + """Top-level container for inference entry points. + + This is the inference counterpart to :class:`PretrainConfigContainer`. It holds only the + configs that inference actually needs and is intentionally shaped differently from the + training container: there is no optimizer, LR schedule, train/validation loop, DDP, rerun + state machine, or straggler detection. + + Explicitly NOT included (relative to ``PretrainConfigContainer``): ``TrainingConfig``, + ``OptimizerConfig``, ``SchedulerConfig``, ``ValidationConfig``, + ``DistributedDataParallelConfig``, ``RerunStateMachineConfig``, ``StragglerDetectionConfig``. + """ + + model: HybridModelConfig | GPTModelConfig + """Which model to load for inference.""" + + checkpoint: CheckpointConfig + """Checkpoint configuration used to load model weights.""" + + inference: InferenceSetupConfig + """Declarative inference settings (the serializable, args-shaped layer). Use + ``InferenceSetupConfig.to_inference_config(model, ...)`` to build the runtime + ``megatron.core.inference.config.InferenceConfig`` consumed by the engine.""" + + dist: DistributedInitConfig = field(default_factory=DistributedInitConfig) + rng: RNGConfig = field(default_factory=RNGConfig) + tokenizer: TokenizerConfig = field(default_factory=TokenizerConfig) + logger: LoggerConfig = field(default_factory=LoggerConfig) + profiling: ProfilingConfig = field(default_factory=ProfilingConfig) diff --git a/megatron/training/config/inference_config.py b/megatron/training/config/inference_config.py new file mode 100644 index 00000000000..782d8a8ab9b --- /dev/null +++ b/megatron/training/config/inference_config.py @@ -0,0 +1,350 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +"""Declarative configuration dataclass for Megatron inference entry points. + +This module defines :class:`InferenceSetupConfig`, the inference counterpart to the +training-oriented config dataclasses (e.g. ``TrainingConfig``, ``OptimizerConfig``). It +holds the inference-specific knobs that today live as loose ``args.`` values +produced by ``_add_inference_args`` in ``megatron.training.arguments``. Field names mirror +the corresponding argparse ``dest`` names one-to-one, so an ``InferenceSetupConfig`` can be +built directly from an ``argparse.Namespace`` via ``_default_config_from_args``. + +Layering note +------------- +``InferenceSetupConfig`` is the *declarative, serializable* layer (primitives/strings, safe +to YAML-serialize, built from args before the model or distributed groups exist). It is the +counterpart to ``megatron.training.models.GPTModelConfig``. + +The *runtime engine* config consumed by the inference context/engine is +``megatron.core.inference.config.InferenceConfig`` -- it holds rich runtime objects +(``ProcessGroupCollection``, ``MambaInferenceStateConfig``, ``torch.dtype``, a wandb module) +and can only be built once the model and process groups exist. + +Use :meth:`InferenceSetupConfig.to_inference_config` to produce the runtime engine config +from this declarative config plus the runtime artifacts. This mirrors the +``GPTModelConfig -> TransformerConfig`` relationship. +""" +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Literal + +if TYPE_CHECKING: + from megatron.core.inference.config import InferenceConfig + from megatron.core.transformer.module import MegatronModule + + +@dataclass(kw_only=True) +class InferenceSetupConfig: + """Declarative configuration settings for inference engines and the dynamic context. + + These fields correspond to the ``inference`` argument group defined by + ``_add_inference_args`` in ``megatron/training/arguments.py``. They cover both + the static and dynamic inference engines, the KV-cache memory buffer, CUDA graph + capture during decode, prefix caching, and inference-time logging. + + This is the serializable, args-shaped layer. The runtime engine config consumed by + the inference context/engine is ``megatron.core.inference.config.InferenceConfig``; + build it via :meth:`to_inference_config`. + """ + + # ---------------- General inference settings ---------------- + + inference_batch_times_seqlen_threshold: int = -1 + """If (batch-size * sequence-length) is smaller than this threshold then batches will not be + split up for pipelining. Requires setting --pipeline-model-parallel-size > 1. Setting this to + -1 indicates that batch pipelining is not used.""" + + max_tokens_to_oom: int = 12000 + """Maximum number of tokens during inference (# in prompt + # to generate). Allows us to throw + an error before OOM crashes server.""" + + output_bert_embeddings: bool = False + """Output Bert embeddings (via mean pooling) from model, rather than its binary head output or + entire hidden batch.""" + + bert_embedder_type: Literal["megatron", "huggingface"] = "megatron" + """Select either Megatron or Huggingface as the Bert embedder.""" + + cuda_graph_modules: list[str] = field(default_factory=list) + """Selects capture coverage within per-layer CUDA graphs (local and transformer_engine + implementations). An empty list means capturing the whole Transformer layer.""" + + use_legacy_static_engine: bool = False + """Use legacy static engine. (Current static engine uses dynamic engine under the hood.)""" + + inference_max_requests: int = 8 + """Maximum number of requests for inference.""" + + inference_max_seq_length: int = 2560 + """Maximum sequence length expected for inference (prefill + decode).""" + + # ---------------- Dynamic batching ---------------- + + inference_dynamic_batching: bool = False + """Enable dynamic batching mode.""" + + inference_dynamic_batching_buffer_size_gb: float = 40.0 + """Amount of on-GPU memory allocated for the KV cache. The total amount of memory allocated for + the KV cache (CPU + GPU memory) depends on the value set for the unified virtual memory (UVM) + level (via inference_dynamic_batching_unified_memory_level).""" + + inference_dynamic_batching_paused_buffer_size_gb: float | None = None + """Amount of memory reserved for paused requests in the dynamic inference context. Active + requests are paused when there are not enough active blocks available to continue generating a + request.""" + + inference_dynamic_batching_mamba_memory_ratio: float | None = None + """Percentage of memory buffer to allocate for Mamba states. If not specified, allocates Mamba + state tensors for each KV cache block. Only used for hybrid models.""" + + inference_dynamic_batching_block_size: int = 256 + """KV cache block size. It should be a multiple of 256.""" + + inference_dynamic_batching_max_requests: int | None = None + """Override the inference context's `max_requests`. By default, `max_requests` is set to the + number of blocks in the context's memory buffer.""" + + inference_dynamic_batching_max_tokens: int | None = None + """Override the inference context's default `max_tokens`.""" + + inference_dynamic_batching_num_cuda_graphs: int = 16 + """Maximum number of cuda graphs to capture, where the cuda graph batch sizes range from 1 to + `max_requests`. The user can also pass -1, in which case we automatically determine the number + of graphs to capture based on the `max_requests`.""" + + inference_dynamic_batching_track_paused_request_events: bool = False + """Track paused request ids by adding 'paused' events to each request's event history. This has + a very minor impact on latency.""" + + inference_dynamic_batching_track_generated_token_events: bool = False + """Track per-token events with timestamps for each generated token. When enabled, each generated + token creates a GENERATED_TOKEN event with a timestamp, useful for per-token latency analysis.""" + + inference_dynamic_batching_unified_memory_level: Literal[0, 1] = 0 + """Set unified memory usage within the dynamic inference context. The levels are: 0) no unified + memory, 1) allocate `memory_buffer` in unified memory.""" + + inference_dynamic_batching_cuda_graph_mixed_prefill_count: int = 16 + """Number of mixed prefill requests to capture in a cuda graph.""" + + inference_dynamic_batching_cuda_graph_sizing_distribution: Literal["exponential", "linear"] = ( + "exponential" + ) + """Spacing of CUDA graph token counts. "exponential" (default) halves from cuda_graph_max_tokens + down to tp_size, giving a log-spaced distribution with bounded relative padding. "linear" uses + varying linear strides across the range.""" + + inference_dynamic_batching_sampling_backend: Literal["torch", "flashinfer"] = "torch" + """Which sampling kernels to use during inference. Falls back to "torch" with a warning if + "flashinfer" is requested but the package is not installed.""" + + # ---------------- CUDA graphs ---------------- + + decode_only_cuda_graphs: bool = False + """Only use cuda graphs for decode-only steps, not prefill and mixed steps.""" + + inference_cuda_graph_all_prefills: bool = False + """Extend prefill/mixed CUDA graph capture up to `max_tokens`. By default, all graphs are + limited by the decode limit of `max_requests * (num_speculative_tokens + 1)`.""" + + # ---------------- Chunked prefill / speculation ---------------- + + enable_chunked_prefill: bool = False + """Enable chunked prefill (disabled by default).""" + + num_speculative_tokens: int = 0 + """Number of speculative tokens generated during decode.""" + + # ---------------- Prefix caching ---------------- + + inference_dynamic_batching_enable_prefix_caching: bool = False + """Enable/disable prefix caching for dynamic batching inference. When disabled, KV cache blocks + cannot be shared between requests with identical prompt prefixes.""" + + inference_dynamic_batching_prefix_caching_eviction_policy: Literal["ref_zero", "lru"] = "ref_zero" + """Eviction policy for prefix caching blocks. "ref_zero" (default) immediately returns blocks to + the free pool when ref_count hits 0. "lru" keeps blocks cached and evicts via LRU only when + space is needed.""" + + inference_dynamic_batching_prefix_caching_coordinator_policy: Literal[ + "longest_prefix", "first_prefix_block", "round_robin" + ] = "first_prefix_block" + """Coordinator routing policy for prefix caching. "first_prefix_block" (default) routes based on + the first block hash only. "longest_prefix" routes to the rank with the longest matching prefix. + "round_robin" ignores prefix affinity and cycles through ranks.""" + + inference_dynamic_batching_prefix_caching_routing_alpha: float = 0.5 + """Weight for prefix-aware routing score: score = alpha * match + (1 - alpha) * normalized_load. + Higher alpha favors prefix cache hits; lower alpha favors load balance.""" + + inference_dynamic_batching_prefix_caching_mamba_gb: float | None = None + """GPU memory budget (in GB) for the Mamba state cache used by prefix caching on hybrid models. + When set, Mamba states at block boundaries are cached for reuse.""" + + # ---------------- Logging ---------------- + + inference_logging_step_interval: int = 0 + """Step interval for logging inference metrics. Default to 0 to disable inference logging.""" + + inference_text_gen_server_logging: bool = False + """Enable per-request logging in the inference text generation server.""" + + inference_wandb_logging: bool = False + """Enable inference wandb logging.""" + + # ---------------- Coordinator / distributed ---------------- + + inference_coordinator_port: int | None = None + """This port will be used to setup the inference coordinator on node-0.""" + + inference_use_synchronous_zmq_collectives: bool = False + """Use synchronous ZMQ collectives for inference. Helps in reducing performance variability for + MoEs.""" + + inference_disable_ep_consensus: bool = False + """Skip the EP-group consensus all-reduce in the inference engine control loop and step on local + state only. Only safe when EP coordination is not required (e.g. ep_world_size == 1).""" + + # ---------------- Mamba inference state dtypes ---------------- + # NOTE: These are provided on the CLI as strings ("bf16"/"fp16"/"fp32") but are mapped to the + # corresponding torch dtype during argument validation (see validate_args in arguments.py). + + mamba_inference_conv_states_dtype: Literal["bf16", "fp16", "fp32"] = "bf16" + """Dtype for the Mamba inference conv states tensor.""" + + mamba_inference_ssm_states_dtype: Literal["bf16", "fp16", "fp32"] = "bf16" + """Dtype for the Mamba inference SSM states tensor.""" + + def to_inference_config( + self, + model: "MegatronModule", + *, + pg_collection: Any = None, + return_log_probs: bool = False, + skip_prompt_log_probs: bool = False, + kv_cache_management_mode: str = "persist", + static_kv_memory_pointers: bool = False, + use_flashinfer_fused_rope: bool = False, + enable_cuda_graphs: bool = True, + metrics_writer: Any = None, + verbose: bool = True, + ) -> "InferenceConfig": + """Build the runtime ``megatron.core.inference.config.InferenceConfig`` from this config. + + This is the bridge from the declarative inference settings to the runtime engine + config consumed by the dynamic inference context/engine. It supplies the fields that + depend on the built model (max sequence length, Mamba state config, process groups) + and the cross-cutting values that do not live on this declarative config. + + Args: + model: The (possibly wrapped) model to run inference with. Used to derive the + effective max sequence length, the Mamba inference state config, and the + process group collection when ``pg_collection`` is not provided. + pg_collection: Process groups for distributed execution. Defaults to the + model's ``pg_collection`` attribute when None. + return_log_probs: Whether generation returns log probs. Controls + ``materialize_only_last_token_logits`` (the engine must materialize all + logits when log probs are requested). + skip_prompt_log_probs: Whether prompt log probs are skipped. When True, only the + last token's logits are needed even if ``return_log_probs`` is True, so + ``materialize_only_last_token_logits`` is True. Sourced from the CLI arg + ``--skip-prompt-log-probs`` at the call site. + kv_cache_management_mode: How large tensors are handled on suspend/resume + ("persist"/"offload"/"recompute"). Sourced from the RL arg + ``rl_kv_cache_management_mode`` at the call site. + static_kv_memory_pointers: Whether the KV cache stays at fixed addresses across + suspend/resume. Sourced from the RL arg ``rl_persist_cuda_graphs``. + use_flashinfer_fused_rope: Whether to use flashinfer's fused rope implementation. + enable_cuda_graphs: When False, ``num_cuda_graphs`` is forced to None (no capture). + Callers typically pass ``inference_cuda_graph_scope != none``. + metrics_writer: Optional wandb module for inference metric logging. + verbose: Whether the context logs detailed configuration at initialization. + + Returns: + A fully-populated runtime ``InferenceConfig``. + """ + from megatron.core.inference.config import ( + CudaGraphSizingDistribution, + InferenceConfig, + KVCacheManagementMode, + MambaInferenceStateConfig, + PrefixCachingCoordinatorPolicy, + PrefixCachingEvictionPolicy, + ) + from megatron.core.utils import get_attr_wrapped_model + + # Effective max sequence length depends on the model's position embedding type. + position_embedding_type = get_attr_wrapped_model(model, "position_embedding_type") + model_max_seq_len = get_attr_wrapped_model(model, "max_sequence_length") + inf_max_seq_len = self.inference_max_seq_length + max_batch_size = self.inference_dynamic_batching_max_requests + + if position_embedding_type == "learned_absolute": + # The context's max_sequence_length must not exceed the model's, otherwise the + # context's position_ids index past the position embedding table. + if inf_max_seq_len: + max_sequence_length = min(model_max_seq_len, inf_max_seq_len) + else: + max_sequence_length = model_max_seq_len + assert max_batch_size is None or max_batch_size <= model_max_seq_len + else: + max_sequence_length = inf_max_seq_len + if max_batch_size is not None: + max_sequence_length = max(max_sequence_length, max_batch_size) + + mamba_inference_state_config = MambaInferenceStateConfig.from_model( + model, + conv_states_dtype=self.mamba_inference_conv_states_dtype, + ssm_states_dtype=self.mamba_inference_ssm_states_dtype, + ) + if pg_collection is None: + pg_collection = get_attr_wrapped_model(model, "pg_collection") + + return InferenceConfig( + verbose=verbose, + block_size_tokens=self.inference_dynamic_batching_block_size, + buffer_size_gb=self.inference_dynamic_batching_buffer_size_gb, + paused_buffer_size_gb=self.inference_dynamic_batching_paused_buffer_size_gb, + mamba_memory_ratio=self.inference_dynamic_batching_mamba_memory_ratio, + num_cuda_graphs=( + self.inference_dynamic_batching_num_cuda_graphs if enable_cuda_graphs else None + ), + max_requests=self.inference_dynamic_batching_max_requests, + max_tokens=self.inference_dynamic_batching_max_tokens, + unified_memory_level=self.inference_dynamic_batching_unified_memory_level, + kv_cache_management_mode=KVCacheManagementMode(kv_cache_management_mode), + cuda_graph_mixed_prefill_count=( + self.inference_dynamic_batching_cuda_graph_mixed_prefill_count + ), + cuda_graph_sizing_distribution=CudaGraphSizingDistribution( + self.inference_dynamic_batching_cuda_graph_sizing_distribution + ), + use_cuda_graphs_for_non_decode_steps=not self.decode_only_cuda_graphs, + cuda_graph_all_prefills=self.inference_cuda_graph_all_prefills, + static_kv_memory_pointers=static_kv_memory_pointers, + max_sequence_length=max_sequence_length, + mamba_inference_state_config=mamba_inference_state_config, + pg_collection=pg_collection, + use_flashinfer_fused_rope=use_flashinfer_fused_rope, + materialize_only_last_token_logits=( + not (return_log_probs and not skip_prompt_log_probs) + ), + track_generated_token_events=( + self.inference_dynamic_batching_track_generated_token_events + ), + track_paused_request_events=self.inference_dynamic_batching_track_paused_request_events, + enable_chunked_prefill=self.enable_chunked_prefill, + enable_prefix_caching=self.inference_dynamic_batching_enable_prefix_caching, + prefix_caching_eviction_policy=PrefixCachingEvictionPolicy( + self.inference_dynamic_batching_prefix_caching_eviction_policy + ), + prefix_caching_coordinator_policy=PrefixCachingCoordinatorPolicy( + self.inference_dynamic_batching_prefix_caching_coordinator_policy + ), + prefix_caching_routing_alpha=self.inference_dynamic_batching_prefix_caching_routing_alpha, + prefix_caching_mamba_gb=self.inference_dynamic_batching_prefix_caching_mamba_gb, + metrics_writer=metrics_writer, + logging_step_interval=self.inference_logging_step_interval, + num_speculative_tokens=self.num_speculative_tokens, + use_synchronous_zmq_collectives=self.inference_use_synchronous_zmq_collectives, + disable_ep_consensus=self.inference_disable_ep_consensus, + sampling_backend=self.inference_dynamic_batching_sampling_backend, + ) From edbb694ecba979da0e8b4e8c4b236d04bd66702c Mon Sep 17 00:00:00 2001 From: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> Date: Tue, 23 Jun 2026 09:38:09 -0700 Subject: [PATCH 4/5] Fold inference-arg-group kwargs into InferenceSetupConfig Address review: InferenceSetupConfig is meant to mirror _add_inference_args 1:1, but return_log_probs, skip_prompt_log_probs, and use_flashinfer_fused_rope were still passed as explicit kwargs to to_inference_config. Move them onto the dataclass and read from self.* inside to_inference_config. The remaining explicit kwargs are intentional: kv_cache_management_mode and static_kv_memory_pointers come from the RL arg group (--rl-*), and enable_cuda_graphs is a derived boolean from args.inference_cuda_graph_scope (also not in the inference group). --- megatron/inference/utils.py | 7 ++-- megatron/training/config/inference_config.py | 39 ++++++++++++-------- 2 files changed, 27 insertions(+), 19 deletions(-) diff --git a/megatron/inference/utils.py b/megatron/inference/utils.py index d6d27df1a38..9eaf27af599 100644 --- a/megatron/inference/utils.py +++ b/megatron/inference/utils.py @@ -344,15 +344,16 @@ def get_inference_config_from_model_and_args(model: MegatronModule, args): "wandb module is available. Inference logging will be disabled.", ) + # Only kwargs that are NOT in the inference argument group are passed explicitly. + # The rest (return_log_probs, skip_prompt_log_probs, use_flashinfer_fused_rope, and the + # inference_dynamic_batching_* / prefix_caching_* / chunked_prefill knobs) live on + # InferenceSetupConfig and are read from ``self`` inside ``to_inference_config``. setup_cfg = inference_cfg_from_args(args) return setup_cfg.to_inference_config( model, - return_log_probs=args.return_log_probs, kv_cache_management_mode=args.rl_kv_cache_management_mode, static_kv_memory_pointers=args.rl_persist_cuda_graphs, - use_flashinfer_fused_rope=args.use_flashinfer_fused_rope, enable_cuda_graphs=(args.inference_cuda_graph_scope != InferenceCudaGraphScope.none), - skip_prompt_log_probs=args.skip_prompt_log_probs, metrics_writer=metrics_writer, ) diff --git a/megatron/training/config/inference_config.py b/megatron/training/config/inference_config.py index 782d8a8ab9b..29d5c848049 100644 --- a/megatron/training/config/inference_config.py +++ b/megatron/training/config/inference_config.py @@ -213,16 +213,28 @@ class InferenceSetupConfig: mamba_inference_ssm_states_dtype: Literal["bf16", "fp16", "fp32"] = "bf16" """Dtype for the Mamba inference SSM states tensor.""" + # ---------------- Log-prob and RoPE knobs from _add_inference_args ---------------- + + return_log_probs: bool = False + """Return the log probabilities of the final output tokens. Mirrors ``--return-log-probs``. + Controls ``materialize_only_last_token_logits`` (the engine must materialize all logits when + log probs are requested, unless ``skip_prompt_log_probs`` is also True).""" + + skip_prompt_log_probs: bool = False + """Skip prompt log probs. Mirrors ``--skip-prompt-log-probs``. When True, only the last + token's logits are needed even if ``return_log_probs`` is True, so + ``materialize_only_last_token_logits`` stays True.""" + + use_flashinfer_fused_rope: bool = False + """Use flashinfer's fused rope implementation. Mirrors ``--use-flashinfer-fused-rope``.""" + def to_inference_config( self, model: "MegatronModule", *, pg_collection: Any = None, - return_log_probs: bool = False, - skip_prompt_log_probs: bool = False, kv_cache_management_mode: str = "persist", static_kv_memory_pointers: bool = False, - use_flashinfer_fused_rope: bool = False, enable_cuda_graphs: bool = True, metrics_writer: Any = None, verbose: bool = True, @@ -240,21 +252,16 @@ def to_inference_config( process group collection when ``pg_collection`` is not provided. pg_collection: Process groups for distributed execution. Defaults to the model's ``pg_collection`` attribute when None. - return_log_probs: Whether generation returns log probs. Controls - ``materialize_only_last_token_logits`` (the engine must materialize all - logits when log probs are requested). - skip_prompt_log_probs: Whether prompt log probs are skipped. When True, only the - last token's logits are needed even if ``return_log_probs`` is True, so - ``materialize_only_last_token_logits`` is True. Sourced from the CLI arg - ``--skip-prompt-log-probs`` at the call site. kv_cache_management_mode: How large tensors are handled on suspend/resume ("persist"/"offload"/"recompute"). Sourced from the RL arg - ``rl_kv_cache_management_mode`` at the call site. + ``rl_kv_cache_management_mode`` at the call site (not part of the inference + argument group, so it is not folded into this dataclass). static_kv_memory_pointers: Whether the KV cache stays at fixed addresses across - suspend/resume. Sourced from the RL arg ``rl_persist_cuda_graphs``. - use_flashinfer_fused_rope: Whether to use flashinfer's fused rope implementation. + suspend/resume. Sourced from the RL arg ``rl_persist_cuda_graphs`` (not part + of the inference argument group). enable_cuda_graphs: When False, ``num_cuda_graphs`` is forced to None (no capture). - Callers typically pass ``inference_cuda_graph_scope != none``. + Callers typically pass ``inference_cuda_graph_scope != none``; derived, not a + 1:1 args field. metrics_writer: Optional wandb module for inference metric logging. verbose: Whether the context logs detailed configuration at initialization. @@ -323,9 +330,9 @@ def to_inference_config( max_sequence_length=max_sequence_length, mamba_inference_state_config=mamba_inference_state_config, pg_collection=pg_collection, - use_flashinfer_fused_rope=use_flashinfer_fused_rope, + use_flashinfer_fused_rope=self.use_flashinfer_fused_rope, materialize_only_last_token_logits=( - not (return_log_probs and not skip_prompt_log_probs) + not (self.return_log_probs and not self.skip_prompt_log_probs) ), track_generated_token_events=( self.inference_dynamic_batching_track_generated_token_events From e83fac93f3f02d8eec66abe1a77b90c1a78d8a27 Mon Sep 17 00:00:00 2001 From: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> Date: Tue, 23 Jun 2026 14:47:17 -0700 Subject: [PATCH 5/5] Addressed keshavs comments --- megatron/inference/utils.py | 4 ---- megatron/training/config/inference_config.py | 3 +-- 2 files changed, 1 insertion(+), 6 deletions(-) diff --git a/megatron/inference/utils.py b/megatron/inference/utils.py index 9eaf27af599..00b931d5eab 100644 --- a/megatron/inference/utils.py +++ b/megatron/inference/utils.py @@ -344,10 +344,6 @@ def get_inference_config_from_model_and_args(model: MegatronModule, args): "wandb module is available. Inference logging will be disabled.", ) - # Only kwargs that are NOT in the inference argument group are passed explicitly. - # The rest (return_log_probs, skip_prompt_log_probs, use_flashinfer_fused_rope, and the - # inference_dynamic_batching_* / prefix_caching_* / chunked_prefill knobs) live on - # InferenceSetupConfig and are read from ``self`` inside ``to_inference_config``. setup_cfg = inference_cfg_from_args(args) return setup_cfg.to_inference_config( model, diff --git a/megatron/training/config/inference_config.py b/megatron/training/config/inference_config.py index 29d5c848049..e8391b70e8f 100644 --- a/megatron/training/config/inference_config.py +++ b/megatron/training/config/inference_config.py @@ -254,8 +254,7 @@ def to_inference_config( model's ``pg_collection`` attribute when None. kv_cache_management_mode: How large tensors are handled on suspend/resume ("persist"/"offload"/"recompute"). Sourced from the RL arg - ``rl_kv_cache_management_mode`` at the call site (not part of the inference - argument group, so it is not folded into this dataclass). + ``rl_kv_cache_management_mode`` at the call site. static_kv_memory_pointers: Whether the KV cache stays at fixed addresses across suspend/resume. Sourced from the RL arg ``rl_persist_cuda_graphs`` (not part of the inference argument group).