feat(sglang): NVFP4 quantized rollout and refit

examples/configs/grpo_math_1B_sglang.yaml

@@ -31,10 +31,14 @@ policy:
       rollout_health_check_timeout: 60
       rollout_health_check_first_wait: 60
       # Weight precision for rollout/refit. scheme=bf16 (default) sends BF16
-      # HF tensors; scheme=mxfp8 boots SGLang from an MXFP8 HF checkpoint and
-      # quantizes refit tensors online (see SglangQuantizationConfig).
+      # HF tensors; scheme=mxfp8/nvfp4 boots SGLang from a matching quantized
+      # HF checkpoint and quantizes refit tensors online.
       quantization:
         scheme: bf16
+        extra_high_precision_layers_hf: []
+        num_layers_at_start_in_bf16: 0
+        num_layers_at_end_in_bf16: 0
+        modules_to_not_convert: []
       sglang_server_config:
         needs_offload: true
         cpu_weight_backup: true

nemo_rl/algorithms/grpo.py

@@ -1266,22 +1266,33 @@ def init_vllm_then_policy():
         if "model_path" not in generation_config["sglang_cfg"]:
             generation_config["sglang_cfg"]["model_path"] = policy_config["model_name"]
 
-        # If MXFP8 is requested, ensure SGLang boots from an MXFP8 HF
-        # checkpoint. This must happen before ``init_sglang`` so the engine
-        # loads quantized weights.
+        # Quantized online refit requires SGLang to boot from a checkpoint
+        # with the exact same tensor layout and high-precision exclusions.
+        # Resolve it before ``init_sglang`` starts the engines.
         sglang_quantization_cfg = (
             generation_config["sglang_cfg"].get("quantization") or {}
         )
-        if sglang_quantization_cfg.get("scheme", "bf16") == "mxfp8":
-            from nemo_rl.models.generation.sglang.mxfp8_setup import (
-                ensure_mxfp8_checkpoint,
-            )
+        from nemo_rl.models.generation.sglang.quantization_utils import (
+            ensure_sglang_quantized_checkpoint,
+            get_sglang_quantization_scheme,
+            validate_sglang_quantized_refit_backend,
+        )
 
-            mxfp8_path = ensure_mxfp8_checkpoint(
+        sglang_quantization_scheme = get_sglang_quantization_scheme(
+            sglang_quantization_cfg
+        )
+        validate_sglang_quantized_refit_backend(
+            scheme=sglang_quantization_scheme,
+            use_megatron=bool(
+                policy_config.get("megatron_cfg", {}).get("enabled", False)
+            ),
+        )
+        generation_config["sglang_cfg"]["model_path"] = (
+            ensure_sglang_quantized_checkpoint(
                 model_path=generation_config["sglang_cfg"]["model_path"],
-                quantization_cfg=sglang_quantization_cfg,
+                quantization_config=sglang_quantization_cfg,
             )
-            generation_config["sglang_cfg"]["model_path"] = mxfp8_path
+        )
 
         policy_generation, policy = initialize_generation_with_policy(
             init_generation_fn=init_sglang,

nemo_rl/algorithms/ppo.py

@@ -634,6 +634,31 @@ def initialize_generation_with_policy(
         if "model_path" not in generation_config["sglang_cfg"]:
             generation_config["sglang_cfg"]["model_path"] = policy_config["model_name"]
 
+        sglang_quantization_cfg = (
+            generation_config["sglang_cfg"].get("quantization") or {}
+        )
+        from nemo_rl.models.generation.sglang.quantization_utils import (
+            ensure_sglang_quantized_checkpoint,
+            get_sglang_quantization_scheme,
+            validate_sglang_quantized_refit_backend,
+        )
+
+        sglang_quantization_scheme = get_sglang_quantization_scheme(
+            sglang_quantization_cfg
+        )
+        validate_sglang_quantized_refit_backend(
+            scheme=sglang_quantization_scheme,
+            use_megatron=bool(
+                policy_config.get("megatron_cfg", {}).get("enabled", False)
+            ),
+        )
+        generation_config["sglang_cfg"]["model_path"] = (
+            ensure_sglang_quantized_checkpoint(
+                model_path=generation_config["sglang_cfg"]["model_path"],
+                quantization_config=sglang_quantization_cfg,
+            )
+        )
+
         policy_generation, policy, value_model = initialize_generation_with_policy(
             init_generation_fn=init_sglang,
             generation_name="SGLang",

nemo_rl/models/generation/sglang/config.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, NotRequired, TypedDict
+from typing import Any, Literal, NotRequired, TypedDict
 
 from nemo_rl.models.generation.interfaces import GenerationConfig
 
@@ -21,15 +21,18 @@ class SglangQuantizationConfig(TypedDict, total=False):
     """SGLang weight-precision config.
 
     ``scheme="bf16"`` (or omitting the block) means BF16 rollout/refit. Set
-    ``scheme="mxfp8"`` to boot SGLang from an MXFP8 HF checkpoint and to send
-    MXFP8 HF tensors during online refit.
+    ``scheme="mxfp8"`` or ``scheme="nvfp4"`` to boot SGLang from the
+    corresponding quantized HF checkpoint and quantize HF tensors during
+    online refit. High-precision exclusions are shared by offline conversion
+    and online refit.
     """
 
-    scheme: str  # "bf16" | "mxfp8"
-    weight_block_size: list[int]
-    scale_fmt: str
+    scheme: Literal["bf16", "mxfp8", "nvfp4"]
+    # HF module-name substrings that the checkpoint loader and refit both skip.
     modules_to_not_convert: list[str]
+    # Additional HF weight-name substrings to keep in high precision.
     extra_high_precision_layers_hf: list[str]
+    # Number of decoder layers at each edge to keep in high precision.
     num_layers_at_start_in_bf16: int
     num_layers_at_end_in_bf16: int
     converted_model_path: str
@@ -97,7 +100,7 @@ class SglangSpecificArgs(TypedDict):
     rollout_health_check_interval: NotRequired[int]
     rollout_health_check_timeout: NotRequired[int]
     rollout_health_check_first_wait: NotRequired[int]
-    # Weight precision and (when scheme=mxfp8) offline-conversion knobs.
+    # Weight precision and quantized-checkpoint conversion/refit knobs.
     quantization: NotRequired[SglangQuantizationConfig]
     sglang_router_config: SGLangRouterConfig
 

nemo_rl/models/generation/sglang/mxfp8_quantization_core.py

@@ -113,21 +113,16 @@ def should_quantize(
 def quantize_mxfp8(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
     """Return ``(qweight, scale)`` in the SGLang MXFP8 layout.
 
-    Uses flashinfer's swizzle-free MXFP8 kernel (``flashinfer.mxfp8_quantize``
-    with ``is_sf_swizzled_layout=False``). flashinfer is a hard requirement
-    here — both the SGLang and Megatron actor environments pin it via
-    ``pyproject.toml``'s global ``flashinfer-python==0.6.4`` constraint, so a
-    missing import means the env was built incorrectly.
+    Uses flashinfer's swizzle-free MXFP8 kernel
+    (``is_sf_swizzled_layout=False``).
     """
     try:
         from flashinfer import mxfp8_quantize as flashinfer_mxfp8_quantize
     except ImportError as e:
         raise ImportError(
             "flashinfer is required for MXFP8 weight quantization but is not "
-            "installed in the current actor environment. Install "
-            "`flashinfer-python==0.6.4` (and `flashinfer-cubin==0.6.4`); "
-            "in NeMo-RL this is normally provided by the `mcore` or `sglang` "
-            "extras (see pyproject.toml constraint-dependencies)."
+            "installed in the current actor environment. In NeMo-RL this is "
+            "normally provided by the pinned `mcore` or `sglang` extras."
         ) from e
 
     weight = weight.contiguous()
@@ -150,45 +145,3 @@ def source_fp8_to_mxfp8_scale_u8(
         SOURCE_FP8_BLOCK_SIZE[0], dim=-2
     ).repeat_interleave(SOURCE_FP8_BLOCK_SIZE[1] // TARGET_MXFP8_BLOCK_SIZE[1], dim=-1)
     return mxfp8_scale_u8[..., :n, : (k // TARGET_MXFP8_BLOCK_SIZE[1])].contiguous()
-
-
-def build_dynamic_skip_substrings(
-    *,
-    quantization_config: dict[str, Any],
-    num_hidden_layers: int,
-) -> tuple[str, ...]:
-    """Compute the dynamic skip substrings for one HF model.
-
-    Combines the static ``SKIP_WEIGHT_SUBSTRINGS`` list with the user-provided
-    ``extra_high_precision_layers_hf`` / ``modules_to_not_convert`` lists from
-    the quantization config, plus per-layer prefixes for the ``head`` / ``tail``
-    BF16-band layers.
-    """
-    extra_high_precision_layers_hf = tuple(
-        quantization_config.get("extra_high_precision_layers_hf", ()) or ()
-    )
-    modules_to_not_convert = tuple(
-        quantization_config.get("modules_to_not_convert", ()) or ()
-    )
-    num_layers_at_start_in_bf16 = int(
-        quantization_config.get("num_layers_at_start_in_bf16", 0) or 0
-    )
-    num_layers_at_end_in_bf16 = int(
-        quantization_config.get("num_layers_at_end_in_bf16", 0) or 0
-    )
-
-    head_end_idx = num_layers_at_start_in_bf16
-    tail_start_idx = num_hidden_layers - num_layers_at_end_in_bf16
-    dynamic_skip_layer_prefixes: set[str] = set()
-    dynamic_skip_layer_prefixes.update(
-        f"model.layers.{i}." for i in range(0, head_end_idx)
-    )
-    dynamic_skip_layer_prefixes.update(
-        f"model.layers.{i}." for i in range(tail_start_idx, num_hidden_layers)
-    )
-    return (
-        *SKIP_WEIGHT_SUBSTRINGS,
-        *extra_high_precision_layers_hf,
-        *modules_to_not_convert,
-        *sorted(dynamic_skip_layer_prefixes),
-    )