chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[RL] refactor flash rl weight reload in sglang (#14870)

Browse Source 
Co-authored-by: eternally-z <zzywzj@gmail.com>
This commit is contained in:
Peng Zhang
2025-12-12 12:50:47 +08:00
committed by GitHub
parent 76743a983e
commit 8f5adac8c6
3 changed files with 15 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
python/sglang/srt/model_loader/loader.py
View File

@@ -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())

37
python/sglang/srt/models/qwen2.py
View File

@@ -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

24
python/sglang/srt/models/qwen3.py
View File

@@ -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