Skip to content

[Bugfix] Restore unloaded FP8 scale params during layerwise reload - #45835

Closed
aoshen02 wants to merge 1 commit into
vllm-project:mainfrom
aoshen02:fix/fp8-layerwise-scale-inv
Closed

[Bugfix] Restore unloaded FP8 scale params during layerwise reload#45835
aoshen02 wants to merge 1 commit into
vllm-project:mainfrom
aoshen02:fix/fp8-layerwise-scale-inv

Conversation

@aoshen02

Copy link
Copy Markdown
Collaborator

Summary

  • Fix NaN inference output when external weight sync provides FP8 weights without scale_inv tensors during layerwise reload
  • After replaying buffered weights, restore unloaded scale parameters from previously saved kernel tensors

Root cause

When RL training frameworks sync FP8 MoE weights to vLLM via start_weight_update/update_weights, they may send only the weight data without scale_inv tensors. During layerwise reload:

  1. materialize_layer allocates scale params as torch.empty (contains NaN)
  2. Weight replay only loads the params that were actually sent — scale params stay NaN
  3. process_weights_after_loading reads NaN scales and passes them to the MoE kernel
  4. All subsequent inference produces NaN

The fix checks for unloaded parameters after replay. If they contain NaN and previously saved kernel tensors exist (from before reload), it restores the valid scale values.

Not duplicating existing PRs

Test plan

  • Reproduce: load FP8 MoE model, call start_weight_update + update_weights with weights only (no scale_inv), verify NaN before fix
  • After fix: same flow produces valid inference output
  • Normal checkpoint reload (safetensors with scale_inv) still works unchanged

AI assistance disclosure

This fix was developed with AI assistance (Claude). The human submitter has reviewed all changed lines and understands the fix end-to-end.

Related: #41670

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 <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant