[RL] support weight reload for low-bit rollout (#9650)

Co-authored-by: Hecate0821 <hec4te0821@gmail.com>
Co-authored-by: eternally-z <zzywzj@gmail.com>
Co-authored-by: Wilboludriver <wilbolu@outlook.com>
Co-authored-by: Wilbolu <81792854+Wilboludriver@users.noreply.github.com>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>

python/sglang/srt/configs/load_config.py

@@ -23,6 +23,7 @@ class LoadFormat(str, enum.Enum):
     BITSANDBYTES = "bitsandbytes"
     MISTRAL = "mistral"
     LAYERED = "layered"
+    FLASH_RL = "flash_rl"  # For RL training with quantized models
     JAX = "jax"
     REMOTE = "remote"
     REMOTE_INSTANCE = "remote_instance"
@@ -46,6 +47,8 @@ class LoadConfig:
         "dummy" will initialize the weights with random values, which is
             mainly for profiling.
         "bitsandbytes" will load nf4 type weights.
+        "flash_rl" will load weights with support for RL training
+            with quantized models, enabling efficient weight reloading.
     ignore_patterns: The list of patterns to ignore when loading the model.
         Default to "original/**/*" to avoid repeated loading of llama's
         checkpoints.
@@ -78,6 +81,11 @@ class LoadConfig:
     # ModelOpt configuration object
     modelopt_config: Optional[ModelOptConfig] = None
 
+    # QuantizedRL-specific options (for FlashRL-style quantization)
+    rl_quant_profile: Optional[str] = (
+        None  # Path to rollout quantization profile (e.g., /root/profile.7b.pt)
+    )
+
     def __post_init__(self):
         model_loader_extra_config = self.model_loader_extra_config or {}
         if isinstance(model_loader_extra_config, str):

python/sglang/srt/layers/linear.py

@@ -419,6 +419,15 @@ class ColumnParallelLinear(LinearBase):
         else:
             # FIXME: This branch is needed to load deepseek v3 awq.
             # However, we should fix this and avoid the branching here.
+            # After QuantizedRL reload, params might still need tp_rank
+            try:
+                param.load_column_parallel_weight(
+                    loaded_weight,
+                    tp_rank=self.tp_rank,
+                    use_presharded_weights=self.use_presharded_weights,
+                )
+            except TypeError:
+                # Fallback for parameters that don't accept additional args
                 param.load_column_parallel_weight(loaded_weight)
 
     def forward(self, input_):
@@ -1360,6 +1369,15 @@ class RowParallelLinear(LinearBase):
         else:
             # `params` is defined in `vllm/model_executor/parameter.py`,
             # It does not support additional parameters.
+            # However, after QuantizedRL reload, params might still need tp_rank
+            try:
+                param.load_row_parallel_weight(
+                    loaded_weight,
+                    tp_rank=self.tp_rank,
+                    use_presharded_weights=self.use_presharded_weights,
+                )
+            except TypeError:
+                # Fallback for parameters that don't accept additional args
                 param.load_row_parallel_weight(loaded_weight)
 
     def forward(self, input_, skip_all_reduce=False):

python/sglang/srt/model_executor/model_runner.py

@@ -764,6 +764,7 @@ class ModelRunner:
             remote_instance_weight_loader_seed_instance_service_port=self.server_args.remote_instance_weight_loader_seed_instance_service_port,
             remote_instance_weight_loader_send_weights_group_ports=self.server_args.remote_instance_weight_loader_send_weights_group_ports,
             modelopt_config=modelopt_config,
+            rl_quant_profile=self.server_args.rl_quant_profile,
         )
         if self.device == "cpu":
             self.model_config = adjust_config_with_unaligned_cpu_tp(

python/sglang/srt/model_loader/loader.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 import collections
 import dataclasses
 import fnmatch
+import gc
 import glob
 import json
 import logging
@@ -25,6 +26,7 @@ from typing import (
     List,
     Optional,
     Tuple,
+    Union,
     cast,
 )
 
@@ -113,6 +115,8 @@ _is_npu = is_npu()
 # ModelOpt: QUANT_CFG_CHOICES is imported from modelopt_utils.py
 # which contains the complete mapping of quantization config choices
 
+logger = logging.getLogger(__name__)
+
 
 @contextmanager
 def device_loading_context(module: torch.nn.Module, target_device: torch.device):
@@ -693,6 +697,466 @@ class LayeredModelLoader(DefaultModelLoader):
         return model.eval()
 
 
+class QuantizedRLModelLoader(DefaultModelLoader):
+    """
+    Model loader for RL training with FP8 quantization (profile-free, native SGLang).
+
+    Workflow:
+      1. Initial load: Load base model → Record state → Apply FP8 quantization
+      2. Training Actor in full precision
+      3. Reload: Trainer sends full precision weights → Quantize to FP8 → Copy to original memory
+      4. Use torch.as_strided to preserve memory locations across reloads
+
+    Usage:
+      --model-path Qwen/Qwen2.5-7B --quantization fp8 --load-format flash_rl
+    """
+
+    # Parameter attributes to record for weight reloading
+    RECORDED_LOADER_KEYS = [
+        "weight_loader",
+        "load_qkv_weight",
+        "load_column_parallel_weight",
+        "load_row_parallel_weight",
+        "load_merged_column_weight",
+        "output_dim",
+        "input_dim",
+        "_assert_and_load",
+    ]
+
+    # Parameters to skip during FP8 quantization (matches FlashRL's exclude_list)
+    SKIP_QUANTIZATION_PARAMS = [
+        "weight_scale",
+        "input_scale",
+        "output_scale",
+        ".bias",
+        "lm_head.weight",
+        "model.norm.weight",
+        "embed_tokens",  # BF16 params
+        "rotary_emb.inv_freq",
+        "rotary_emb.cos_cached",
+        "rotary_emb.sin_cached",
+        "projector",
+        "input_layernorm.weight",
+        "post_attention_layernorm.weight",  # LayerNorms
+    ]
+
+    # Stacked parameters (Qwen2): shards loaded separately, then combined
+    STACKED_PARAMS_MAPPING = [
+        ("qkv_proj", ["q_proj", "k_proj", "v_proj"]),
+        ("gate_up_proj", ["gate_proj", "up_proj"]),
+    ]
+    _QKV_SHARD_ALIASES = {
+        "q_proj": "q",
+        "k_proj": "k",
+        "v_proj": "v",
+    }
+
+    def __init__(self, load_config: LoadConfig):
+        super().__init__(load_config)
+        logger.info("[QuantizedRL] Profile-free FP8 quantization enabled")
+        self._initial_load_complete = False
+
+    def _prepare_weights(
+        self, model_name_or_path: str, revision: Optional[str], fall_back_to_pt: bool
+    ):
+        """Standard weight preparation using base model path."""
+        logger.info(f"[QuantizedRL] Loading from base model: {model_name_or_path}")
+        temp_config = LoadConfig(load_format=LoadFormat.AUTO)
+        temp_loader = DefaultModelLoader(temp_config)
+        return temp_loader._prepare_weights(
+            model_name_or_path, revision, fall_back_to_pt
+        )
+
+    @staticmethod
+    def _bind_method_to_cls(func, obj):
+        """Bind function to object instance (for weight_loader methods)."""
+        import types
+
+        if hasattr(func, "__self__") or not callable(func):
+            return func
+        return types.MethodType(func, obj)
+
+    def load_weights_and_postprocess(self, model, weights, target_device):
+        """
+        Initial load: Load BF16 → Record state → Apply FP8 quantization.
+        Called ONCE during model initialization.
+        """
+        logger.info("[QuantizedRL] Initial load with FP8 quantization")
+
+        model.load_weights(weights)
+        original_weights = dict(model.named_parameters())
+
+        # Record pre-quantization state (shape/stride) for torch.as_strided reset
+
+        model.original_weights_rebuild_keys = {}
+        for name, p in original_weights.items():
+            model.original_weights_rebuild_keys[name] = {
+                "shape": p.shape,
+                "stride": p.stride(),
+                "dtype": p.dtype,
+                "nbytes": p.untyped_storage().nbytes(),
+            }
+
+        # Record parameter attributes (weight_loader, etc.) before quantization
+        recorded_loader = {
+            k: dict() for k in QuantizedRLModelLoader.RECORDED_LOADER_KEYS
+        }
+        for name, p in original_weights.items():
+            for key in QuantizedRLModelLoader.RECORDED_LOADER_KEYS:
+                if hasattr(p, key):
+                    attr = getattr(p, key)
+                    if not callable(attr):
+                        recorded_loader[key][name] = attr
+                    elif hasattr(attr, "__self__") and p is attr.__self__:
+                        recorded_loader[key][name] = attr.__func__  # Store unbound
+                    else:
+                        recorded_loader[key][name] = attr
+        model.recorded_loader = recorded_loader
+
+        # Apply FP8 quantization (creates new Parameters, loses attributes)
+        for _, module in model.named_modules():
+            quant_method = getattr(module, "quant_method", None)
+            if quant_method is not None:
+                with device_loading_context(module, target_device):
+                    quant_method.process_weights_after_loading(module)
+
+        model.flash_rl_initial_load_complete = True
+        self._initial_load_complete = True
+        logger.info("[QuantizedRL] Initial load complete")
+
+    @staticmethod
+    def is_reload_scenario(model):
+        """Check if model is ready for reloading (initial load completed)."""
+        return (
+            hasattr(model, "original_weights_rebuild_keys")
+            and hasattr(model, "recorded_loader")
+            and getattr(model, "flash_rl_initial_load_complete", False)
+        )
+
+    @staticmethod
+    def _is_stacked_param(name):
+        """Check if parameter is stacked (qkv_proj, gate_up_proj)."""
+        for stacked_name, _ in QuantizedRLModelLoader.STACKED_PARAMS_MAPPING:
+            if stacked_name in name:
+                return True
+        return False
+
+    @staticmethod
+    def _resolve_stacked_info(name: str) -> Tuple[str, Optional[str], Optional[Any]]:
+        for target, shard_names in QuantizedRLModelLoader.STACKED_PARAMS_MAPPING:
+            for idx, shard in enumerate(shard_names):
+                if shard in name:
+                    shard_id = (
+                        QuantizedRLModelLoader._QKV_SHARD_ALIASES.get(shard, shard)
+                        if target == "qkv_proj"
+                        else idx
+                    )
+                    return name.replace(shard, target), target, shard_id
+        return name, None, None
+
+    @staticmethod
+    def _store_quantized_scale(
+        scale_store: Dict[str, Union[torch.Tensor, Dict[Any, torch.Tensor]]],
+        name: str,
+        scale: torch.Tensor,
+    ) -> None:
+        param_name, stacked_key, shard_id = (
+            QuantizedRLModelLoader._resolve_stacked_info(name)
+        )
+        if stacked_key is None:
+            scale_store[param_name] = scale
+        else:
+            shard_dict = scale_store.setdefault(param_name, {})
+            assert isinstance(shard_dict, dict)
+            shard_dict[shard_id] = scale
+
+    @staticmethod
+    def _apply_scale_update(
+        all_params: Dict[str, torch.nn.Parameter],
+        param_name: str,
+        scale_info: Union[torch.Tensor, Dict[Any, torch.Tensor], None],
+    ) -> None:
+        if scale_info is None:
+            return
+        # Get tp rank and size
+        tp_rank = get_tensor_model_parallel_rank()
+        tp_size = get_tensor_model_parallel_world_size()
+
+        def _get_tp_sharded_scale(full_scale_tensor):
+            """Get tp sharded scale from full scale tensor"""
+            if tp_size == 1:
+                return full_scale_tensor
+
+            full_dim = full_scale_tensor.shape[0]
+            shard_dim = full_dim // tp_size
+            start_idx = tp_rank * shard_dim
+            end_idx = start_idx + shard_dim
+            return full_scale_tensor[start_idx:end_idx]
+
+        if param_name.endswith(".weight"):
+            scale_param_name = f"{param_name[:-7]}.weight_scale"
+        else:
+            scale_param_name = f"{param_name}.weight_scale"
+
+        scale_param = all_params.get(scale_param_name)
+        if scale_param is None:
+            logger.warning(
+                "[QuantizedRL] Scale parameter not found: %s", scale_param_name
+            )
+            return
+        if isinstance(scale_info, torch.Tensor):
+            new_scale = scale_info.t().contiguous()
+            if scale_param.data.shape == new_scale.shape:
+                scale_param.data.copy_(new_scale)
+            else:
+                logger.warning(
+                    "[QuantizedRL] Scale shape mismatch for %s: expected %s, got %s",
+                    scale_param_name,
+                    scale_param.data.shape,
+                    new_scale.shape,
+                )
+        else:
+            stacked_key = next(
+                (
+                    target
+                    for target, _ in QuantizedRLModelLoader.STACKED_PARAMS_MAPPING
+                    if target in param_name
+                ),
+                None,
+            )
+            shard_names = next(
+                (
+                    names
+                    for target, names in QuantizedRLModelLoader.STACKED_PARAMS_MAPPING
+                    if target == stacked_key
+                ),
+                [],
+            )
+            rows_per_shard = scale_param.data.shape[-1] // max(len(shard_names), 1)
+            if rows_per_shard * len(shard_names) != scale_param.data.shape[-1]:
+                logger.warning(
+                    f"Scale param shape {scale_param.data.shape[-1]} not divisible by {len(shard_names)}"
+                )
+            offset = 0
+            for idx, shard in enumerate(shard_names):
+                shard_id = (
+                    QuantizedRLModelLoader._QKV_SHARD_ALIASES.get(shard, shard)
+                    if stacked_key == "qkv_proj"
+                    else idx
+                )
+                shard_scale = scale_info.get(shard_id)
+                shard_scale = _get_tp_sharded_scale(shard_scale)
+                if shard_scale is None:
+                    offset += rows_per_shard
+                    continue
+                shard_rows = shard_scale.shape[0]
+                start = offset
+                end = start + shard_rows
+                scale_param.data[..., start:end] = shard_scale.t().contiguous()
+                offset = end
+
+    @staticmethod
+    def rebinding_and_load_weights(model, first_time_load_weights, weights):
+        """
+        Reload: VERL sends BF16 → Quantize to FP8 → Copy to original memory.
+
+        Flow: Reset params → Restore attributes → Quantize in iterator → Load → Copy back
+        """
+        logger.info("[QuantizedRL] Reload: Updating weights with FP8 quantization")
+
+        weights_list = list(weights)
+        updated_param_names, is_last_update = (
+            QuantizedRLModelLoader._get_updated_params(weights_list, model)
+        )
+
+        # Save current FP8 parameter data pointers
+        existing_params = dict(model.named_parameters())
+        current_param_data = {}
+        for name in updated_param_names:
+            if name in existing_params:
+                current_param_data[name] = existing_params[name].data
+
+        # Reset to pre-quantization shape using torch.as_strided
+        # Keeps same storage, just changes view - critical for memory preservation
+        for name, rebuild_info in model.original_weights_rebuild_keys.items():
+            if name in updated_param_names and name in existing_params:
+                existing_params[name].data = torch.as_strided(
+                    # Note: avoid clone here
+                    existing_params[name].data.clone(),
+                    rebuild_info["shape"],
+                    rebuild_info["stride"],
+                )
+
+        # Restore weight loader attributes (only if missing)
+        for k, loader_dict in model.recorded_loader.items():
+            for param_name, loader in loader_dict.items():
+                if param_name in updated_param_names and param_name in existing_params:
+                    param = existing_params[param_name]
+                    if not hasattr(param, k):
+                        if callable(loader):
+                            if hasattr(loader, "__self__"):
+                                setattr(param, k, loader)
+                            else:
+                                setattr(
+                                    param,
+                                    k,
+                                    QuantizedRLModelLoader._bind_method_to_cls(
+                                        loader, param
+                                    ),
+                                )
+                        else:
+                            setattr(param, k, loader)
+
+        del existing_params
+
+        # Quantize BF16 weights to FP8 in iterator (before weight_loader)
+        # Store scales for later update
+        quantized_scales: Dict[str, Union[torch.Tensor, Dict[Any, torch.Tensor]]] = {}
+
+        def quantize_weights_iterator(weights_iter):
+            """Quantize individual shards before weight_loader stacks them."""
+            from sglang.srt.layers.quantization.fp8_kernel import (
+                per_token_group_quant_fp8,
+            )
+
+            for name, weight in weights_iter:
+                if any(
+                    skip in name
+                    for skip in QuantizedRLModelLoader.SKIP_QUANTIZATION_PARAMS
+                ):
+                    logger.info(f"[QuantizedRL] Skip: {name} ({weight.dtype})")
+                    yield (name, weight)
+                elif weight.dtype in [torch.bfloat16, torch.float32, torch.float16]:
+                    qweight, scale = per_token_group_quant_fp8(weight, weight.shape[-1])
+                    logger.info(f"[QuantizedRL] Quantize: {name} {weight.dtype}→FP8")
+                    QuantizedRLModelLoader._store_quantized_scale(
+                        quantized_scales, name, scale
+                    )
+                    yield (name, qweight)
+                else:
+                    logger.info(f"[QuantizedRL] Keep: {name} ({weight.dtype})")
+                    yield (name, weight)
+
+        # Load quantized weights (weight_loader stacks FP8 shards)
+        first_time_load_weights(quantize_weights_iterator(iter(weights_list)))
+
+        # Copy back to original FP8 memory locations and update scales
+        all_params = dict(model.named_parameters())
+
+        for name in updated_param_names:
+            if name not in all_params or name not in current_param_data:
+                continue
+            if any(
+                skip in name for skip in QuantizedRLModelLoader.SKIP_QUANTIZATION_PARAMS
+            ):
+                continue
+
+            new_param = all_params[name]
+            old_fp8_data = current_param_data[name]
+
+            # Handle embeddings/lm_head (BF16) and quantized weights (FP8)
+            if "embed_tokens" in name or "lm_head" in name:
+                old_fp8_data.copy_(new_param.data)
+                new_param.data = old_fp8_data
+            elif (
+                new_param.dtype == torch.float8_e4m3fn
+                and old_fp8_data.dtype == torch.float8_e4m3fn
+            ):
+                # FP8: Use strided view for transposed storage
+                strided_data = torch.as_strided(
+                    new_param.data, old_fp8_data.shape, old_fp8_data.stride()
+                )
+                old_fp8_data.copy_(strided_data)
+                new_param.data = old_fp8_data
+                QuantizedRLModelLoader._apply_scale_update(
+                    all_params,
+                    name,
+                    quantized_scales.get(name),
+                )
+            elif new_param.dtype == old_fp8_data.dtype:
+                # Same dtype (LayerNorm, etc.): Direct copy
+                old_fp8_data.copy_(new_param.data)
+                new_param.data = old_fp8_data
+            else:
+                raise RuntimeError(
+                    f"Unexpected dtype mismatch for {name}: "
+                    f"new={new_param.dtype}, old={old_fp8_data.dtype}"
+                )
+
+        # Cleanup
+        del current_param_data
+        if is_last_update:
+            gc.collect()
+            torch.cuda.empty_cache()
+
+        logger.info("[QuantizedRL] Reload complete")
+        return updated_param_names, is_last_update
+
+    @staticmethod
+    def _get_updated_params(weights_list, model):
+        """Identify which parameters need updating from incoming weights."""
+        stacked_params_mapping = [
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+
+        params_dict = dict(model.named_parameters())
+        updated_params = set()
+        is_last_update = False
+
+        for name, _ in weights_list:
+            if name == "lm_head.weight":
+                is_last_update = True
+
+            if any(
+                skip in name for skip in QuantizedRLModelLoader.SKIP_QUANTIZATION_PARAMS
+            ):
+                continue
+
+            from sglang.srt.layers.utils import get_layer_id
+
+            # Skip params outside layer range (for pipeline parallelism)
+            layer_id = get_layer_id(name)
+            if (
+                layer_id is not None
+                and hasattr(model, "start_layer")
+                and (layer_id < model.start_layer or layer_id >= model.end_layer)
+            ):
+                continue
+
+            # Skip tied embeddings and vision tower params
+            if (
+                hasattr(model, "config")
+                and model.config.tie_word_embeddings
+                and "lm_head.weight" in name
+            ):
+                continue
+            if name.startswith("model.vision_tower") and name not in params_dict:
+                continue
+
+            # Map stacked param shards (q/k/v_proj → qkv_proj)
+            mapped = False
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name in name:
+                    name = name.replace(weight_name, param_name)
+                    if name.endswith(".bias") and name not in params_dict:
+                        continue
+                    updated_params.add(name)
+                    mapped = True
+                    break
+
+            if not mapped:
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                if name in params_dict:
+                    updated_params.add(name)
+
+        return list(updated_params), is_last_update
+
+
 class DummyModelLoader(BaseModelLoader):
     """Model loader that will set model weights to random values."""
 
@@ -2094,6 +2558,27 @@ def get_model_loader(
     if load_config.load_format == LoadFormat.LAYERED:
         return LayeredModelLoader(load_config)
 
+    # Check for FLASH_RL format early
+    # FP8 approach: BF16/FP16 model with native FP8 quantization
+    if load_config.load_format == LoadFormat.FLASH_RL:
+        logger.info(
+            "Using QuantizedRLModelLoader for RL training with native FP8 quantization."
+        )
+        logger.info(
+            "FP8 approach: Model loads with native SGLang FP8 quantization. "
+            "Same model path for both training and inference."
+        )
+
+        # Set quantization to FP8 for native SGLang support
+        if model_config and not model_config.quantization:
+            logger.info(
+                "QuantizedRL: Setting quantization to fp8 (native SGLang support). "
+                "Model will be loaded with FP8 infrastructure"
+            )
+            model_config.quantization = "fp8"
+
+        return QuantizedRLModelLoader(load_config)
+
     if load_config.load_format == LoadFormat.REMOTE:
         return RemoteModelLoader(load_config)
 

python/sglang/srt/models/qwen2.py

@@ -566,7 +566,8 @@ class Qwen2ForCausalLM(nn.Module):
     def end_layer(self):
         return self.model.end_layer
 
-    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+    def _load_weights_impl(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+        """Internal implementation of weight loading without reload scenario handling."""
         stacked_params_mapping = [
             # (param_name, shard_name, shard_id)
             ("qkv_proj", "q_proj", "q"),
@@ -577,6 +578,7 @@ class Qwen2ForCausalLM(nn.Module):
         ]
 
         params_dict = dict(self.named_parameters())
+        updated_params = set()
         for name, loaded_weight in weights:
             layer_id = get_layer_id(name)
             if (
@@ -620,6 +622,7 @@ class Qwen2ForCausalLM(nn.Module):
                 param = params_dict[name]
                 weight_loader = param.weight_loader
                 weight_loader(param, loaded_weight, shard_id)
+                updated_params.add(name)
                 break
             else:
                 # Skip loading extra bias for GPTQ models.
@@ -632,9 +635,41 @@ class Qwen2ForCausalLM(nn.Module):
                         param, "weight_loader", default_weight_loader
                     )
                     weight_loader(param, loaded_weight)
+                    updated_params.add(name)
                 else:
                     logger.warning(f"Parameter {name} not found in params_dict")
 
+        return updated_params
+
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+        """
+        Load weights into the model, with support for RL training reload scenarios.
+
+        Args:
+            weights: Iterator of (name, tensor) tuples with weights to load
+
+        Note: quantize_fn and quant_profile are stored on the model during initialization,
+              so they don't need to be passed as arguments.
+        """
+        import logging
+
+        from sglang.srt.model_loader.loader import QuantizedRLModelLoader
+
+        logger = logging.getLogger(__name__)
+
+        # Check if this is a reload scenario for RL training with quantized models
+        is_reload = QuantizedRLModelLoader.is_reload_scenario(self)
+        if is_reload:
+            logger.info("RELOAD SCENARIO - Using rebinding_and_load_weights")
+            # Use the fast path for RL training reloads
+            # quantize_fn and quant_profile are retrieved from model inside rebinding_and_load_weights
+            QuantizedRLModelLoader.rebinding_and_load_weights(
+                self, self._load_weights_impl, weights
+            )
+        else:
+            # Standard weight loading path
+            self._load_weights_impl(weights)
+
     def get_embed_and_head(self):
         return self.model.embed_tokens.weight, self.lm_head.weight
 

python/sglang/srt/models/qwen3.py

@@ -495,7 +495,8 @@ class Qwen3ForCausalLM(nn.Module):
     def end_layer(self):
         return self.model.end_layer
 
-    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+    def _load_weights_impl(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+        """Internal implementation of weight loading without reload scenario handling."""
         stacked_params_mapping = [
             # (param_name, shard_name, shard_id)
             ("qkv_proj", "q_proj", "q"),
@@ -506,6 +507,7 @@ class Qwen3ForCausalLM(nn.Module):
         ]
 
         params_dict = dict(self.named_parameters())
+        updated_params = set()
         for name, loaded_weight in weights:
             if "Embedding" in self.config.name_or_path:
                 name = add_prefix(name, "model")
@@ -552,6 +554,7 @@ class Qwen3ForCausalLM(nn.Module):
                 param = params_dict[name]
                 weight_loader = param.weight_loader
                 weight_loader(param, loaded_weight, shard_id)
+                updated_params.add(name)
                 break
             else:
                 # Skip loading extra bias for GPTQ models.
@@ -564,9 +567,28 @@ class Qwen3ForCausalLM(nn.Module):
                         param, "weight_loader", default_weight_loader
                     )
                     weight_loader(param, loaded_weight)
+                    updated_params.add(name)
                 else:
                     logger.warning(f"Parameter {name} not found in params_dict")
 
+        return updated_params
+
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+        """Load weights into the model, with support for RL training reload scenarios."""
+        from sglang.srt.model_loader.loader import QuantizedRLModelLoader
+
+        # Check if this is a reload scenario for RL training with quantized models
+        is_reload = QuantizedRLModelLoader.is_reload_scenario(self)
+        if is_reload:
+            # Use the fast path for RL training reloads
+            logger.info("[QuantizedRL] Using fast path reload in load_weights")
+            QuantizedRLModelLoader.rebinding_and_load_weights(
+                self, self._load_weights_impl, weights
+            )
+        else:
+            # Standard weight loading path
+            self._load_weights_impl(weights)
+
     def get_embed_and_head(self):
         return self.model.embed_tokens.weight, self.lm_head.weight
 

python/sglang/srt/server_args.py

@@ -79,6 +79,7 @@ LOAD_FORMAT_CHOICES = [
     "gguf",
     "bitsandbytes",
     "layered",
+    "flash_rl",
     "remote",
     "remote_instance",
 ]
@@ -250,6 +251,7 @@ class ServerArgs:
     skip_tokenizer_init: bool = False
     load_format: str = "auto"
     model_loader_extra_config: str = "{}"
+    rl_quant_profile: Optional[str] = None  # For flash_rl load format
     trust_remote_code: bool = False
     context_length: Optional[int] = None
     is_embedding: bool = False
@@ -2169,6 +2171,12 @@ class ServerArgs:
             "This will be passed to the model loader corresponding to the chosen load_format.",
             default=ServerArgs.model_loader_extra_config,
         )
+        parser.add_argument(
+            "--rl-quant-profile",
+            type=str,
+            default=ServerArgs.rl_quant_profile,
+            help="Path to the FlashRL quantization profile. Required when using --load-format flash_rl.",
+        )
         parser.add_argument(
             "--trust-remote-code",
             action="store_true",