Skip to content

[Bugfix] Fix layerwise weight reload: VllmConfig context + kernel-tensor copy #40647

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

Closed
Rorical wants to merge 2 commits into from
Closed

[Bugfix] Fix layerwise weight reload: VllmConfig context + kernel-tensor copy #40647

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2f39f06
[Bugfix] Re-enter VllmConfig context during layerwise weight reload
Rorical Apr 22, 2026
fb87ecf
[Bugfix] Only restore kernel tensors for weights actually reloaded
Rorical Apr 22, 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
52 changes: 45 additions & 7 deletions vllm/model_executor/model_loader/reload/layerwise.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,13 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import inspect
from collections.abc import Callable
from contextlib import nullcontext
from functools import wraps
from weakref import WeakKeyDictionary

import torch

from vllm.config import ModelConfig
from vllm.config import ModelConfig, get_current_vllm_config, set_current_vllm_config

@gemini-code-assist gemini-code-assist Bot Apr 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

It is recommended to use get_current_vllm_config_or_none to avoid broad exception handling when checking for the existence of a global configuration context.

Suggested change
from vllm.config import ModelConfig, get_current_vllm_config, set_current_vllm_config
from vllm.config import ModelConfig, get_current_vllm_config, set_current_vllm_config, get_current_vllm_config_or_none

from vllm.logger import init_logger
from vllm.model_executor.layers.attention import Attention, MLAAttention
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
Expand Down Expand Up @@ -43,6 +44,28 @@
WeakKeyDictionary()
)

# Capture VllmConfig at init so reload-path process_weights_after_loading
# (which runs outside set_current_vllm_config) can re-enter the context.
_cached_vllm_config = None


def _capture_vllm_config() -> None:
global _cached_vllm_config
try:
_cached_vllm_config = get_current_vllm_config()
except Exception:
pass


def _vllm_config_ctx():
if _cached_vllm_config is None:
return nullcontext()
try:
get_current_vllm_config()
return nullcontext()
except Exception:
return set_current_vllm_config(_cached_vllm_config)
Comment on lines +52 to +67

@gemini-code-assist gemini-code-assist Bot Apr 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Using try-except Exception: pass for flow control is generally discouraged and can lead to stale state. If get_current_vllm_config() fails during a subsequent capture (e.g., between different model initializations in the same process), the global _cached_vllm_config will retain its previous value instead of being reset to None. Using get_current_vllm_config_or_none() provides a cleaner implementation and ensures the cached state correctly reflects the current environment.

Suggested change
def _capture_vllm_config() -> None:
global _cached_vllm_config
try:
_cached_vllm_config = get_current_vllm_config()
except Exception:
pass
def _vllm_config_ctx():
if _cached_vllm_config is None:
return nullcontext()
try:
get_current_vllm_config()
return nullcontext()
except Exception:
return set_current_vllm_config(_cached_vllm_config)
def _capture_vllm_config() -> None:
global _cached_vllm_config
_cached_vllm_config = get_current_vllm_config_or_none()
def _vllm_config_ctx():
if _cached_vllm_config is None or get_current_vllm_config_or_none() is not None:
return nullcontext()
return set_current_vllm_config(_cached_vllm_config)



def get_layerwise_info(layer: torch.nn.Module) -> LayerReloadingInfo:
"""
Expand All @@ -65,6 +88,7 @@ def record_metadata_for_reloading(model: torch.nn.Module):
Stores parameter and buffer metadata as meta tensors for restoration.
Must be called before `initialize_layerwise_reload`.
"""
_capture_vllm_config()
for layer in model.modules():
info = get_layerwise_info(layer)
info.restore_metadata = capture_layer_to_meta(layer)
Expand All @@ -91,6 +115,10 @@ def initialize_layerwise_reload(model: torch.nn.Module):
model._original_do_torchao_reload = getattr(model, "_do_torchao_reload", False)
model._do_torchao_reload = False

# Fallback capture if record_metadata_for_reloading ran before ctx was set.
if _cached_vllm_config is None:
_capture_vllm_config()

for layer in model.modules():
info = get_layerwise_info(layer)

Expand Down Expand Up @@ -259,7 +287,8 @@ def _finalize_attention_layer(
)
else:
_place_kernel_tensors(layer, info)
layer.process_weights_after_loading(model_config.dtype)
with _vllm_config_ctx():
layer.process_weights_after_loading(model_config.dtype)


def _reload_attention_scales(layer: torch.nn.Module, info: LayerReloadingInfo) -> None:
Expand All @@ -282,7 +311,8 @@ def _reload_attention_scales(layer: torch.nn.Module, info: LayerReloadingInfo) -
args.arguments["param"] = param
_get_weight_loader(param)(*args.args, **args.kwargs)

quant_method.process_weights_after_loading(layer)
with _vllm_config_ctx():
quant_method.process_weights_after_loading(layer)

_copy_and_restore_kernel_tensors(layer, info)

Expand Down Expand Up @@ -318,7 +348,8 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo):
# Process weights (quantization, repacking, etc.)
quant_method = getattr(layer, "quant_method", None)
if isinstance(quant_method, QuantizeMethodBase):
quant_method.process_weights_after_loading(layer)
with _vllm_config_ctx():
quant_method.process_weights_after_loading(layer)

# Copy processed values into original tensor storage (preserves cudagraph refs)
# this code is a no-op if not reloading (because kernel tensors is empty)
Expand All @@ -344,13 +375,20 @@ def _get_weight_loader(tensor: torch.Tensor):

def _copy_and_restore_kernel_tensors(layer: torch.nn.Module, info: LayerReloadingInfo):
"""Copy processed values into original kernel tensor storage and restore
kernel tensor references on the layer. Preserves cudagraph references."""
kernel tensor references on the layer. Preserves cudagraph references.

Only copies tensors actually loaded this round; others (e.g. MambaMixer2's
`conv_weights` view, KV-scale sentinels) keep prior values to avoid
stamping uninitialized materialize_layer() data into shared storage."""
assert info.kernel_tensors is not None
loaded_names = {n for n, _ in info.loaded_weights}
parameters, buffers = info.kernel_tensors
for name, param in parameters.items():
param.data.copy_(getattr(layer, name))
if name in loaded_names:
param.data.copy_(getattr(layer, name))
for name, buffer in buffers.items():
buffer.data.copy_(getattr(layer, name))
if name in loaded_names:
buffer.data.copy_(getattr(layer, name))

_place_kernel_tensors(layer, info)

Expand Down
Loading