From 21028b55072a161c8669a7d565677e2dd1c8f6ef Mon Sep 17 00:00:00 2001 From: Peng Zhang Date: Wed, 10 Dec 2025 15:44:01 +0800 Subject: [PATCH] [RL] support weight reload for low-bit rollout (#9650) Co-authored-by: Hecate0821 Co-authored-by: eternally-z Co-authored-by: Wilboludriver Co-authored-by: Wilbolu <81792854+Wilboludriver@users.noreply.github.com> Co-authored-by: Ke Bao --- python/sglang/srt/configs/load_config.py | 8 + python/sglang/srt/layers/linear.py | 22 +- .../sglang/srt/model_executor/model_runner.py | 1 + python/sglang/srt/model_loader/loader.py | 485 ++++++++++++++++++ python/sglang/srt/models/qwen2.py | 37 +- python/sglang/srt/models/qwen3.py | 24 +- python/sglang/srt/server_args.py | 8 + 7 files changed, 581 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/configs/load_config.py b/python/sglang/srt/configs/load_config.py index 042eb322a..e87bb21e4 100644 --- a/python/sglang/srt/configs/load_config.py +++ b/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): diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index f3500540d..b63680d02 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -419,7 +419,16 @@ class ColumnParallelLinear(LinearBase): else: # FIXME: This branch is needed to load deepseek v3 awq. # However, we should fix this and avoid the branching here. - param.load_column_parallel_weight(loaded_weight) + # 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_): bias = self.bias if not self.skip_bias_add else None @@ -1360,7 +1369,16 @@ class RowParallelLinear(LinearBase): else: # `params` is defined in `vllm/model_executor/parameter.py`, # It does not support additional parameters. - param.load_row_parallel_weight(loaded_weight) + # 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): if self.input_is_parallel: diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index ac5e363ef..b8b1b9627 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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( diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 19d0fec10..119967978 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/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) diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py index a7dbadec6..c09daebf7 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/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 diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py index 30b92acbd..7fcfa3fe4 100644 --- a/python/sglang/srt/models/qwen3.py +++ b/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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index ddd0d0ca4..8fc9062ca 100644 --- a/python/sglang/srt/server_args.py +++ b/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",