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

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

Browse Source 
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>
This commit is contained in:
Peng Zhang
2025-12-10 15:44:01 +08:00
committed by GitHub
parent b0a25d0913
commit 21028b5507
7 changed files with 581 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/sglang/srt/configs/load_config.py
View File

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

22
python/sglang/srt/layers/linear.py
View File

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

1
python/sglang/srt/model_executor/model_runner.py
View File

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

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

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

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

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

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

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

8
python/sglang/srt/server_args.py
View File

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