diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index cf3553bd57..0f6dccd847 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -6,8 +6,15 @@ from weakref import WeakKeyDictionary, ref import pytest import torch +from torch.nn.parameter import UninitializedParameter +import vllm.model_executor.model_loader.reload.meta as reload_meta from vllm.model_executor.layers.linear import QKVParallelLinear +from vllm.model_executor.model_loader.reload.layerwise import ( + finalize_layerwise_reload, + initialize_layerwise_reload, + record_metadata_for_reloading, +) from vllm.model_executor.model_loader.reload.meta import ( capture_layer_to_meta, get_numel_loaded, @@ -21,6 +28,38 @@ from vllm.model_executor.model_loader.reload.utils import get_layer_tensors from vllm.platforms import current_platform +class _AliasedBufferLayer(torch.nn.Module): + def __init__(self): + super().__init__() + weight = torch.arange(6, dtype=torch.float32).reshape(2, 3) + self.weight = torch.nn.Parameter(weight) + self.register_buffer( + "weight_view", self.weight.detach().view(-1), persistent=False + ) + + +class _ParentAliasedChildBufferLayer(torch.nn.Module): + def __init__(self): + super().__init__() + self.scale = torch.nn.Parameter(torch.ones(1)) + self.conv1d = torch.nn.Linear(3, 2, bias=False) + self.conv1d.weight.data.copy_( + torch.arange(6, dtype=torch.float32).reshape(2, 3) + ) + self.register_buffer( + "conv_weights", self.conv1d.weight.detach().view(-1), persistent=False + ) + + +class _AliasedBufferWithUninitializedChildLayer(_AliasedBufferLayer): + def __init__(self): + super().__init__() + self.child = torch.nn.Module() + self.child.register_parameter( + "lazy_weight", UninitializedParameter(requires_grad=False) + ) + + def test_move_metatensors(): tensor = torch.empty((1, 2, 3)) meta_tensor = to_meta_tensor(tensor) @@ -123,6 +162,81 @@ def test_get_numel_loaded(): assert ret == "value" +def test_layerwise_reload_skips_non_persistent_parameter_alias_buffers(monkeypatch): + layer = _AliasedBufferLayer() + model = torch.nn.Sequential(layer) + loaded_weight = torch.full_like(layer.weight, 7.0) + + def materialize_with_sentinel(meta_tensor): + tensor = torch.empty_strided( + size=tuple(meta_tensor.size()), + stride=tuple(meta_tensor.stride()), + dtype=meta_tensor.dtype, + requires_grad=False, + ) + tensor.fill_(-123.0) + tensor.__class__ = meta_tensor.__class__ + tensor.__dict__ = meta_tensor.__dict__.copy() + return tensor + + monkeypatch.setattr( + reload_meta, "materialize_meta_tensor", materialize_with_sentinel + ) + + record_metadata_for_reloading(model) + initialize_layerwise_reload(model) + layer.weight.weight_loader(layer.weight, loaded_weight) + finalize_layerwise_reload(model, model_config=None) + + assert torch.equal(layer.weight, loaded_weight) + assert layer.weight_view.untyped_storage().data_ptr() == ( + layer.weight.untyped_storage().data_ptr() + ) + + +def test_capture_layer_to_meta_skips_uninitialized_parameter_storage_ptrs(): + layer = _AliasedBufferWithUninitializedChildLayer() + + _, buffers = capture_layer_to_meta(layer) + + assert "weight_view" not in buffers + + +def test_layerwise_reload_skips_child_parameter_alias_buffers(monkeypatch): + layer = _ParentAliasedChildBufferLayer() + model = torch.nn.Sequential(layer) + loaded_conv = torch.full_like(layer.conv1d.weight, 7.0) + loaded_scale = torch.full_like(layer.scale, 3.0) + + def materialize_with_sentinel(meta_tensor): + tensor = torch.empty_strided( + size=tuple(meta_tensor.size()), + stride=tuple(meta_tensor.stride()), + dtype=meta_tensor.dtype, + requires_grad=False, + ) + tensor.fill_(-123.0) + tensor.__class__ = meta_tensor.__class__ + tensor.__dict__ = meta_tensor.__dict__.copy() + return tensor + + monkeypatch.setattr( + reload_meta, "materialize_meta_tensor", materialize_with_sentinel + ) + + record_metadata_for_reloading(model) + initialize_layerwise_reload(model) + layer.conv1d.weight.weight_loader(layer.conv1d.weight, loaded_conv) + layer.scale.weight_loader(layer.scale, loaded_scale) + finalize_layerwise_reload(model, model_config=None) + + assert torch.equal(layer.conv1d.weight, loaded_conv) + assert torch.equal(layer.conv_weights, loaded_conv.view(-1)) + assert layer.conv_weights.untyped_storage().data_ptr() == ( + layer.conv1d.weight.untyped_storage().data_ptr() + ) + + @pytest.mark.parametrize( "tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])] ) diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py index f8d8304199..40dd6dc9f3 100644 --- a/vllm/model_executor/model_loader/reload/layerwise.py +++ b/vllm/model_executor/model_loader/reload/layerwise.py @@ -380,6 +380,8 @@ def _copy_and_restore_kernel_tensors(layer: torch.nn.Module, info: LayerReloadin for name, param in parameters.items(): param.data.copy_(getattr(layer, name)) for name, buffer in buffers.items(): + if name not in layer._buffers: + continue buffer.data.copy_(getattr(layer, name)) _place_kernel_tensors(layer, info) diff --git a/vllm/model_executor/model_loader/reload/meta.py b/vllm/model_executor/model_loader/reload/meta.py index baa2081d58..397a458cbd 100644 --- a/vllm/model_executor/model_loader/reload/meta.py +++ b/vllm/model_executor/model_loader/reload/meta.py @@ -4,6 +4,7 @@ import inspect from collections.abc import Callable import torch +from torch.nn.parameter import UninitializedParameter from torch.utils._python_dispatch import TorchDispatchMode from .sanitize import restore_layer_refs, sanitize_layer_refs @@ -55,11 +56,45 @@ def materialize_meta_tensor(meta_tensor: torch.Tensor) -> torch.Tensor: return tensor +def _is_non_persistent_parameter_alias_buffer( + layer: torch.nn.Module, + name: str, + buffer: torch.Tensor, + parameter_storage_ptrs: set[int], +) -> bool: + if name not in layer._non_persistent_buffers_set: + return False + + buffer_storage_ptr = _tensor_storage_ptr(buffer) + return ( + buffer_storage_ptr is not None and buffer_storage_ptr in parameter_storage_ptrs + ) + + +def _tensor_storage_ptr(tensor: torch.Tensor) -> int | None: + if isinstance(tensor, UninitializedParameter): + return None + + try: + return tensor.untyped_storage().data_ptr() + except (RuntimeError, ValueError): + return None + + +def _parameter_storage_ptrs(layer: torch.nn.Module) -> set[int]: + return { + storage_ptr + for param in layer.parameters(recurse=True) + if (storage_ptr := _tensor_storage_ptr(param)) is not None + } + + def capture_layer_to_meta(layer: torch.nn.Module) -> LayerTensors: if layer.__class__.__name__ in SKIP_MODULES: return ({}, {}) params, buffers = get_layer_params_buffers(layer) + parameter_storage_ptrs = _parameter_storage_ptrs(layer) return ( { name: sanitize_layer_refs(to_meta_tensor(param), layer) @@ -70,6 +105,9 @@ def capture_layer_to_meta(layer: torch.nn.Module) -> LayerTensors: name: sanitize_layer_refs(to_meta_tensor(buffer), layer) for name, buffer in buffers.items() if name not in SKIP_TENSORS + and not _is_non_persistent_parameter_alias_buffer( + layer, name, buffer, parameter_storage_ptrs + ) }, )