From 8f5adac8c64d16aa3a68795ca6f27e9e7bf94f19 Mon Sep 17 00:00:00 2001 From: Peng Zhang Date: Fri, 12 Dec 2025 12:50:47 +0800 Subject: [PATCH] [RL] refactor flash rl weight reload in sglang (#14870) Co-authored-by: eternally-z --- python/sglang/srt/model_loader/loader.py | 13 +++++++++ python/sglang/srt/models/qwen2.py | 37 +----------------------- python/sglang/srt/models/qwen3.py | 24 +-------------- 3 files changed, 15 insertions(+), 59 deletions(-) diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 5e41aea5f..a981bcf1e 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -788,6 +788,19 @@ class QuantizedRLModelLoader(DefaultModelLoader): """ logger.info("[QuantizedRL] Initial load with FP8 quantization") + original_load_weights = model.load_weights + + def load_weights_proxy(weights): + if QuantizedRLModelLoader.is_reload_scenario(model): + logger.info("[QuantizedRL] Using fast path reload in load_weights") + QuantizedRLModelLoader.rebinding_and_load_weights( + model, original_load_weights, weights + ) + else: + original_load_weights(weights) + + model.load_weights = load_weights_proxy + model.load_weights(weights) original_weights = dict(model.named_parameters()) diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py index c09daebf7..a7dbadec6 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/python/sglang/srt/models/qwen2.py @@ -566,8 +566,7 @@ class Qwen2ForCausalLM(nn.Module): def end_layer(self): return self.model.end_layer - def _load_weights_impl(self, weights: Iterable[Tuple[str, torch.Tensor]]): - """Internal implementation of weight loading without reload scenario handling.""" + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), @@ -578,7 +577,6 @@ 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 ( @@ -622,7 +620,6 @@ 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. @@ -635,41 +632,9 @@ 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 c527201c0..6ab04da65 100644 --- a/python/sglang/srt/models/qwen3.py +++ b/python/sglang/srt/models/qwen3.py @@ -496,8 +496,7 @@ class Qwen3ForCausalLM(nn.Module): def end_layer(self): return self.model.end_layer - def _load_weights_impl(self, weights: Iterable[Tuple[str, torch.Tensor]]): - """Internal implementation of weight loading without reload scenario handling.""" + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), @@ -508,7 +507,6 @@ 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") @@ -555,7 +553,6 @@ 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. @@ -568,28 +565,9 @@ 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