Skip to content

[Quantization] Fix per-tensor FP8 MoE parameter identity for EPLB hot-reload #19499

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Socratesa wants to merge 2 commits into
base: main
Choose a base branch
from
+53 −27
Open

[Quantization] Fix per-tensor FP8 MoE parameter identity for EPLB hot-reload #19499

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7ce24c3
fix(quant): preserve FP8 MoE parameter identity for EPLB hot-reload
Socratesa Feb 27, 2026
94c2ca6
Merge branch 'main' into ds_coder_v2_lite_fp8_pr
Socratesa Feb 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 34 additions & 25 deletions python/sglang/srt/layers/quantization/fp8.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1147,25 +1147,23 @@ def process_weights_after_loading(self, layer: Module) -> None:
w13_weight = torch.empty_like(layer.w13_weight.data, dtype=fp8_dtype)
w2_weight = torch.empty_like(layer.w2_weight.data, dtype=fp8_dtype)

# Re-initialize w13_scale because we directly quantize
# merged w13 weights and generate a single scaling factor.
layer.w13_weight_scale = torch.nn.Parameter(
torch.ones(
layer.num_local_experts,
dtype=torch.float32,
device=w13_weight.device,
),
requires_grad=False,
)
# Quantize merged w13 weights and generate a single scaling factor
# per expert. We fill both columns of the original [E, 2] scale
# Parameter in-place (same value) to preserve its object identity
# and custom attributes (weight_loader, quant_method) for EPLB
# hot-reload, following the same principle as _copy_or_rebind in
# process_weights_after_loading_block_quant.
for expert in range(layer.num_local_experts):
w13_weight[expert, :, :], layer.w13_weight_scale[expert] = (
scaled_fp8_quant(layer.w13_weight.data[expert, :, :])
w13_weight[expert, :, :], scale = scaled_fp8_quant(
layer.w13_weight.data[expert, :, :]
)
w2_weight[expert, :, :], layer.w2_weight_scale[expert] = (
layer.w13_weight_scale.data[expert].fill_(scale)
w2_weight[expert, :, :], layer.w2_weight_scale.data[expert] = (
scaled_fp8_quant(layer.w2_weight.data[expert, :, :])
)
layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
# Rebind underlying tensor data while preserving Parameter objects.
layer.w13_weight.data = w13_weight
layer.w2_weight.data = w2_weight

if _is_hip:
self.process_weights_hip_scale_padding(layer)
Expand All @@ -1191,12 +1189,12 @@ def process_weights_after_loading(self, layer: Module) -> None:
"fp8 MoE layer. Using the maximum across experts "
"for each layer. "
)
layer.w13_input_scale = torch.nn.Parameter(
layer.w13_input_scale.max(), requires_grad=False
)
layer.w2_input_scale = torch.nn.Parameter(
layer.w2_input_scale.max(), requires_grad=False
)
# Fill in-place to preserve the Parameter object and its
# custom attributes (weight_loader, quant_method, etc.)
# needed for EPLB hot-reload, following the same principle
# as _copy_or_rebind in process_weights_after_loading_block_quant.
layer.w13_input_scale.data.fill_(layer.w13_input_scale.data.max())
layer.w2_input_scale.data.fill_(layer.w2_input_scale.data.max())

# If ROCm, normalize the weights and scales to e4m3fnuz
if _is_fp8_fnuz:
Expand Down Expand Up @@ -1246,9 +1244,14 @@ def process_weights_after_loading(self, layer: Module) -> None:
) = scaled_fp8_quant(dq_weight, max_w13_scales[expert_id])
start += shard_size

layer.w13_weight_scale = torch.nn.Parameter(
max_w13_scales, requires_grad=False
)
# Fill in-place to preserve the Parameter object and its
# custom attributes (weight_loader, quant_method, etc.)
# needed for EPLB hot-reload. Both columns of the [E, 2]
# tensor are set to the same per-expert max so that the
# weight_loader can reuse the original shape on reload.
# The runtime extracts a 1-D [E] view via [:, 0] in apply().
for expert_id in range(layer.num_local_experts):
layer.w13_weight_scale.data[expert_id].fill_(max_w13_scales[expert_id])

if _is_hip:
self.process_weights_hip_scale_padding(layer)
Expand Down Expand Up @@ -1459,8 +1462,14 @@ def apply(
block_shape = [scale_block_size, scale_block_size]
w13_scale_n = (w13_weight.shape[1] - 1) // scale_block_size + 1
w13_scale_k = (w13_weight.shape[2] - 1) // scale_block_size + 1
# w13_weight_scale is kept as [E, 2] for EPLB hot-reload
# compatibility; both columns hold the same per-expert max.
# Collapse to [E] before expanding to block shape.
_w13_ws = layer.w13_weight_scale
if _w13_ws.ndim == 2:
_w13_ws = _w13_ws[:, 0]
w13_scale = (
layer.w13_weight_scale.unsqueeze(1)
_w13_ws.unsqueeze(1)
.repeat_interleave(w13_scale_n, dim=1)
.unsqueeze(2)
.repeat_interleave(w13_scale_k, dim=2)
Expand Down
21 changes: 19 additions & 2 deletions python/sglang/srt/model_executor/model_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,11 @@
from sglang.srt.model_executor.piecewise_cuda_graph_runner import (
PiecewiseCudaGraphRunner,
)
from sglang.srt.model_loader.loader import DefaultModelLoader, get_model_loader
from sglang.srt.model_loader.loader import (
DefaultModelLoader,
device_loading_context,
get_model_loader,
)
from sglang.srt.model_loader.remote_instance_weight_loader_utils import (
RemoteInstanceWeightLoaderBackend,
register_memory_region,
Expand Down Expand Up @@ -1158,7 +1162,20 @@ def get_weight_iter(config):
return iter

def model_load_weights(model, iter):
loader.load_weights_and_postprocess(model, iter, target_device)
if weight_name_filter is not None:
# Partial reload (e.g. EPLB expert rebalancing): only load the
# filtered weights and run process_weights_after_loading on MoE
# modules whose weights were actually updated. Running it on
# ALL modules would double-process unrelated layers (e.g. the
# FP8 linear transpose in attention) and corrupt their weights.
model.load_weights(iter)
for name, module in model.named_modules():
quant_method = getattr(module, "quant_method", None)
if quant_method is not None and weight_name_filter(name):
with device_loading_context(module, target_device):
quant_method.process_weights_after_loading(module)
else:
loader.load_weights_and_postprocess(model, iter, target_device)
return model

with set_default_torch_dtype(self.model_config.dtype):
Expand Down
Loading