From f6f770b390a0e2528111bddf8ba61d9de87bf57c Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 16 Jun 2026 15:06:12 +0000 Subject: [PATCH] [Bugfix] Restore unloaded FP8 scale params during layerwise reload When external weight sync (e.g. from RL training frameworks) provides FP8 weights without scale_inv tensors, layerwise reload leaves these scale parameters as uninitialized torch.empty (NaN). The subsequent process_weights_after_loading reads the NaN scales and passes them to the MoE kernel, producing all-NaN inference output. Fix: after replaying buffered weights, check for parameters that were not loaded during this reload cycle. If they contain NaN and we have previously saved kernel tensors (from before reload), restore the valid scale values. This ensures process_weights_after_loading always sees valid scales regardless of what the external weight provider sends. Affects: FP8 blockwise MoE models (e.g. Qwen3-30B-A3B-FP8) when used with RL frameworks that sync trainer weights to vLLM via start_weight_update / update_weights APIs. Related: #41670 Co-authored-by: Claude Signed-off-by: aoshen02 --- .../model_loader/reload/layerwise.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py index 6cf1c19cba43..3d1cb860ad99 100644 --- a/vllm/model_executor/model_loader/reload/layerwise.py +++ b/vllm/model_executor/model_loader/reload/layerwise.py @@ -342,10 +342,34 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo): param.weight_loader = _get_original_loader(param) # Load all buffered weights into materialized layer (using original loaders) + loaded_names = set() for name, args in info.loaded_weights: param = getattr(layer, name) args.arguments["param"] = param param.weight_loader(*args.args, **args.kwargs) + loaded_names.add(name) + + # During reload, scale parameters may not be provided by external weight + # sync (e.g. when an RL framework sends FP8 weights without scale_inv). + # These unloaded scales are still torch.empty (NaN) after materialization. + # Restore them from the previously saved kernel tensors so that + # process_weights_after_loading sees valid values. + if info.kernel_tensors is not None: + parameters, buffers = info.kernel_tensors + for name in list(parameters) + list(buffers): + if name in loaded_names: + continue + materialized = getattr(layer, name, None) + if materialized is None: + continue + if not materialized.is_floating_point(): + continue + if torch.isnan(materialized).any(): + saved = parameters.get(name) or buffers.get(name) + if saved is not None: + materialized.data.copy_(saved.data) + logger.debug("%s.%s: restored from kernel tensors", + layer.__class__.__name__, name) # Process weights (quantization, repacking, etc.) quant_method = getattr(layer, "quant_method", None)