Addressing new inference updates in mcore (#706)

Signed-off-by: Onur Yilmaz <oyilmaz@nvidia.com>
Signed-off-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>
Co-authored-by: nemo-autobot-origin[bot] <297562586+nemo-autobot-origin[bot]@users.noreply.github.com>

Author: 5 people

nemo_deploy/deploy_ray.py

@@ -30,7 +30,7 @@
     from nemo_deploy.multimodal.megatron_multimodal_deployable_ray import MegatronMultimodalRayDeployable
 
     HAVE_RAY = True
-except (ImportError, ModuleNotFoundError):
+except (ImportError, ModuleNotFoundError, AssertionError):
     from unittest.mock import MagicMock
 
     ray = MagicMock()

nemo_deploy/llm/inference/inference_base.py

@@ -27,14 +27,8 @@
     get_default_load_sharded_strategy,
 )
 from megatron.core.dist_checkpointing.validation import StrictHandling
-from megatron.core.inference.contexts.static_context import StaticInferenceContext
-from megatron.core.inference.engines.mcore_engine import MCoreEngine
-from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
-    GPTInferenceWrapper,
-)
-from megatron.core.inference.text_generation_controllers.text_generation_controller import (
-    TextGenerationController,
-)
+from megatron.core.inference.apis import MegatronLLM
+from megatron.core.inference.config import InferenceConfig
 from megatron.core.transformer.enums import AttnBackend
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.transformer_config import MLATransformerConfig
@@ -390,27 +384,24 @@ def setup_model_and_tokenizer_for_inference(
 
 
 class MCoreEngineWithCleanup:
-    """Wrapper around MCoreEngine that ensures proper cleanup of distributed resources.
+    """Wrapper around MegatronLLM that ensures proper cleanup of distributed resources.
 
-    This class delegates all operations to the underlying MCoreEngine while ensuring that
-    distributed resources are properly cleaned up when the engine is destroyed.
+    This class delegates all operations to the underlying MegatronLLM engine while ensuring
+    that distributed resources are properly cleaned up when the engine is destroyed.
     """
 
     def __init__(
         self,
-        mcore_engine: MCoreEngine,
-        model_inference_wrapper: GPTInferenceWrapper,
+        llm: MegatronLLM,
         tokenizer: Union[MCoreTokenizerWrappper, MegatronTokenizer],
     ):
         """Initialize the MCoreEngineWithCleanup.
 
         Args:
-            mcore_engine (MCoreEngine): The underlying MCoreEngine instance
-            model_inference_wrapper (GPTInferenceWrapper): The model inference wrapper
+            llm (MegatronLLM): The underlying MegatronLLM instance
             tokenizer (Union[MCoreTokenizerWrappper, MegatronTokenizer]): The tokenizer instance
         """
-        self.mcore_engine = mcore_engine
-        self.model_inference_wrapper = model_inference_wrapper
+        self.mcore_engine = llm
         self.tokenizer = tokenizer
 
     def __del__(self):
@@ -446,16 +437,16 @@ def create_mcore_engine(
     buffer_size_gb: float = 10.0,
     legacy_model_format: bool = False,
     **model_config_kwargs,
-) -> Tuple[MCoreEngineWithCleanup, GPTInferenceWrapper, Union[MCoreTokenizerWrappper, MegatronTokenizer]]:
-    """Set up the model, tokenizer and MCoreEngine for inference.
+) -> Tuple[MCoreEngineWithCleanup, Union[MCoreTokenizerWrappper, MegatronTokenizer]]:
+    """Set up the model, tokenizer and MegatronLLM engine for inference.
 
     Args:
         path (Path): Path to the checkpoint file
         params_dtype (torch.dtype): Data type for model parameters (default: torch.bfloat16)
         inference_batch_times_seqlen_threshold (int): Threshold for batch size times sequence length
         inference_max_seq_length (int): Maximum sequence length for inference
         max_batch_size (int): Maximum batch size for inference
-        random_seed (Optional[int]): Random seed for reproducibility
+        random_seed (Optional[int]): Random seed for reproducibility (set globally during init)
         tensor_model_parallel_size (Optional[int]): Size of tensor model parallelism
         pipeline_model_parallel_size (Optional[int]): Size of pipeline model parallelism
         context_parallel_size (Optional[int]): Size of context parallelism
@@ -466,11 +457,10 @@ def create_mcore_engine(
         model_type (str): Type of model to load (default: "gpt")
         model_format (str): Format of model to load (default: "nemo")
         micro_batch_size (Optional[int]): Micro batch size for model execution
-        legacy_model_format (bool): Whether to use the legacy StaticInferenceEngine path in MCoreEngine (default: False)
+        legacy_model_format (bool): Deprecated; no longer used (DynamicInferenceEngine is always used)
     Returns:
-        Tuple[MCoreEngineWithCleanup, GPTInferenceWrapper, Union[MCoreTokenizerWrappper, MegatronTokenizer]]: Tuple containing:
+        Tuple[MCoreEngineWithCleanup, Union[MCoreTokenizerWrappper, MegatronTokenizer]]: Tuple containing:
             - MCoreEngineWithCleanup: Engine for text generation with proper cleanup
-            - GPTInferenceWrapper: Inference-wrapped model
             - Union[MCoreTokenizerWrappper, MegatronTokenizer]: Tokenizer instance
     """
     # Default to 1 for any parallelism dimension that's None
@@ -512,34 +502,32 @@ def create_mcore_engine(
     else:
         raise ValueError(f"Model format {model_format} not supported.")
 
-    # MLA models require block_size_tokens=64 for the dynamic engine, which is not
-    # configurable in the current Megatron-LM version. Fall back to the legacy static
-    # engine so MLA inference works correctly without touching Megatron-LM.
+    model.eval()
+
+    # MLA models require block_size_tokens=64 for correct KV cache operation with the
+    # dynamic inference engine. Set the attention backend to flash if not already set.
+    block_size_tokens = 256
     model_config = getattr(model, "config", None)
     if isinstance(model_config, MLATransformerConfig):
-        legacy_model_format = True
-        # The legacy static engine requires an explicit attention backend.
-        # MLA models use flash attention (attention_mask is handled internally).
+        block_size_tokens = 64
         if not model_config.attention_backend:
             model_config.attention_backend = AttnBackend.flash
 
-    inference_context = StaticInferenceContext(
-        max_batch_size=max_batch_size,
+    inference_config = InferenceConfig(
         max_sequence_length=inference_max_seq_length,
+        buffer_size_gb=int(buffer_size_gb),
+        max_requests=max_batch_size,
+        block_size_tokens=block_size_tokens,
+        materialize_only_last_token_logits=True,
     )
-    model_inference_wrapper = GPTInferenceWrapper(model, inference_context)
-    text_generation_controller = TextGenerationController(
-        inference_wrapped_model=model_inference_wrapper, tokenizer=tokenizer
-    )
-    mcore_engine = MCoreEngine(
-        text_generation_controller=text_generation_controller,
-        max_batch_size=max_batch_size,
-        random_seed=random_seed,
-        buffer_size_gb=buffer_size_gb,
-        legacy=legacy_model_format,
+
+    llm = MegatronLLM(
+        model=model,
+        tokenizer=tokenizer,
+        inference_config=inference_config,
     )
 
     # Wrap the engine to ensure cleanup
-    wrapped_engine = MCoreEngineWithCleanup(mcore_engine, model_inference_wrapper, tokenizer)
+    wrapped_engine = MCoreEngineWithCleanup(llm, tokenizer)
 
-    return wrapped_engine, model_inference_wrapper, tokenizer
+    return wrapped_engine, tokenizer

nemo_deploy/llm/megatronllm_deployable.py

@@ -21,8 +21,8 @@
 import torch
 import torch.distributed
 from jinja2 import Template
-from megatron.core.inference.common_inference_params import CommonInferenceParams
-from megatron.core.inference.inference_request import InferenceRequest
+from megatron.core.inference.inference_request import DynamicInferenceRequest
+from megatron.core.inference.sampling_params import SamplingParams
 
 from nemo_deploy import ITritonDeployable
 from nemo_deploy.llm.inference.inference_base import create_mcore_engine
@@ -113,7 +113,7 @@ def __init__(
         if model_type not in ["gpt", "mamba"]:
             raise ValueError(f"Model type {model_type} not supported for Megatron models.")
 
-        self.mcore_engine, self.inference_wrapped_model, self.mcore_tokenizer = create_mcore_engine(
+        self.mcore_engine, self.mcore_tokenizer = create_mcore_engine(
             num_devices=num_devices,
             num_nodes=num_nodes,
             path=Path(megatron_checkpoint_filepath),
@@ -144,18 +144,18 @@ def __init__(
     def generate(
         self,
         prompts: List[str],
-        inference_params: Optional[CommonInferenceParams] = None,
-    ) -> List[InferenceRequest]:
+        inference_params: Optional[SamplingParams] = None,
+    ) -> List[DynamicInferenceRequest]:
         """Generates text based on the provided input prompts.
 
         Args:
             prompts (List[str]): A list of input strings.
-            inference_params (Optional[CommonInferenceParams]): Parameters for controlling the inference process.
+            inference_params (Optional[SamplingParams]): Parameters for controlling the inference process.
 
         Returns:
-            List[InferenceRequest]: A list containing the generated results.
+            List[DynamicInferenceRequest]: A list containing the generated results.
         """
-        inference_params = inference_params or CommonInferenceParams()
+        inference_params = inference_params or SamplingParams()
 
         # Store the original number of prompts
         orig_num_prompts = len(prompts)
@@ -173,17 +173,15 @@ def generate(
 
             results = self.mcore_engine.generate(
                 prompts=padded_prompts,
-                add_BOS=False,
-                common_inference_params=inference_params,
+                sampling_params=inference_params,
             )
 
             # Only return results for the original prompts
             return list(results)[:orig_num_prompts]
         else:
             results = self.mcore_engine.generate(
                 prompts=prompts,
-                add_BOS=False,
-                common_inference_params=inference_params,
+                sampling_params=inference_params,
             )
             return list(results)
 
@@ -198,7 +196,7 @@ def generate_other_ranks(self):
                     data=[None], src=0
                 )
 
-                inference_params = CommonInferenceParams(
+                inference_params = SamplingParams(
                     temperature=temperature,
                     top_k=int(top_k),
                     top_p=float(top_p),
@@ -208,7 +206,7 @@ def generate_other_ranks(self):
                 )
 
                 if log_probs:
-                    dynamic_engine = getattr(self.mcore_engine, "dynamic_engine", None)
+                    dynamic_engine = getattr(self.mcore_engine, "engine", None)
                     if dynamic_engine is not None:
                         dynamic_engine.materialize_only_last_token_logits = False
                         dynamic_engine.context.config.materialize_only_last_token_logits = False
@@ -419,15 +417,15 @@ def _infer_fn(
                 )
 
         # cast top_k,top_p to native int, float since typecheck assert statements added in MCore0.13 error otherwise
-        # return_prompt_top_n_logprobs returns top_logprobs for prompt tokens too when top_logprobs>0.
-        inference_params = CommonInferenceParams(
+        # skip_prompt_log_probs=False (default) includes prompt tokens in top-N logprobs when top_logprobs>0.
+        inference_params = SamplingParams(
             temperature=temperature,
             top_k=int(top_k),
             top_p=float(top_p),
             num_tokens_to_generate=num_tokens_to_generate,
             return_log_probs=log_probs,
             top_n_logprobs=top_logprobs,
-            return_prompt_top_n_logprobs=bool(top_logprobs),
+            skip_prompt_log_probs=not bool(top_logprobs),
             stop_words=stop_words,
         )
 
@@ -436,7 +434,7 @@ def _infer_fn(
         # (prompt log probs are required for logprob eval benchmarks).
         # Toggle it on both the engine and the context config (controls the
         # model forward pass and log prob calculations).
-        dynamic_engine = getattr(self.mcore_engine, "dynamic_engine", None)
+        dynamic_engine = getattr(self.mcore_engine, "engine", None)
         needs_all_logits = log_probs or bool(top_logprobs)
         if dynamic_engine is not None and needs_all_logits:
             dynamic_engine.materialize_only_last_token_logits = False

pyproject.toml

@@ -110,8 +110,11 @@ vllm = [
     { index = "pytorch-cu130", marker = "python_version < '3.9' and platform_machine == 'x86_64'" },
     { index = "pypi", marker = "platform_machine == 'aarch64'" },
 ]
-megatron-bridge = { git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git", rev = "859a18d001eca4423ad79e7771632d839f50fa77" }
-nvidia-resiliency-ext = { git = "https://github.com/NVIDIA/nvidia-resiliency-ext.git", rev = "b2bb3d728a18795807d9f76c535e005a609a1b01" }
+megatron-bridge = { git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git", rev = "97e8dba4060ae74a110b9da6255eb0a66b016b2d" }
+# Use the released nvidia-resiliency-ext (>=0.6.0) rather than a git dev build:
+# megatron-core's runtime is_nvrx_min_version() asserts the installed package
+# version is >=0.6.0, and a git build reports 0.6.0.devN (< 0.6.0), which fails
+# the assertion at import time.
 # nemo-toolkit = { git = "https://github.com/NVIDIA/NeMo.git", rev = "main" }
 
 [tool.uv]
@@ -123,7 +126,9 @@ no-build-isolation-package = [
     "mamba-ssm",
     "causal-conv1d",
     "nv-grouped-gemm",
+    "fast-hadamard-transform",
 ]
+
 # Always apply the build group since dependencies like TE/mcore/nemo-run require build dependencies
 # and this lets us assume they are implicitly installed with a simply `uv sync`. Ideally, we'd
 # avoid including these in the default dependency set, but for now it's required.
@@ -149,7 +154,7 @@ override-dependencies = [
     "datasets>=3.3.0",
     "flash-linear-attention>=0.3.0,<0.4.dev0",
     "patchelf; sys_platform=='never'",
-    "nvidia-resiliency-ext>=0.3.0,<0.6.0",
+    "nvidia-resiliency-ext>=0.6.0",
     "transformer-engine[pytorch,core_cu13]>=2.14.0a0,<2.17.0; sys_platform != 'darwin'",
     "transformer-engine-cu13>=2.14.0a0,<2.17.0; sys_platform != 'darwin'",
     "transformer-engine-cu12; sys_platform == 'never'",
@@ -162,6 +167,10 @@ override-dependencies = [
 ]
 prerelease = "allow"
 
+[[tool.uv.dependency-metadata]]
+name = "nvidia-resiliency-ext"
+version = "0.6.0"
+
 [[tool.uv.index]]
 name = "pypi"
 url = "https://pypi.org/simple"

tests/functional_tests/utils/run_nemo_export.py

@@ -35,13 +35,13 @@
 
 in_framework_supported = True
 try:
-    from megatron.core.inference.common_inference_params import CommonInferenceParams
+    from megatron.core.inference.sampling_params import SamplingParams
 
     from nemo_deploy.llm import NemoQueryLLMPyTorch
     from nemo_deploy.llm.megatronllm_deployable import MegatronLLMDeployable
 except Exception as e:
     LOGGER.warning(
-        "Cannot import MegatronLLMDeployable class, or NemoQueryLLMPyTorch, or CommonInferenceParams, "
+        "Cannot import MegatronLLMDeployable class, or NemoQueryLLMPyTorch, or SamplingParams, "
         f"in-framework inference will not be available. Reason: {type(e).__name__}: {e}"
     )
     in_framework_supported = False
@@ -98,7 +98,7 @@ def get_accuracy_with_lambada(model, nq, lora_uids, test_data_path, use_vllm: bo
                 if in_framework_supported and isinstance(model, MegatronLLMDeployable):
                     model_output = model.generate(
                         prompts=[prompt],
-                        inference_params=CommonInferenceParams(
+                        inference_params=SamplingParams(
                             temperature=0.1,
                             top_k=1,
                             top_p=0.0,