diff --git a/megatron/inference/utils.py b/megatron/inference/utils.py index 567d48ffc3b..00b931d5eab 100644 --- a/megatron/inference/utils.py +++ b/megatron/inference/utils.py @@ -1,21 +1,12 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging -from argparse import ArgumentParser -from functools import partial -from typing import Optional +import warnings +from argparse import ArgumentParser, Namespace +from typing import Literal, Optional + import torch -from gpt_builders import gpt_builder -from hybrid_builders import hybrid_builder -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 ( @@ -25,41 +16,80 @@ 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 -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 +from megatron.training.argument_utils import gpt_config_from_args, hybrid_config_from_args from megatron.training.checkpointing import load_checkpoint -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_for_inference() -> MegatronModule: - """Initialize model and load checkpoint for inference.""" +def get_model_builder( + args: Namespace, provider: Optional[Literal["gpt", "hybrid", "mamba"]] = 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 get_model_for_inference() -> MegatronModule: + """Initialize model and load checkpoint for inference.""" - # Build model. - model = _get_model(partial(model_provider, model_builder), wrap_with_ddp=False) + args = get_args() + + 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 @@ -289,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() @@ -341,48 +344,13 @@ 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, + 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), 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, - logprobs_mode=args.inference_dynamic_batching_logprobs_mode, ) 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..12a424e377d --- /dev/null +++ b/megatron/training/config/inference_config.py @@ -0,0 +1,363 @@ +# 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.""" + + inference_dynamic_batching_logprobs_mode: Literal["raw_logprobs", "processed_logprobs"] = ( + "raw_logprobs" + ) + """How returned inference log-probs are computed engine-wide. "raw_logprobs" (default) uses the + unmodified model logits; "processed_logprobs" uses temperature and filters by top-k/top-p.""" + + # ---------------- 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.""" + + # ---------------- 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, + kv_cache_management_mode: str = "persist", + static_kv_memory_pointers: 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. + 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`` (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``; 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. + + 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=self.use_flashinfer_fused_rope, + materialize_only_last_token_logits=( + not (self.return_log_probs and not self.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, + logprobs_mode=self.inference_dynamic_batching_logprobs_mode, + ) 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..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 @@ -14,8 +13,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 @@ -28,10 +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 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)) @@ -39,8 +44,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 +142,18 @@ 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 + 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: - 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) + 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)