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,15 @@ policy:
       precision: fp8
       max_model_len: 4096
       use_deep_gemm: true
+    vllm_kwargs:
+      # At the same gpu_memory_utilization, vLLM 0.25's profiling sizes the
+      # KV cache several GiB larger than 0.20 did, 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"; lowering gpu_memory_utilization to 0.5
+      # still OOM'ed). Pin the KV cache to a 0.20-equivalent size instead
+      # (takes precedence over gpu_memory_utilization).
+      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

@@ -23,7 +23,6 @@
 method-owned MoE kernel references across layerwise reload.
 """
 
-import copy
 from types import MethodType
 from typing import Any
 
@@ -88,17 +87,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 +103,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 +113,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 +123,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."""
@@ -297,30 +236,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 +270,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 +312,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

@@ -178,16 +178,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 +574,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 +658,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: