build: bump vllm 0.20.0 → 0.25.1

examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-1n8g-megatron-fp8-rollouts.v3.yaml

@@ -34,6 +34,13 @@ policy:
       precision: fp8
       max_model_len: 4096
       use_deep_gemm: true
+    vllm_kwargs:
+      # Empirical value that stops the CuMemAllocator wake-up OOM under vLLM
+      # 0.25 (takes precedence over gpu_memory_utilization). Not a 0.20-
+      # equivalent: 0.20 gave this engine 36.13 GiB, so this is a ~45% cut.
+      # Read https://github.com/NVIDIA-NeMo/RL/issues/3412 before retuning --
+      # this recipe's driver cannot detect a too-small pin.
+      kv_cache_memory_bytes: 21474836480  # 20 GiB
 data:
   max_input_seq_length: 4096
 logger:

examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron-fp8-e2e.yaml

@@ -43,6 +43,14 @@ policy:
       quantization_ignored_layer_kws:
       - a_proj
       - b_proj
+    vllm_kwargs:
+      # At the same gpu_memory_utilization, vLLM 0.25's profiling sizes the
+      # KV cache to ~34 GiB where 0.20 sized it to ~20 GiB, growing the
+      # CuMemAllocator sleep pool past what can be re-mapped at wake-up next
+      # to the colocated Megatron policy ("CUDA Error: out of memory at
+      # csrc/cumem_allocator.cpp"). Pin the KV cache to the 0.20-proven size
+      # (takes precedence over gpu_memory_utilization).
+      kv_cache_memory_bytes: 21474836480  # 20 GiB
 logger:
   monitor_gpus: false
   wandb:

nemo_rl/modelopt/models/generation/vllm_modelopt.py

@@ -17,13 +17,13 @@
 
 vLLM owns checkpoint-layout restoration, layerwise post-load processing,
 CUDA-graph-stable tensor placement, and KV-cache scale reload.  This module
-only supplies the vLLM 0.20 gaps needed here: ModelOpt W4A16 NVFP4 methods,
-rank-local Marlin padding, per-projection ModelOpt MoE input-scale loading,
-materialization of FlashInfer's global-scale views, and retention of
-method-owned MoE kernel references across layerwise reload.
+only supplies the vLLM 0.25 gaps needed here: ModelOpt W4A16 NVFP4 methods,
+per-projection ModelOpt MoE input-scale loading, materialization of
+FlashInfer's global-scale views, and retention of method-owned MoE kernel
+references across layerwise reload.  (Rank-local Marlin padding is gone: 0.25's
+``prepare_nvfp4_moe_layer_for_marlin`` pads natively.)
 """
 
-import copy
 from types import MethodType
 from typing import Any
 
@@ -58,12 +58,14 @@ def _load_modelopt_moe_input_scale(
 ) -> bool | None:
     """Load a ModelOpt input scale without losing the gate/up shard.
 
-    Replaces the ``input_scale`` branch of vLLM v0.20.0's
-    ``FusedMoE.weight_loader``, whose ``_load_single_value`` writes
-    ``param.data[expert_id]`` and drops the gate/up (w1/w3) shard index:
-    https://github.com/vllm-project/vllm/blob/v0.20.0/vllm/model_executor/layers/fused_moe/layer.py#L1025-L1031
-    Delete once upstream loads per-projection ModelOpt MoE input scales
-    correctly.
+    Replaces the ``input_scale`` branch of vLLM's MoE weight loader, which
+    drops the gate/up (w1/w3) shard index.  vLLM 0.25 handles this natively at
+    https://github.com/vllm-project/vllm/blob/v0.25.1/vllm/model_executor/layers/fused_moe/routed_experts.py#L699-L712
+    but this override is still required: it uses
+    ``shard_index = 0 if w1 else min(1, param.shape[-1] - 1)`` where upstream
+    hardcodes ``1``, so it also handles the single-scale (non-gated) layout,
+    and it copies with ``reshape_as`` rather than ``_to_scalar()``.
+    Delete once upstream covers the single-column case too.
     """
     del weight_name
     global_expert_id = expert_id
@@ -88,17 +90,14 @@ def _load_modelopt_moe_input_scale(
     return True if return_success else None
 
 
-def _normalized_w4a16_config(config: dict[str, Any]) -> dict[str, Any]:
-    normalized = copy.deepcopy(config)
-    quantization = normalized.get("quantization")
-    target = quantization if isinstance(quantization, dict) else normalized
+def _validated_w4a16_config(config: dict[str, Any]) -> dict[str, Any]:
+    quantization = config.get("quantization")
+    target = quantization if isinstance(quantization, dict) else config
     if str(target.get("quant_algo", "")).upper() != _W4A16_ALGO:
         raise ValueError(f"{NEMO_MODELOPT_W4A16} requires quant_algo={_W4A16_ALGO!r}")
-    # vLLM 0.20 validates known ModelOpt algorithms before dispatching to a
-    # custom subclass. Normalize only for its parser; class identity selects
-    # the W4A16 methods below.
-    target["quant_algo"] = _W4A4_ALGO
-    return normalized
+    # vLLM 0.25 understands W4A16_NVFP4 natively, so the algo passes through
+    # unchanged (the base __init__ keys use_a16/LinearMethodCls off it).
+    return config
 
 
 def _canonicalize_nvfp4_scale_(scale: torch.Tensor) -> None:
@@ -107,54 +106,6 @@ def _canonicalize_nvfp4_scale_(scale: torch.Tensor) -> None:
         scale.copy_(scale.to(torch.float32).abs().to(scale.dtype))
 
 
-def _pad_nvfp4_moe_for_marlin(
-    w13: torch.Tensor,
-    w13_scale: torch.Tensor,
-    w2: torch.Tensor,
-    w2_scale: torch.Tensor,
-    *,
-    is_act_and_mul: bool,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:
-    """Apply rank-local post-load padding required by the Marlin MoE kernel."""
-    num_experts = w13.shape[0]
-    num_shards = 2 if is_act_and_mul else 1
-    intermediate_size = w13.shape[1] // num_shards
-    hidden_size = w13.shape[2] * 2
-    if hidden_size % 128 == 0:
-        tile_size = 64
-    elif hidden_size % 64 == 0:
-        tile_size = 128
-    else:
-        raise ValueError(
-            f"W4A16 Marlin MoE requires hidden_size divisible by 64, got {hidden_size}"
-        )
-    padded_size = (intermediate_size + tile_size - 1) // tile_size * tile_size
-    if padded_size == intermediate_size:
-        return w13, w13_scale, w2, w2_scale, intermediate_size
-
-    def pad_w13(tensor: torch.Tensor) -> torch.Tensor:
-        tensor = tensor.view(
-            num_experts,
-            num_shards,
-            intermediate_size,
-            tensor.shape[-1],
-        )
-        tensor = torch.nn.functional.pad(
-            tensor,
-            (0, 0, 0, padded_size - intermediate_size),
-        )
-        return tensor.reshape(num_experts, num_shards * padded_size, -1)
-
-    w13 = pad_w13(w13)
-    w13_scale = pad_w13(w13_scale)
-    w2 = torch.nn.functional.pad(w2, (0, (padded_size - intermediate_size) // 2))
-    w2_scale = torch.nn.functional.pad(
-        w2_scale,
-        (0, (padded_size - intermediate_size) // 16),
-    )
-    return w13, w13_scale, w2, w2_scale, padded_size
-
-
 def register_nemo_modelopt_nvfp4() -> None:
     """Register NeMo's two ModelOpt NVFP4 configs through vLLM's public API."""
     global _registered
@@ -165,14 +116,7 @@ def register_nemo_modelopt_nvfp4() -> None:
         MarlinNvFp4LinearKernel,
         NvFp4LinearLayerConfig,
     )
-    from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
-        FusedMoEMethodBase,
-    )
-    from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
-        NvFp4MoeBackend,
-        is_global_sf_supported_for_nvfp4_backend,
-        select_nvfp4_moe_backend,
-    )
+    from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend
     from vllm.model_executor.layers.linear import (
         register_weight_loader_v2_supported_method,
     )
@@ -182,8 +126,6 @@ def register_nemo_modelopt_nvfp4() -> None:
         ModelOptNvFp4FusedMoE,
         ModelOptNvFp4LinearMethod,
     )
-    from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Static
-    from vllm.model_executor.utils import replace_parameter
 
     class NemoModelOptNvFp4FusedMoE(ModelOptNvFp4FusedMoE):
         """Native W4A4 MoE plus the vLLM 0.20 input-scale loader fix."""
@@ -273,12 +215,15 @@ def create_weights(self, layer: Any, *args: Any, **kwargs: Any) -> None:
             super().create_weights(layer, *args, **kwargs)
             del layer.input_scale
 
-        # Adapted from vLLM v0.20.0 ModelOptNvFp4LinearMethod
+        # Adapted from vLLM's ModelOptNvFp4LinearMethod
         # .process_weights_after_loading/.apply with input-scale/alpha handling
-        # removed for weight-only W4A16:
-        # https://github.com/vllm-project/vllm/blob/v0.20.0/vllm/model_executor/layers/quantization/modelopt.py#L1169-L1208
-        # Re-sync on vLLM bumps; delete if upstream gains a native W4A16
-        # NVFP4 method.
+        # removed for weight-only W4A16.  vLLM 0.25.1 does now ship a native
+        # ModelOptNvFp4W4A16LinearMethod:
+        # https://github.com/vllm-project/vllm/blob/v0.25.1/vllm/model_executor/layers/quantization/modelopt.py#L1245
+        # This override is kept because 0.25's ModelOptNvFp4Config installs
+        # LinearMethodCls as an *instance* attribute keyed off the quant algo,
+        # which shadows a subclass override (see from_config below).
+        # Re-sync on vLLM bumps.
         def process_weights_after_loading(self, layer: Any) -> None:
             layer.weight_global_scale = Parameter(
                 layer.weight_scale_2.max().to(torch.float32),
@@ -297,30 +242,17 @@ def apply(
             return self.kernel.apply_weights(layer=layer, x=x, bias=bias)
 
     class NemoModelOptW4A16FusedMoE(ModelOptNvFp4FusedMoE):
-        """ModelOpt W4A16 MoE using vLLM's NVFP4 Marlin implementation."""
+        """ModelOpt W4A16 MoE using vLLM's NVFP4 Marlin implementation.
+
+        vLLM 0.25's base __init__ keys weight-only mode off
+        quant_config.quant_method == "W4A16_NVFP4" (activation_key=None), so
+        the 0.20-era duplicated __init__ is gone; the algo passes through
+        from_config unchanged.
+        """
 
         moe_kernel: Any
         moe_quant_config: Any
 
-        def __init__(self, quant_config: object, moe_config: object) -> None:
-            # Duplicates vLLM v0.20.0 ModelOptNvFp4FusedMoE.__init__ except
-            # activation_key=None (weight-only); the base hard-wires
-            # kNvfp4Dynamic and offers no hook:
-            # https://github.com/vllm-project/vllm/blob/v0.20.0/vllm/model_executor/layers/quantization/modelopt.py#L1218-L1234
-            # Intentionally calls FusedMoEMethodBase.__init__ to skip the
-            # parent __init__; do not replace it with super().__init__().
-            # Re-sync on vLLM bumps.
-            FusedMoEMethodBase.__init__(self, moe_config)
-            self.quant_config = quant_config
-            self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
-                config=self.moe,
-                weight_key=kNvfp4Static,
-                activation_key=None,
-            )
-            self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(
-                self.nvfp4_backend
-            )
-
         def create_weights(
             self,
             layer: Any,
@@ -344,35 +276,20 @@ def create_weights(
         def process_weights_after_loading(self, layer: Any) -> None:
             reload_kernel = self.moe_kernel
             reload_quant_config = self.moe_quant_config
-            original_intermediate_size = (
-                layer.moe_config.intermediate_size_per_partition
-            )
             if self.nvfp4_backend == NvFp4MoeBackend.MARLIN:
-                w13, w13_scale, w2, w2_scale, padded_size = _pad_nvfp4_moe_for_marlin(
-                    layer.w13_weight,
-                    layer.w13_weight_scale,
-                    layer.w2_weight,
-                    layer.w2_weight_scale,
-                    is_act_and_mul=self.moe.is_act_and_mul,
-                )
-                replace_parameter(layer, "w13_weight", w13)
-                replace_parameter(layer, "w13_weight_scale", w13_scale)
-                replace_parameter(layer, "w2_weight", w2)
-                replace_parameter(layer, "w2_weight_scale", w2_scale)
+                # vLLM 0.25's prepare_nvfp4_moe_layer_for_marlin pads the
+                # rank-local intermediate tiles itself (and asserts on the
+                # unpadded checkpoint shapes), so no NeMo-side pre-padding.
+                # Only the E4M3 sign-bit canonicalization of the ModelOpt
+                # export remains our concern.
                 _canonicalize_nvfp4_scale_(layer.w13_weight_scale)
                 _canonicalize_nvfp4_scale_(layer.w2_weight_scale)
-                layer.moe_config.intermediate_size_per_partition = padded_size
             # W4A16 checkpoint metadata deliberately omits activation scales so
             # layerwise reload never waits for tensors that do not exist. The
             # native Marlin converter accepts None and removes these attributes.
             layer.w13_input_scale = None
             layer.w2_input_scale = None
-            try:
-                super().process_weights_after_loading(layer)
-            finally:
-                layer.moe_config.intermediate_size_per_partition = (
-                    original_intermediate_size
-                )
+            super().process_weights_after_loading(layer)
             if reload_kernel is not None:
                 self.moe_kernel = reload_kernel
                 self.moe_quant_config = reload_quant_config
@@ -401,7 +318,13 @@ def override_quantization_method(
 
         @classmethod
         def from_config(cls, config: dict[str, Any]) -> Any:
-            return super().from_config(_normalized_w4a16_config(config))
+            instance = super().from_config(_validated_w4a16_config(config))
+            # vLLM 0.25's ModelOptNvFp4Config.__init__ selects LinearMethodCls
+            # from the quant algo as an *instance* attribute, which shadows
+            # the class attribute above; rebind the NeMo method explicitly so
+            # W4A16 linears keep the refit-friendly Marlin implementation.
+            instance.LinearMethodCls = NemoModelOptW4A16LinearMethod
+            return instance
 
     register_quantization_config(NEMO_MODELOPT_W4A4)(NemoModelOptNvFp4Config)
     register_quantization_config(NEMO_MODELOPT_W4A16)(NemoModelOptW4A16Config)

nemo_rl/modelopt/models/generation/vllm_quant_backend.py

@@ -144,10 +144,15 @@ def _batch_fused_modelopt_moe_weights(
 ) -> list[tuple[str, torch.Tensor]]:
     """Map fused ModelOpt payloads to vLLM per-projection checkpoint names.
 
-    Large expert weights and block scales stay batched so vLLM can
+    ``w2`` weights and block scales stay batched so vLLM can
     tensor-parallel-shard the full ``[E, ...]`` tensor at once.  Its scalar
     loader still requires an expert id, so only the tiny per-expert global
     scales are exposed as scalar views.
+
+    Gated ``w13`` payloads are the exception on vLLM >= 0.25: they are emitted
+    as per-expert 2-D shards instead, because ``RoutedExperts.load_weights``'
+    fused-3D branch mis-transposes packed NVFP4. See the comment at the
+    emission site below.
     """
     batched: list[tuple[str, torch.Tensor]] = []
     for name, tensor in weights:
@@ -178,16 +183,24 @@ def _batch_fused_modelopt_moe_weights(
                     f"Expected fused gate/up tensor with an even projection "
                     f"dimension for {name}, got {tuple(tensor.shape)}"
                 )
+            # Emit per-expert 2-D shards rather than batched 3-D tensors:
+            # gated models (e.g. Qwen3-MoE) route batched tensors through
+            # vLLM 0.25's RoutedExperts.load_weights fused branch, whose
+            # orientation heuristic compares the last dim against the
+            # unpacked hidden size and mis-transposes packed NVFP4 weights
+            # (K/2 uint8) and block scales (K/16). Per-expert 2-D loads take
+            # the same weight_loader path as the initial disk load.
             gate, up = tensor.chunk(2, dim=1)
             batched.extend(
                 (
-                    f"{prefix}.experts.0.{projection}.{target_suffix}",
-                    shard,
+                    f"{prefix}.experts.{expert_id}.{projection}.{target_suffix}",
+                    expert_weight,
                 )
                 for projection, shard in (
                     ("gate_proj", gate),
                     ("up_proj", up),
                 )
+                for expert_id, expert_weight in enumerate(shard.unbind(0))
             )
             continue
 
@@ -566,6 +579,36 @@ def new_named_parameters(self, *args, **kwargs):
             for buf in patched_quantizer_buffers:
                 del buf.weight_loader
 
+    @contextmanager
+    def _attach_input_quantizer_amax_loaders(self, model):
+        """Eagerly attach weight_loaders to input_quantizer amax buffers.
+
+        vLLM >= 0.25 loads refit weights through per-module
+        ``load_weights`` (e.g. ``LinearBase.load_weights``), which resolves
+        targets via ``getattr`` and calls ``param.weight_loader(param,
+        loaded_weight, shard_id)`` directly — it never iterates
+        ``model.named_parameters()``, so the lazy attach in
+        ``_patch_named_parameters_to_include_buffers`` no longer fires and
+        quantizer amax buffers arrive without a loader (AttributeError:
+        'Tensor' object has no attribute 'weight_loader').
+        """
+
+        def input_amax_loader(param, loaded_weight, *args, **kwargs):
+            param.copy_(torch.max(param, loaded_weight))
+
+        attached = []
+        for name, buf in model.named_buffers():
+            if "input_quantizer" not in name:
+                continue
+            if not hasattr(buf, "weight_loader"):
+                buf.weight_loader = input_amax_loader
+                attached.append(buf)
+        try:
+            yield
+        finally:
+            for buf in attached:
+                del buf.weight_loader
+
     def _load_weights(self, weights):
         """Load pre-folded weights and input_quantizer amax buffers.
 
@@ -620,6 +663,9 @@ def _load_weights(self, weights):
                 contexts.enter_context(
                     self._patch_named_parameters_to_include_buffers(child)
                 )
+            contexts.enter_context(
+                self._attach_input_quantizer_amax_loaders(self.model_runner.model)
+            )
             return super()._load_weights(weights)
 
     def get_weight_snapshot(self, name: str) -> torch.Tensor: