From 7ce24c315e8c2106402ba33eaf2e79149470705d Mon Sep 17 00:00:00 2001 From: Socratesa Date: Fri, 27 Feb 2026 09:15:01 +0000 Subject: [PATCH] fix(quant): preserve FP8 MoE parameter identity for EPLB hot-reload The per-tensor FP8 MoE path in process_weights_after_loading() replaced Parameter objects with new ones via torch.nn.Parameter(), destroying custom attributes (weight_loader, quant_method) needed by EPLB weight hot-reload. The block-quant path was already fixed (PR #17449) but the per-tensor path was missed. Additionally, update_weights_from_disk() unconditionally called process_weights_after_loading() on ALL modules even during partial reloads (e.g. EPLB expert rebalancing), causing non-expert layers like FP8 attention to be double-processed (double transpose -> shape mismatch). Changes in fp8.py: - Dynamic-quant path: use .data= for weight rebinding and .data[expert].fill_() for scale updates instead of Parameter replacement. - Checkpoint-FP8 path: use .data.fill_() for input_scale merging; fill both columns of the [E,2] w13_weight_scale in-place instead of replacing with a new [E] Parameter. - DeepGemm apply(): add ndim==2 guard to collapse w13_weight_scale from [E,2] to [E] via [:,0] before expanding to block shape. Changes in model_runner.py: - When weight_name_filter is set (EPLB expert rebalancing), split load_weights and process_weights_after_loading into two steps, only calling the latter on modules whose names match the filter. Signed-off-by: Socratesa --- python/sglang/srt/layers/quantization/fp8.py | 59 +++++++++++-------- .../sglang/srt/model_executor/model_runner.py | 21 ++++++- 2 files changed, 53 insertions(+), 27 deletions(-) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 6f6ec68d8eb3..4de92bf5e634 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -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) @@ -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: @@ -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) @@ -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) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 4fe3be6d4ee3..fe10517f84f4 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -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, @@ -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):