| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import gc |
| 4 | import importlib.machinery |
| 5 | import inspect |
| 6 | import sys |
| 7 | import types |
| 8 | from types import SimpleNamespace |
| 9 | from unittest.mock import Mock |
| 10 | from weakref import WeakKeyDictionary, ref |
| 11 | |
| 12 | import pytest |
| 13 | import torch |
| 14 | from torch.nn.parameter import UninitializedParameter |
| 15 | |
| 16 | import vllm.model_executor.model_loader.reload.layerwise as reload_layerwise |
| 17 | import vllm.model_executor.model_loader.reload.meta as reload_meta |
| 18 | from vllm.config import ModelConfig |
| 19 | from vllm.model_executor.layers.attention import MMEncoderAttention |
| 20 | from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase |
| 21 | from vllm.model_executor.layers.linear import QKVParallelLinear |
| 22 | from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase |
| 23 | from vllm.model_executor.model_loader.reload.layerwise import ( |
| 24 | finalize_layerwise_reload, |
| 25 | initialize_layerwise_reload, |
| 26 | initialize_online_processing, |
| 27 | record_metadata_for_reloading, |
| 28 | ) |
| 29 | from vllm.model_executor.model_loader.reload.meta import ( |
| 30 | capture_layer_to_meta, |
| 31 | get_numel_loaded, |
| 32 | materialize_layer, |
| 33 | materialize_meta_tensor, |
| 34 | restore_layer_on_meta, |
| 35 | to_meta_tensor, |
| 36 | ) |
| 37 | from vllm.model_executor.model_loader.reload.types import LayerReloadingInfo |
| 38 | from vllm.model_executor.model_loader.reload.utils import get_layer_tensors |
| 39 | from vllm.model_executor.model_loader.weight_utils import ( |
| 40 | composed_weight_loader, |
| 41 | default_weight_loader, |
| 42 | ) |
| 43 | from vllm.platforms import current_platform |
| 44 | |
| 45 | |
| 46 | def _fp8_reload_unsupported() -> bool: |
| 47 | """Whether the FP8 reload/online-quantize tests should be skipped. |
| 48 | |
| 49 | ``supports_fp8()`` returns True on MI250 (gfx90a) because the general |
| 50 | quantization paths upcast FP8 weights, but gfx90a has no native FP8 and |
| 51 | cannot run these reload models, so treat it as unsupported here. |
| 52 | """ |
| 53 | if not current_platform.supports_fp8(): |
| 54 | return True |
| 55 | if current_platform.is_rocm(): |
| 56 | from vllm.platforms.rocm import on_gfx90a |
| 57 | |
| 58 | return on_gfx90a() |
| 59 | return False |
| 60 | |
| 61 | |
| 62 | class _AliasedBufferLayer(torch.nn.Module): |
| 63 | def __init__(self): |
| 64 | super().__init__() |
| 65 | weight = torch.arange(6, dtype=torch.float32).reshape(2, 3) |
| 66 | self.weight = torch.nn.Parameter(weight) |
| 67 | self.register_buffer( |
| 68 | "weight_view", self.weight.detach().view(-1), persistent=False |
| 69 | ) |
| 70 | |
| 71 | |
| 72 | class _ParentAliasedChildBufferLayer(torch.nn.Module): |
| 73 | def __init__(self): |
| 74 | super().__init__() |
| 75 | self.scale = torch.nn.Parameter(torch.ones(1)) |
| 76 | self.conv1d = torch.nn.Linear(3, 2, bias=False) |
| 77 | self.conv1d.weight.data.copy_( |
| 78 | torch.arange(6, dtype=torch.float32).reshape(2, 3) |
| 79 | ) |
| 80 | self.register_buffer( |
| 81 | "conv_weights", self.conv1d.weight.detach().view(-1), persistent=False |
| 82 | ) |
| 83 | |
| 84 | |
| 85 | class _ChildAliasOnlyBufferLayer(torch.nn.Module): |
| 86 | def __init__(self): |
| 87 | super().__init__() |
| 88 | self.conv1d = torch.nn.Linear(3, 2, bias=False) |
| 89 | self.conv1d.weight.data.copy_( |
| 90 | torch.arange(6, dtype=torch.float32).reshape(2, 3) |
| 91 | ) |
| 92 | self.register_buffer( |
| 93 | "conv_weights", self.conv1d.weight.detach().view(-1), persistent=False |
| 94 | ) |
| 95 | |
| 96 | |
| 97 | class _AliasedBufferWithUninitializedChildLayer(_AliasedBufferLayer): |
| 98 | def __init__(self): |
| 99 | super().__init__() |
| 100 | self.child = torch.nn.Module() |
| 101 | self.child.register_parameter( |
| 102 | "lazy_weight", UninitializedParameter(requires_grad=False) |
| 103 | ) |
| 104 | |
| 105 | |
| 106 | class _NonPersistentBufferLayer(torch.nn.Module): |
| 107 | def __init__(self): |
| 108 | super().__init__() |
| 109 | self.weight = torch.nn.Parameter(torch.ones(2, 2)) |
| 110 | self.register_buffer("scale", torch.tensor(0.25), persistent=False) |
| 111 | |
| 112 | |
| 113 | class _ReloadableMMEncoderAttention(MMEncoderAttention): |
| 114 | """Minimal stand-in to test reload lifecycle without encoder initialization.""" |
| 115 | |
| 116 | def __init__(self): |
| 117 | torch.nn.Module.__init__(self) |
| 118 | self.weight = torch.nn.Parameter(torch.ones(2, 2)) |
| 119 | self.weight.weight_loader = default_weight_loader |
| 120 | self.post_load_called = False |
| 121 | |
| 122 | def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: |
| 123 | self.post_load_called = True |
| 124 | |
| 125 | |
| 126 | class _ReloadableAttentionLayer( |
| 127 | torch.nn.Module, |
| 128 | AttentionLayerBase, |
| 129 | ): |
| 130 | def __init__(self): |
| 131 | super().__init__() |
| 132 | self.weight = torch.nn.Parameter(torch.ones(2, 2)) |
| 133 | self.weight.weight_loader = default_weight_loader |
| 134 | self.post_load_called = False |
| 135 | |
| 136 | def get_attn_backend(self): |
| 137 | raise NotImplementedError |
| 138 | |
| 139 | def get_kv_cache_spec(self, vllm_config): |
| 140 | return None |
| 141 | |
| 142 | def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: |
| 143 | self.post_load_called = True |
| 144 | |
| 145 | |
| 146 | def test_move_metatensors(): |
| 147 | tensor = torch.empty((1, 2, 3)) |
| 148 | meta_tensor = to_meta_tensor(tensor) |
| 149 | materialized_tensor = materialize_meta_tensor(meta_tensor) |
| 150 | |
| 151 | assert meta_tensor.device.type == "meta" |
| 152 | assert tensor.device == materialized_tensor.device |
| 153 | |
| 154 | assert tensor.dtype == meta_tensor.dtype == materialized_tensor.dtype |
| 155 | assert tensor.shape == meta_tensor.shape == materialized_tensor.shape |
| 156 | assert tensor.__class__ == meta_tensor.__class__ == materialized_tensor.__class__ |
| 157 | assert tensor.__dict__ == meta_tensor.__dict__ == materialized_tensor.__dict__ |
| 158 | |
| 159 | |
| 160 | @pytest.mark.parametrize( |
| 161 | "layer_cls", |
| 162 | [_ReloadableMMEncoderAttention, _ReloadableAttentionLayer], |
| 163 | ) |
| 164 | def test_attention_reload_defers_post_load(default_vllm_config, layer_cls): |
| 165 | default_vllm_config.model_config = ModelConfig() |
| 166 | layer = layer_cls() |
| 167 | model = torch.nn.Sequential(layer) |
| 168 | loaded_weight = torch.full_like(layer.weight, 7.0) |
| 169 | |
| 170 | record_metadata_for_reloading(model) |
| 171 | initialize_layerwise_reload(model) |
| 172 | layer.weight.weight_loader(layer.weight, loaded_weight) |
| 173 | |
| 174 | assert not layer.post_load_called |
| 175 | |
| 176 | finalize_layerwise_reload(model, default_vllm_config.model_config) |
| 177 | |
| 178 | assert layer.post_load_called |
| 179 | assert torch.equal(layer.weight, loaded_weight) |
| 180 | |
| 181 | |
| 182 | @pytest.mark.parametrize( |
| 183 | "layer_cls", |
| 184 | [_ReloadableMMEncoderAttention, _ReloadableAttentionLayer], |
| 185 | ) |
| 186 | def test_attention_first_load_processes_weights(default_vllm_config, layer_cls): |
| 187 | default_vllm_config.model_config = ModelConfig() |
| 188 | layer = layer_cls() |
| 189 | model = torch.nn.Sequential(layer) |
| 190 | loaded_weight = torch.full_like(layer.weight, 7.0) |
| 191 | |
| 192 | initialize_online_processing(layer) |
| 193 | layer.weight.weight_loader(layer.weight, loaded_weight) |
| 194 | |
| 195 | finalize_layerwise_reload(model, default_vllm_config.model_config) |
| 196 | |
| 197 | assert layer.post_load_called |
| 198 | assert torch.equal(layer.weight, loaded_weight) |
| 199 | |
| 200 | |
| 201 | def test_reload_lifecycle(): |
| 202 | layer = torch.nn.Linear(2, 3) |
| 203 | info = LayerReloadingInfo( |
| 204 | restore_metadata=capture_layer_to_meta(layer), |
| 205 | restore_device=torch.device("cpu"), |
| 206 | ) |
| 207 | |
| 208 | restore_layer_on_meta(layer, info) |
| 209 | for name, tensor in get_layer_tensors(layer).items(): |
| 210 | meta_tensor = getattr(layer, name) |
| 211 | assert tensor.dtype == meta_tensor.dtype |
| 212 | assert tensor.shape == meta_tensor.shape |
| 213 | assert tensor.__class__ == meta_tensor.__class__ |
| 214 | assert tensor.__dict__ == meta_tensor.__dict__ |
| 215 | |
| 216 | materialize_layer(layer, info) |
| 217 | for name, tensor in get_layer_tensors(layer).items(): |
| 218 | materialized_tensor = getattr(layer, name) |
| 219 | assert tensor.dtype == materialized_tensor.dtype |
| 220 | assert tensor.shape == materialized_tensor.shape |
| 221 | assert tensor.__class__ == materialized_tensor.__class__ |
| 222 | assert tensor.__dict__ == materialized_tensor.__dict__ |
| 223 | |
| 224 | |
| 225 | def test_restore_layer_replaces_postprocessed_tensor_attribute(): |
| 226 | layer = torch.nn.Linear(2, 3, bias=False) |
| 227 | info = LayerReloadingInfo( |
| 228 | restore_metadata=capture_layer_to_meta(layer), |
| 229 | restore_device=torch.device("cpu"), |
| 230 | ) |
| 231 | del layer.weight |
| 232 | layer.weight = torch.empty(3, 2) |
| 233 | |
| 234 | restore_layer_on_meta(layer, info) |
| 235 | |
| 236 | assert isinstance(layer.weight, torch.nn.Parameter) |
| 237 | assert layer.weight.is_meta |
| 238 | |
| 239 | |
| 240 | def test_materialize_layer_preserves_non_meta_tensors(): |
| 241 | """Ensure that materialize_layer does not overwrite non meta tensors.""" |
| 242 | layer = torch.nn.Linear(2, 3, bias=True) |
| 243 | |
| 244 | # Create a non meta bias tensor and meta weight, which can happen with FP8 |
| 245 | bias_values = torch.ones(3) |
| 246 | layer.bias.data.copy_(bias_values) |
| 247 | layer.weight = torch.nn.Parameter(layer.weight.data.to("meta")) |
| 248 | |
| 249 | assert layer.weight.is_meta |
| 250 | assert not layer.bias.is_meta |
| 251 | |
| 252 | # materialize the layer weights after the bias is initialized |
| 253 | info = LayerReloadingInfo( |
| 254 | restore_metadata=({}, {}), |
| 255 | restore_device=torch.device("cpu"), |
| 256 | ) |
| 257 | materialize_layer(layer, info) |
| 258 | |
| 259 | # Ensure the weight materialized off meta |
| 260 | assert not layer.weight.is_meta |
| 261 | assert layer.weight.device.type == "cpu" |
| 262 | |
| 263 | # Ensure that the bias is (still) not meta and values are unchanged |
| 264 | assert not layer.bias.is_meta |
| 265 | assert torch.equal(layer.bias.data, bias_values) |
| 266 | |
| 267 | |
| 268 | _MARLIN_SIZE_K, _MARLIN_SIZE_N, _MARLIN_GROUP_SIZE = 128, 64, 64 |
| 269 | |
| 270 | |
| 271 | def _stub_marlin_ops(monkeypatch): |
| 272 | from vllm import _custom_ops as ops |
| 273 | from vllm.model_executor.layers.quantization.utils import marlin_utils |
| 274 | |
| 275 | monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4) |
| 276 | monkeypatch.setattr( |
| 277 | ops, |
| 278 | "gptq_marlin_repack", |
| 279 | lambda w, size_k, size_n, num_bits, is_a_8bit=False: torch.zeros( |
| 280 | size_k // 16, size_n * 2, dtype=torch.int32 |
| 281 | ), |
| 282 | ) |
| 283 | |
| 284 | |
| 285 | def _make_marlin_kernel(): |
| 286 | from vllm.model_executor.kernels.linear.mixed_precision.marlin import ( |
| 287 | MarlinLinearKernel, |
| 288 | ) |
| 289 | from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( |
| 290 | MPLinearLayerConfig, |
| 291 | ) |
| 292 | from vllm.scalar_type import scalar_types |
| 293 | |
| 294 | kernel = object.__new__(MarlinLinearKernel) |
| 295 | kernel.config = MPLinearLayerConfig( |
| 296 | full_weight_shape=(_MARLIN_SIZE_K, _MARLIN_SIZE_N), |
| 297 | partition_weight_shape=(_MARLIN_SIZE_K, _MARLIN_SIZE_N), |
| 298 | weight_type=scalar_types.uint4b8, |
| 299 | act_type=torch.float16, |
| 300 | group_size=_MARLIN_GROUP_SIZE, |
| 301 | zero_points=False, |
| 302 | ) |
| 303 | kernel.w_q_name = "qweight" |
| 304 | kernel.w_s_name = "scales" |
| 305 | kernel.w_zp_name = None |
| 306 | return kernel |
| 307 | |
| 308 | |
| 309 | def _load_marlin_checkpoint_format_weights(layer): |
| 310 | from vllm.model_executor.parameter import ( |
| 311 | GroupQuantScaleParameter, |
| 312 | PackedvLLMParameter, |
| 313 | ) |
| 314 | |
| 315 | layer.qweight = PackedvLLMParameter( |
| 316 | data=torch.zeros(_MARLIN_SIZE_K // 8, _MARLIN_SIZE_N, dtype=torch.int32), |
| 317 | input_dim=0, |
| 318 | output_dim=1, |
| 319 | packed_dim=0, |
| 320 | packed_factor=8, |
| 321 | weight_loader=default_weight_loader, |
| 322 | ) |
| 323 | layer.scales = GroupQuantScaleParameter( |
| 324 | data=torch.ones( |
| 325 | _MARLIN_SIZE_K // _MARLIN_GROUP_SIZE, _MARLIN_SIZE_N, dtype=torch.float16 |
| 326 | ), |
| 327 | input_dim=0, |
| 328 | output_dim=1, |
| 329 | weight_loader=default_weight_loader, |
| 330 | ) |
| 331 | |
| 332 | |
| 333 | def test_marlin_post_load_does_not_own_workspace(monkeypatch, dist_init): |
| 334 | """Weight reload must not create layer-owned Marlin lock storage.""" |
| 335 | _stub_marlin_ops(monkeypatch) |
| 336 | kernel = _make_marlin_kernel() |
| 337 | |
| 338 | layer = torch.nn.Module() |
| 339 | _load_marlin_checkpoint_format_weights(layer) |
| 340 | kernel.process_weights_after_loading(layer) |
| 341 | |
| 342 | assert not hasattr(kernel, "workspace") |
| 343 | |
| 344 | _load_marlin_checkpoint_format_weights(layer) |
| 345 | kernel.process_weights_after_loading(layer) |
| 346 | |
| 347 | assert not hasattr(kernel, "workspace") |
| 348 | |
| 349 | |
| 350 | @pytest.mark.parametrize("variant", ["fp8", "mxfp8", "nvfp4"]) |
| 351 | def test_marlin_prepare_layer_does_not_own_workspace(monkeypatch, variant): |
| 352 | """Weight preparation must not attach runtime workspace to model layers.""" |
| 353 | from vllm import _custom_ops as ops |
| 354 | from vllm.model_executor.layers.quantization.utils import ( |
| 355 | marlin_utils, |
| 356 | marlin_utils_fp4, |
| 357 | marlin_utils_fp8, |
| 358 | ) |
| 359 | |
| 360 | size_k, size_n = 128, 64 |
| 361 | |
| 362 | monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4) |
| 363 | monkeypatch.setattr( |
| 364 | ops, |
| 365 | "gptq_marlin_repack", |
| 366 | lambda b_q_weight, size_k, size_n, num_bits, is_a_8bit=False: torch.zeros( |
| 367 | size_k // 16, size_n * 2, dtype=torch.int32 |
| 368 | ), |
| 369 | ) |
| 370 | |
| 371 | layer = torch.nn.Module() |
| 372 | layer.output_size_per_partition = size_n |
| 373 | layer.input_size_per_partition = size_k |
| 374 | layer.orig_dtype = torch.float16 |
| 375 | layer.params_dtype = torch.float16 |
| 376 | |
| 377 | if variant == "fp8": |
| 378 | prepare = marlin_utils_fp8.prepare_fp8_layer_for_marlin |
| 379 | |
| 380 | def load_checkpoint_format_weights(): |
| 381 | layer.weight = torch.nn.Parameter( |
| 382 | torch.zeros(size_k, size_n, dtype=torch.float8_e4m3fn), |
| 383 | requires_grad=False, |
| 384 | ) |
| 385 | layer.weight_scale = torch.nn.Parameter( |
| 386 | torch.ones(1, dtype=torch.float32), requires_grad=False |
| 387 | ) |
| 388 | elif variant == "mxfp8": |
| 389 | prepare = marlin_utils_fp8.prepare_mxfp8_layer_for_marlin |
| 390 | |
| 391 | def load_checkpoint_format_weights(): |
| 392 | layer.weight = torch.nn.Parameter( |
| 393 | torch.zeros(size_n, size_k, dtype=torch.float8_e4m3fn), |
| 394 | requires_grad=False, |
| 395 | ) |
| 396 | layer.weight_scale = torch.nn.Parameter( |
| 397 | torch.full((size_n, size_k // 32), 127, dtype=torch.uint8), |
| 398 | requires_grad=False, |
| 399 | ) |
| 400 | else: |
| 401 | prepare = marlin_utils_fp4.prepare_fp4_layer_for_marlin |
| 402 | |
| 403 | def load_checkpoint_format_weights(): |
| 404 | layer.weight = torch.nn.Parameter( |
| 405 | torch.zeros(size_n, size_k // 2, dtype=torch.uint8), |
| 406 | requires_grad=False, |
| 407 | ) |
| 408 | layer.weight_scale = torch.nn.Parameter( |
| 409 | torch.ones(size_n, size_k // 16, dtype=torch.float8_e4m3fn), |
| 410 | requires_grad=False, |
| 411 | ) |
| 412 | layer.weight_global_scale = torch.nn.Parameter( |
| 413 | torch.ones(1, dtype=torch.float32), requires_grad=False |
| 414 | ) |
| 415 | |
| 416 | load_checkpoint_format_weights() |
| 417 | prepare(layer) |
| 418 | assert not hasattr(layer, "workspace") |
| 419 | |
| 420 | # Reload: fresh checkpoint-format tensors, prepare runs again |
| 421 | load_checkpoint_format_weights() |
| 422 | prepare(layer) |
| 423 | |
| 424 | assert not hasattr(layer, "workspace") |
| 425 | |
| 426 | |
| 427 | def test_marlin_workspace_uses_persistent_workspace_manager(monkeypatch): |
| 428 | """Calls reuse initialized locks; independent streams get separate storage.""" |
| 429 | from vllm.model_executor.layers.quantization.utils import marlin_utils |
| 430 | from vllm.utils import torch_utils |
| 431 | from vllm.v1.worker import workspace as workspace_module |
| 432 | |
| 433 | monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4) |
| 434 | device = torch.device("cpu") |
| 435 | manager = workspace_module.WorkspaceManager(device) |
| 436 | monkeypatch.setattr(workspace_module, "_manager", manager) |
| 437 | stream = "main" |
| 438 | monkeypatch.setattr(torch_utils, "current_stream", lambda: stream) |
| 439 | |
| 440 | workspace = marlin_utils.get_marlin_workspace(device) |
| 441 | assert workspace.shape == (4 * marlin_utils.MARLIN_MAX_BLOCKS_PER_SM,) |
| 442 | assert workspace.dtype == torch.int32 |
| 443 | assert torch.count_nonzero(workspace) == 0 |
| 444 | |
| 445 | workspace.fill_(1) |
| 446 | assert marlin_utils.get_marlin_workspace(device) is workspace |
| 447 | assert torch.all(workspace == 1) |
| 448 | |
| 449 | stream = "aux" |
| 450 | aux_workspace = marlin_utils.get_marlin_workspace(device) |
| 451 | assert aux_workspace.data_ptr() != workspace.data_ptr() |
| 452 | assert torch.count_nonzero(aux_workspace) == 0 |
| 453 | |
| 454 | manager.lock() |
| 455 | assert marlin_utils.get_marlin_workspace(device) is aux_workspace |
| 456 | |
| 457 | |
| 458 | def test_marlin_workspace_without_manager(monkeypatch): |
| 459 | from vllm.model_executor.layers.quantization.utils import marlin_utils |
| 460 | from vllm.v1.worker import workspace as workspace_module |
| 461 | |
| 462 | monkeypatch.setattr(workspace_module, "_manager", None) |
| 463 | monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4) |
| 464 | first = marlin_utils.get_marlin_workspace(torch.device("cpu")) |
| 465 | second = marlin_utils.get_marlin_workspace(torch.device("cpu")) |
| 466 | assert first.data_ptr() != second.data_ptr() |
| 467 | assert first.shape == (4 * marlin_utils.MARLIN_MAX_BLOCKS_PER_SM,) |
| 468 | assert torch.count_nonzero(first) == 0 |
| 469 | |
| 470 | |
| 471 | def test_model_cleanup(dist_init, default_vllm_config): |
| 472 | layer = QKVParallelLinear(2, 3, 4) |
| 473 | assert layer.weight.weight_loader.__self__ is layer |
| 474 | info = LayerReloadingInfo( |
| 475 | restore_metadata=capture_layer_to_meta(layer), |
| 476 | restore_device=torch.device("cpu"), |
| 477 | ) |
| 478 | |
| 479 | mock_info_dict: WeakKeyDictionary[torch.nn.Module, LayerReloadingInfo] = ( |
| 480 | WeakKeyDictionary() |
| 481 | ) |
| 482 | mock_info_dict[layer] = info |
| 483 | layer_ref = ref(layer) |
| 484 | |
| 485 | del layer |
| 486 | gc.collect() |
| 487 | |
| 488 | assert layer_ref() is None |
| 489 | assert len(mock_info_dict) == 0 |
| 490 | |
| 491 | |
| 492 | @pytest.mark.parametrize("is_gated", [False, True]) |
| 493 | @pytest.mark.parametrize("has_bias", [False, True]) |
| 494 | @pytest.mark.parametrize("tp_rank", [0, 1]) |
| 495 | @pytest.mark.parametrize("padded", [False, True]) |
| 496 | def test_padded_moe_reload_releases_each_layer( |
| 497 | monkeypatch, is_gated, has_bias, tp_rank, padded |
| 498 | ): |
| 499 | """Checkpoint-sized copies finish each layer without global finalization.""" |
| 500 | from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts |
| 501 | from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( |
| 502 | UnquantizedFusedMoEMethod, |
| 503 | ) |
| 504 | |
| 505 | hidden, intermediate, experts = 4, 3, 2 |
| 506 | stored_hidden, stored_intermediate = (8, 8) if padded else (hidden, intermediate) |
| 507 | config = SimpleNamespace( |
| 508 | hidden_dim_unpadded=hidden, |
| 509 | intermediate_size_per_partition_unpadded=intermediate, |
| 510 | is_act_and_mul=is_gated, |
| 511 | has_bias=has_bias, |
| 512 | tp_rank=tp_rank, |
| 513 | tp_shard_with_padding=False, |
| 514 | moe_parallel_config=SimpleNamespace(tp_size=2), |
| 515 | ) |
| 516 | model = torch.nn.ModuleList() |
| 517 | processed: list[torch.nn.Module] = [] |
| 518 | for _ in range(2): |
| 519 | method = object.__new__(UnquantizedFusedMoEMethod) |
| 520 | torch.nn.Module.__init__(method) |
| 521 | method.moe = config |
| 522 | # The regression concerns streaming reload, not kernel conversion. |
| 523 | monkeypatch.setattr(method, "process_weights_after_loading", processed.append) |
| 524 | layer = object.__new__(RoutedExperts) |
| 525 | torch.nn.Module.__init__(layer) |
| 526 | layer.moe_config = config |
| 527 | layer.quant_config = None |
| 528 | layer.quant_method = method |
| 529 | layer.expert_map_manager = SimpleNamespace(map_global_to_local=lambda i: i) |
| 530 | layer._loaded_expert_biases = set() |
| 531 | method.create_weights( |
| 532 | layer, |
| 533 | experts, |
| 534 | stored_hidden, |
| 535 | stored_intermediate, |
| 536 | torch.float32, |
| 537 | weight_loader=layer.weight_loader, |
| 538 | ) |
| 539 | model.append(layer) |
| 540 | |
| 541 | record_metadata_for_reloading(model) |
| 542 | original_params = [dict(layer.named_parameters()) for layer in model] |
| 543 | shards = ["w1", "w3", "w2"] if is_gated else ["w1", "w2"] |
| 544 | for cycle in range(2): |
| 545 | initialize_layerwise_reload(model) |
| 546 | for layer_index, layer in enumerate(model): |
| 547 | info = reload_layerwise.get_layerwise_info(layer) |
| 548 | inputs = [] |
| 549 | expected = { |
| 550 | name: torch.full_like(p, float("nan")) |
| 551 | for name, p in original_params[layer_index].items() |
| 552 | } |
| 553 | params = dict(layer.named_parameters()) |
| 554 | calls = [ |
| 555 | (e, s, b) |
| 556 | for e in range(experts) |
| 557 | for s in shards |
| 558 | for b in ([False, True] if has_bias else [False]) |
| 559 | ] |
| 560 | for call_index, (expert, shard, bias) in enumerate(calls): |
| 561 | name = ("w2" if shard == "w2" else "w13") + ( |
| 562 | "_bias" if bias else "_weight" |
| 563 | ) |
| 564 | shape = ( |
| 565 | ((hidden,) if bias else (hidden, 2 * intermediate)) |
| 566 | if shard == "w2" |
| 567 | else ((2 * intermediate,) if bias else (2 * intermediate, hidden)) |
| 568 | ) |
| 569 | weight = torch.arange( |
| 570 | torch.Size(shape).numel(), dtype=torch.float32 |
| 571 | ).reshape(shape) |
| 572 | weight = weight + 100 * (1 + call_index + cycle) |
| 573 | inputs.append(ref(weight)) |
| 574 | # Direct checkpoint loading is the reference for deferred reload. |
| 575 | layer.weight_loader(expected[name], weight, name, shard, expert) |
| 576 | param = params[name] |
| 577 | param.weight_loader(param, weight, name, shard, expert) |
| 578 | del weight |
| 579 | if call_index != len(calls) - 1: |
| 580 | assert info.can_load(), ( |
| 581 | "Layer processed before its final checkpoint shard" |
| 582 | ) |
| 583 | |
| 584 | assert not info.can_load(), ( |
| 585 | "Padding must not defer the layer until finalization" |
| 586 | ) |
| 587 | assert not info.loaded_weights |
| 588 | assert len(processed) == cycle * len(model) + layer_index + 1 |
| 589 | assert all(source() is None for source in inputs) |
| 590 | for name, original in original_params[layer_index].items(): |
| 591 | assert getattr(layer, name) is original |
| 592 | # Kernel-specific tests cover padding, which is not checkpoint data. |
| 593 | mask = torch.isfinite(expected[name]) |
| 594 | assert torch.equal(original[mask], expected[name][mask]) |
| 595 | |
| 596 | |
| 597 | def test_get_numel_loaded(): |
| 598 | param = torch.empty(10, device="meta") |
| 599 | loaded_weight = torch.empty(10) |
| 600 | |
| 601 | def complex_weight_loader(param, loaded_weight): |
| 602 | param[:3] = loaded_weight[:3] |
| 603 | param[5:8] = loaded_weight[5:8] |
| 604 | return "value" |
| 605 | |
| 606 | args = inspect.signature(complex_weight_loader).bind(param, loaded_weight) |
| 607 | num_loaded, ret = get_numel_loaded(complex_weight_loader, args) |
| 608 | assert num_loaded == 6 |
| 609 | assert ret == "value" |
| 610 | |
| 611 | |
| 612 | def test_get_numel_loaded_caps_at_param_size(): |
| 613 | # composed_weight_loader copies into the param twice (the load and the |
| 614 | # in-place post-load transform), but only param.numel() distinct elements |
| 615 | # are loaded. get_numel_loaded must not double-count, otherwise a layer's |
| 616 | # loaded-element total can be reached early and trailing params get dropped. |
| 617 | param = torch.empty(10) |
| 618 | loaded_weight = torch.ones(10) |
| 619 | loader = composed_weight_loader(default_weight_loader, lambda x: x + 1) |
| 620 | |
| 621 | args = inspect.signature(loader).bind(param, loaded_weight) |
| 622 | num_loaded, _ = get_numel_loaded(loader, args) |
| 623 | assert num_loaded == 10 |
| 624 | |
| 625 | |
| 626 | def test_layerwise_loading_warning_only_checks_new_layers(monkeypatch): |
| 627 | layers = [torch.nn.Linear(16, 1, bias=False) for _ in range(2)] |
| 628 | |
| 629 | def partial_weight_loader(param, loaded_weight): |
| 630 | param.view(-1)[: loaded_weight.numel()].copy_(loaded_weight) |
| 631 | |
| 632 | for layer in layers: |
| 633 | layer.weight.requires_grad_(False) |
| 634 | layer.weight.weight_loader = partial_weight_loader |
| 635 | reload_layerwise.initialize_online_processing(layer) |
| 636 | |
| 637 | monkeypatch.setattr(reload_layerwise, "has_device_tensors", lambda _: True) |
| 638 | get_info_size = Mock(return_value=0) |
| 639 | warning_once = Mock() |
| 640 | monkeypatch.setattr(reload_layerwise, "get_info_size", get_info_size) |
| 641 | monkeypatch.setattr(reload_layerwise.logger, "warning_once", warning_once) |
| 642 | |
| 643 | reload_layerwise.LOADING_LAYERS.clear() |
| 644 | try: |
| 645 | for layer in layers: |
| 646 | for _ in range(3): |
| 647 | layer.weight.weight_loader(layer.weight, torch.ones(1)) |
| 648 | finally: |
| 649 | reload_layerwise.LOADING_LAYERS.clear() |
| 650 | |
| 651 | assert get_info_size.call_count == 2 |
| 652 | warning_once.assert_called_once() |
| 653 | |
| 654 | |
| 655 | class _ComposedLoaderLayer(torch.nn.Module): |
| 656 | """Mimics a Mamba2 mixer's equal-numel direct params (A, D, dt_bias). |
| 657 | |
| 658 | ``A`` uses ``composed_weight_loader`` (an extra in-place transform copy), |
| 659 | matching ``MambaMixer2`` where ``A`` is loaded as ``-exp(A_log)``. |
| 660 | """ |
| 661 | |
| 662 | def __init__(self): |
| 663 | super().__init__() |
| 664 | self.A = torch.nn.Parameter(torch.empty(4, dtype=torch.float32)) |
| 665 | self.D = torch.nn.Parameter(torch.ones(4)) |
| 666 | self.dt_bias = torch.nn.Parameter(torch.ones(4)) |
| 667 | self.A.weight_loader = composed_weight_loader( |
| 668 | default_weight_loader, lambda x: -torch.exp(x.float()) |
| 669 | ) |
| 670 | self.D.weight_loader = default_weight_loader |
| 671 | self.dt_bias.weight_loader = default_weight_loader |
| 672 | |
| 673 | |
| 674 | def test_layerwise_reload_composed_loader_does_not_drop_params(monkeypatch): |
| 675 | # Regression test: a composed_weight_loader param (A) used to double-count |
| 676 | # its elements, finalizing the layer before the trailing param (D) was |
| 677 | # loaded and leaving it as uninitialized materialized memory. |
| 678 | layer = _ComposedLoaderLayer() |
| 679 | model = torch.nn.Sequential(layer) |
| 680 | |
| 681 | def materialize_with_sentinel(meta_tensor): |
| 682 | tensor = torch.empty_strided( |
| 683 | size=tuple(meta_tensor.size()), |
| 684 | stride=tuple(meta_tensor.stride()), |
| 685 | dtype=meta_tensor.dtype, |
| 686 | requires_grad=False, |
| 687 | ) |
| 688 | tensor.fill_(float("nan")) |
| 689 | tensor.__class__ = meta_tensor.__class__ |
| 690 | tensor.__dict__ = meta_tensor.__dict__.copy() |
| 691 | return tensor |
| 692 | |
| 693 | monkeypatch.setattr( |
| 694 | reload_meta, "materialize_meta_tensor", materialize_with_sentinel |
| 695 | ) |
| 696 | |
| 697 | loaded = { |
| 698 | "A": torch.full((4,), 0.5), |
| 699 | "dt_bias": torch.full((4,), 3.0), |
| 700 | "D": torch.full((4,), 7.0), |
| 701 | } |
| 702 | |
| 703 | record_metadata_for_reloading(model) |
| 704 | initialize_layerwise_reload(model) |
| 705 | # Mimic real load_weights: resolve params once, then load in checkpoint |
| 706 | # order with D last (the param that was dropped). |
| 707 | params = dict(layer.named_parameters()) |
| 708 | for name in ("A", "dt_bias", "D"): |
| 709 | param = params[name] |
| 710 | param.weight_loader(param, loaded[name]) |
| 711 | finalize_layerwise_reload(model, model_config=None) |
| 712 | |
| 713 | assert torch.equal(layer.A, -torch.exp(loaded["A"])) |
| 714 | assert torch.equal(layer.dt_bias, loaded["dt_bias"]) |
| 715 | assert torch.equal(layer.D, loaded["D"]) |
| 716 | |
| 717 | |
| 718 | class _RecordingQuantMethod(QuantizeMethodBase): |
| 719 | """Records the layer's bias at the moment processing runs.""" |
| 720 | |
| 721 | uses_meta_device = True |
| 722 | |
| 723 | def __init__(self): |
| 724 | self.bias_at_process = None |
| 725 | |
| 726 | def create_weights(self, layer, *weight_args, **extra_weight_attrs): |
| 727 | pass |
| 728 | |
| 729 | def apply(self, layer, *args, **kwargs): |
| 730 | raise NotImplementedError |
| 731 | |
| 732 | def process_weights_after_loading(self, layer): |
| 733 | self.bias_at_process = layer.bias.detach().clone() |
| 734 | |
| 735 | |
| 736 | class _LateBiasLayer(torch.nn.Module): |
| 737 | """Mimics an online-quantized linear: `weight` is created on meta by |
| 738 | `create_weights()`, which wraps the loaders, and the linear base registers |
| 739 | `bias` afterwards.""" |
| 740 | |
| 741 | def __init__(self, quant_method): |
| 742 | super().__init__() |
| 743 | self.quant_method = quant_method |
| 744 | weight = torch.nn.Parameter(torch.empty(4, 2, device="meta")) |
| 745 | weight.weight_loader = default_weight_loader |
| 746 | self.register_parameter("weight", weight) |
| 747 | initialize_online_processing(self) |
| 748 | bias = torch.nn.Parameter(torch.zeros(4)) |
| 749 | bias.weight_loader = default_weight_loader |
| 750 | self.register_parameter("bias", bias) |
| 751 | |
| 752 | |
| 753 | def test_online_processing_waits_for_late_registered_bias(): |
| 754 | # Regression test: `bias` is skipped by the meta device paths, but it is |
| 755 | # still loaded by a weight loader. Excluding it from the processing trigger |
| 756 | # finalized the layer one load early, so the trailing bias was written into |
| 757 | # an already-processed layer (e.g. over FP8 Marlin's permuted bias). |
| 758 | quant_method = _RecordingQuantMethod() |
| 759 | layer = _LateBiasLayer(quant_method) |
| 760 | loaded_bias = torch.full((4,), 3.0) |
| 761 | |
| 762 | layer.weight.weight_loader(layer.weight, torch.full((4, 2), 2.0)) |
| 763 | assert quant_method.bias_at_process is None |
| 764 | |
| 765 | layer.bias.weight_loader(layer.bias, loaded_bias) |
| 766 | assert quant_method.bias_at_process is not None |
| 767 | assert torch.equal(quant_method.bias_at_process, loaded_bias) |
| 768 | |
| 769 | |
| 770 | def test_layerwise_reload_skips_non_persistent_parameter_alias_buffers(monkeypatch): |
| 771 | layer = _AliasedBufferLayer() |
| 772 | model = torch.nn.Sequential(layer) |
| 773 | loaded_weight = torch.full_like(layer.weight, 7.0) |
| 774 | |
| 775 | def materialize_with_sentinel(meta_tensor): |
| 776 | tensor = torch.empty_strided( |
| 777 | size=tuple(meta_tensor.size()), |
| 778 | stride=tuple(meta_tensor.stride()), |
| 779 | dtype=meta_tensor.dtype, |
| 780 | requires_grad=False, |
| 781 | ) |
| 782 | tensor.fill_(-123.0) |
| 783 | tensor.__class__ = meta_tensor.__class__ |
| 784 | tensor.__dict__ = meta_tensor.__dict__.copy() |
| 785 | return tensor |
| 786 | |
| 787 | monkeypatch.setattr( |
| 788 | reload_meta, "materialize_meta_tensor", materialize_with_sentinel |
| 789 | ) |
| 790 | |
| 791 | record_metadata_for_reloading(model) |
| 792 | initialize_layerwise_reload(model) |
| 793 | layer.weight.weight_loader(layer.weight, loaded_weight) |
| 794 | finalize_layerwise_reload(model, model_config=None) |
| 795 | |
| 796 | assert torch.equal(layer.weight, loaded_weight) |
| 797 | assert layer.weight_view.untyped_storage().data_ptr() == ( |
| 798 | layer.weight.untyped_storage().data_ptr() |
| 799 | ) |
| 800 | assert "weight_view" in layer._non_persistent_buffers_set |
| 801 | assert "0.weight_view" not in model.state_dict() |
| 802 | |
| 803 | |
| 804 | def test_capture_layer_to_meta_skips_uninitialized_parameter_storage_ptrs(): |
| 805 | layer = _AliasedBufferWithUninitializedChildLayer() |
| 806 | |
| 807 | _, buffers = capture_layer_to_meta(layer) |
| 808 | |
| 809 | assert "weight_view" not in buffers |
| 810 | |
| 811 | |
| 812 | def test_layerwise_reload_skips_child_parameter_alias_buffers(monkeypatch): |
| 813 | layer = _ParentAliasedChildBufferLayer() |
| 814 | model = torch.nn.Sequential(layer) |
| 815 | loaded_conv = torch.full_like(layer.conv1d.weight, 7.0) |
| 816 | loaded_scale = torch.full_like(layer.scale, 3.0) |
| 817 | |
| 818 | def materialize_with_sentinel(meta_tensor): |
| 819 | tensor = torch.empty_strided( |
| 820 | size=tuple(meta_tensor.size()), |
| 821 | stride=tuple(meta_tensor.stride()), |
| 822 | dtype=meta_tensor.dtype, |
| 823 | requires_grad=False, |
| 824 | ) |
| 825 | tensor.fill_(-123.0) |
| 826 | tensor.__class__ = meta_tensor.__class__ |
| 827 | tensor.__dict__ = meta_tensor.__dict__.copy() |
| 828 | return tensor |
| 829 | |
| 830 | monkeypatch.setattr( |
| 831 | reload_meta, "materialize_meta_tensor", materialize_with_sentinel |
| 832 | ) |
| 833 | |
| 834 | record_metadata_for_reloading(model) |
| 835 | initialize_layerwise_reload(model) |
| 836 | layer.conv1d.weight.weight_loader(layer.conv1d.weight, loaded_conv) |
| 837 | layer.scale.weight_loader(layer.scale, loaded_scale) |
| 838 | finalize_layerwise_reload(model, model_config=None) |
| 839 | |
| 840 | assert torch.equal(layer.conv1d.weight, loaded_conv) |
| 841 | assert torch.equal(layer.conv_weights, loaded_conv.view(-1)) |
| 842 | assert layer.conv_weights.untyped_storage().data_ptr() == ( |
| 843 | layer.conv1d.weight.untyped_storage().data_ptr() |
| 844 | ) |
| 845 | assert "conv_weights" in layer._non_persistent_buffers_set |
| 846 | assert "0.conv_weights" not in model.state_dict() |
| 847 | |
| 848 | |
| 849 | def test_layerwise_reload_restores_alias_buffer_on_zero_size_layer(monkeypatch): |
| 850 | layer = _ChildAliasOnlyBufferLayer() |
| 851 | model = torch.nn.Sequential(layer) |
| 852 | loaded_conv = torch.full_like(layer.conv1d.weight, 7.0) |
| 853 | |
| 854 | def materialize_with_sentinel(meta_tensor): |
| 855 | tensor = torch.empty_strided( |
| 856 | size=tuple(meta_tensor.size()), |
| 857 | stride=tuple(meta_tensor.stride()), |
| 858 | dtype=meta_tensor.dtype, |
| 859 | requires_grad=False, |
| 860 | ) |
| 861 | tensor.fill_(-123.0) |
| 862 | tensor.__class__ = meta_tensor.__class__ |
| 863 | tensor.__dict__ = meta_tensor.__dict__.copy() |
| 864 | return tensor |
| 865 | |
| 866 | monkeypatch.setattr( |
| 867 | reload_meta, "materialize_meta_tensor", materialize_with_sentinel |
| 868 | ) |
| 869 | |
| 870 | record_metadata_for_reloading(model) |
| 871 | initialize_layerwise_reload(model) |
| 872 | layer.conv1d.weight.weight_loader(layer.conv1d.weight, loaded_conv) |
| 873 | finalize_layerwise_reload(model, model_config=None) |
| 874 | |
| 875 | assert torch.equal(layer.conv_weights, loaded_conv.view(-1)) |
| 876 | assert layer.conv_weights.untyped_storage().data_ptr() == ( |
| 877 | layer.conv1d.weight.untyped_storage().data_ptr() |
| 878 | ) |
| 879 | assert "conv_weights" in layer._non_persistent_buffers_set |
| 880 | assert "0.conv_weights" not in model.state_dict() |
| 881 | |
| 882 | |
| 883 | def test_layerwise_reload_preserves_unloaded_non_persistent_buffers(monkeypatch): |
| 884 | layer = _NonPersistentBufferLayer() |
| 885 | model = torch.nn.Sequential(layer) |
| 886 | loaded_weight = torch.full_like(layer.weight, 7.0) |
| 887 | original_scale = layer.scale.clone() |
| 888 | |
| 889 | def materialize_with_sentinel(meta_tensor): |
| 890 | tensor = torch.empty_strided( |
| 891 | size=tuple(meta_tensor.size()), |
| 892 | stride=tuple(meta_tensor.stride()), |
| 893 | dtype=meta_tensor.dtype, |
| 894 | requires_grad=False, |
| 895 | ) |
| 896 | tensor.fill_(-123.0) |
| 897 | tensor.__class__ = meta_tensor.__class__ |
| 898 | tensor.__dict__ = meta_tensor.__dict__.copy() |
| 899 | return tensor |
| 900 | |
| 901 | monkeypatch.setattr( |
| 902 | reload_meta, "materialize_meta_tensor", materialize_with_sentinel |
| 903 | ) |
| 904 | |
| 905 | record_metadata_for_reloading(model) |
| 906 | initialize_layerwise_reload(model) |
| 907 | layer.weight.weight_loader(layer.weight, loaded_weight) |
| 908 | finalize_layerwise_reload(model, model_config=None) |
| 909 | |
| 910 | assert torch.equal(layer.weight, loaded_weight) |
| 911 | assert torch.equal(layer.scale, original_scale) |
| 912 | assert "scale" in layer._non_persistent_buffers_set |
| 913 | assert "0.scale" not in model.state_dict() |
| 914 | |
| 915 | |
| 916 | def test_layerwise_reload_updates_loaded_non_persistent_buffers(monkeypatch): |
| 917 | layer = _NonPersistentBufferLayer() |
| 918 | model = torch.nn.Sequential(layer) |
| 919 | loaded_weight = torch.full_like(layer.weight, 7.0) |
| 920 | loaded_scale = torch.full_like(layer.scale, 0.5) |
| 921 | |
| 922 | def materialize_with_sentinel(meta_tensor): |
| 923 | tensor = torch.empty_strided( |
| 924 | size=tuple(meta_tensor.size()), |
| 925 | stride=tuple(meta_tensor.stride()), |
| 926 | dtype=meta_tensor.dtype, |
| 927 | requires_grad=False, |
| 928 | ) |
| 929 | tensor.fill_(-123.0) |
| 930 | tensor.__class__ = meta_tensor.__class__ |
| 931 | tensor.__dict__ = meta_tensor.__dict__.copy() |
| 932 | return tensor |
| 933 | |
| 934 | monkeypatch.setattr( |
| 935 | reload_meta, "materialize_meta_tensor", materialize_with_sentinel |
| 936 | ) |
| 937 | |
| 938 | record_metadata_for_reloading(model) |
| 939 | initialize_layerwise_reload(model) |
| 940 | layer.weight.weight_loader(layer.weight, loaded_weight) |
| 941 | layer.scale.weight_loader(layer.scale, loaded_scale) |
| 942 | finalize_layerwise_reload(model, model_config=None) |
| 943 | |
| 944 | assert torch.equal(layer.weight, loaded_weight) |
| 945 | assert torch.equal(layer.scale, loaded_scale) |
| 946 | assert "scale" in layer._non_persistent_buffers_set |
| 947 | assert "0.scale" not in model.state_dict() |
| 948 | |
| 949 | |
| 950 | @pytest.fixture |
| 951 | def hpc_rope_norm(monkeypatch, default_vllm_config): |
| 952 | """Import HpcRopeNorm with the external ``hpc`` package stubbed out.""" |
| 953 | if "hpc" not in sys.modules: |
| 954 | stub = types.ModuleType("hpc") |
| 955 | stub.__spec__ = importlib.machinery.ModuleSpec("hpc", loader=None) |
| 956 | stub.QuantType = types.SimpleNamespace( # type: ignore[attr-defined] |
| 957 | QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR=types.SimpleNamespace(value=0) |
| 958 | ) |
| 959 | monkeypatch.setitem(sys.modules, "hpc", stub) |
| 960 | from vllm.model_executor.layers.hpc import rope_norm |
| 961 | |
| 962 | monkeypatch.setattr(rope_norm, "_hpc_rope_norm_instances", {}) |
| 963 | return rope_norm |
| 964 | |
| 965 | |
| 966 | def test_hpc_rope_norm_kernel_sees_refit_norm_weights(monkeypatch, hpc_rope_norm): |
| 967 | """The fused HPC kernel is handed the live QK-norm weights after a refit. |
| 968 | |
| 969 | Drives the production ``_forward_impl`` with a recording ``hpc`` stub. The |
| 970 | Q/K norm weights it receives must be the model's own float32 parameters, |
| 971 | so a layerwise reload that rewrites them in place (same storage) is what |
| 972 | the kernel sees. Previously the kernel read separate mirrors that no |
| 973 | reload path refreshed. |
| 974 | """ |
| 975 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 976 | |
| 977 | head_dim, num_heads, num_kv_heads, block_size = 128, 8, 1, 4 |
| 978 | layer = torch.nn.Module() |
| 979 | layer.q_norm = RMSNorm(head_dim, 1e-6, dtype=torch.float32) |
| 980 | layer.k_norm = RMSNorm(head_dim, 1e-6, dtype=torch.float32) |
| 981 | layer.hpc_rope_norm = hpc_rope_norm.HpcRopeNorm( |
| 982 | num_heads=num_heads, |
| 983 | num_kv_heads=num_kv_heads, |
| 984 | head_dim=head_dim, |
| 985 | cos_sin_cache=torch.ones(16, head_dim), |
| 986 | use_qk_norm=True, |
| 987 | fallback_qnorm=layer.q_norm, |
| 988 | fallback_knorm=layer.k_norm, |
| 989 | kv_cache_dtype="auto", |
| 990 | layer_name="hpc_test_layer", |
| 991 | ) |
| 992 | model = torch.nn.Sequential(layer) |
| 993 | rnorm = layer.hpc_rope_norm |
| 994 | |
| 995 | calls: list[dict] = [] |
| 996 | |
| 997 | def record(*args, **kwargs): |
| 998 | calls.append(kwargs) |
| 999 | |
| 1000 | monkeypatch.setattr(sys.modules["hpc"], "rope_norm_store_kv", record, raising=False) |
| 1001 | |
| 1002 | def kernel_norm_weights(): |
| 1003 | q_size, kv_size = num_heads * head_dim, num_kv_heads * head_dim |
| 1004 | qkv = torch.zeros(1, q_size + 2 * kv_size, dtype=torch.bfloat16) |
| 1005 | kv_cache = torch.zeros( |
| 1006 | 2, num_kv_heads, block_size, 2 * head_dim, dtype=torch.bfloat16 |
| 1007 | ) |
| 1008 | attn_layer = types.SimpleNamespace( |
| 1009 | _k_scale=torch.ones(1), _v_scale=torch.ones(1) |
| 1010 | ) |
| 1011 | attn_metadata = types.SimpleNamespace( |
| 1012 | num_actual_tokens=1, |
| 1013 | num_decodes=1, |
| 1014 | num_decode_tokens=1, |
| 1015 | num_prefills=0, |
| 1016 | num_prefill_tokens=0, |
| 1017 | max_query_len=1, |
| 1018 | decode_query_len=1, |
| 1019 | qo_indptr=None, |
| 1020 | qo_indptr_decode=None, |
| 1021 | slot_mapping=torch.tensor([4]), |
| 1022 | seq_lens=torch.tensor([1]), |
| 1023 | block_table_tensor=torch.tensor([[1]]), |
| 1024 | hpc_kv_written=False, |
| 1025 | ) |
| 1026 | output = torch.zeros(1, q_size, dtype=torch.bfloat16) |
| 1027 | rnorm._forward_impl(qkv, kv_cache, attn_metadata, attn_layer, output) |
| 1028 | return calls[-1]["q_norm_weight"], calls[-1]["k_norm_weight"] |
| 1029 | |
| 1030 | def loaded(value): |
| 1031 | return torch.full((head_dim,), value, dtype=torch.bfloat16) |
| 1032 | |
| 1033 | default_weight_loader(layer.q_norm.weight, loaded(0.5)) |
| 1034 | default_weight_loader(layer.k_norm.weight, loaded(0.25)) |
| 1035 | q, k = kernel_norm_weights() |
| 1036 | assert torch.equal(q, loaded(0.5).float()) |
| 1037 | assert torch.equal(k, loaded(0.25).float()) |
| 1038 | q_ptr, k_ptr = q.data_ptr(), k.data_ptr() |
| 1039 | |
| 1040 | record_metadata_for_reloading(model) |
| 1041 | initialize_layerwise_reload(model) |
| 1042 | layer.q_norm.weight.weight_loader(layer.q_norm.weight, loaded(2.0)) |
| 1043 | layer.k_norm.weight.weight_loader(layer.k_norm.weight, loaded(3.0)) |
| 1044 | finalize_layerwise_reload(model, model_config=None) |
| 1045 | |
| 1046 | q, k = kernel_norm_weights() |
| 1047 | assert q.dtype == k.dtype == torch.float32 |
| 1048 | assert torch.equal(q, loaded(2.0).float()) |
| 1049 | assert torch.equal(k, loaded(3.0).float()) |
| 1050 | assert (q.data_ptr(), k.data_ptr()) == (q_ptr, k_ptr) |
| 1051 | assert not hasattr(rnorm, "qnorm_weight") |
| 1052 | |
| 1053 | |
| 1054 | @pytest.mark.parametrize( |
| 1055 | "tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])] |
| 1056 | ) |
| 1057 | @pytest.mark.parametrize( |
| 1058 | "base_model,mul_model,add_model", |
| 1059 | [ |
| 1060 | pytest.param( |
| 1061 | "Qwen/Qwen3-0.6B", |
| 1062 | "inference-optimization/Qwen3-0.6B-debug-multiply", |
| 1063 | "inference-optimization/Qwen3-0.6B-debug-add", |
| 1064 | marks=[pytest.mark.slow_test], |
| 1065 | ), |
| 1066 | pytest.param( |
| 1067 | "inference-optimization/Qwen3-0.6B-FP8_BLOCK", |
| 1068 | "inference-optimization/Qwen3-0.6B-debug-multiply-FP8_BLOCK", |
| 1069 | "inference-optimization/Qwen3-0.6B-debug-add-FP8_BLOCK", |
| 1070 | marks=[pytest.mark.slow_test], |
| 1071 | ), |
| 1072 | pytest.param( |
| 1073 | "inference-optimization/Qwen3-0.6B-W4A16-G128", |
| 1074 | "inference-optimization/Qwen3-0.6B-debug-multiply-W4A16-G128", |
| 1075 | "inference-optimization/Qwen3-0.6B-debug-add-W4A16-G128", |
| 1076 | marks=[pytest.mark.slow_test], |
| 1077 | ), |
| 1078 | pytest.param( |
| 1079 | "inference-optimization/DeepSeek-V3-debug-empty", |
| 1080 | "inference-optimization/DeepSeek-V3-debug-multiply", |
| 1081 | "inference-optimization/DeepSeek-V3-debug-add", |
| 1082 | marks=[pytest.mark.slow_test], |
| 1083 | ), |
| 1084 | pytest.param( |
| 1085 | "inference-optimization/DeepSeek-V3-debug-empty-FP8_DYNAMIC", |
| 1086 | "inference-optimization/DeepSeek-V3-debug-multiply-FP8_DYNAMIC", |
| 1087 | "inference-optimization/DeepSeek-V3-debug-add-FP8_DYNAMIC", |
| 1088 | ), |
| 1089 | pytest.param( |
| 1090 | "inference-optimization/DeepSeek-V3-debug-empty-NVFP4A16", |
| 1091 | "inference-optimization/DeepSeek-V3-debug-multiply-NVFP4A16", |
| 1092 | "inference-optimization/DeepSeek-V3-debug-add-NVFP4A16", |
| 1093 | marks=[pytest.mark.slow_test], |
| 1094 | ), |
| 1095 | ], |
| 1096 | ) |
| 1097 | def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner): |
| 1098 | if current_platform.device_count() < tp_size: |
| 1099 | pytest.skip(reason="Not enough CUDA devices") |
| 1100 | |
| 1101 | if "FP8" in base_model and _fp8_reload_unsupported(): |
| 1102 | pytest.skip(reason="Requires FP8 support") |
| 1103 | |
| 1104 | with vllm_runner( |
| 1105 | model_name=base_model, |
| 1106 | tensor_parallel_size=tp_size, |
| 1107 | enable_expert_parallel=(tp_size > 1 and "DeepSeek" in base_model), |
| 1108 | enable_prefix_caching=False, |
| 1109 | max_model_len=16, |
| 1110 | max_num_seqs=1, |
| 1111 | ) as llm: |
| 1112 | llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_model}) |
| 1113 | mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0] |
| 1114 | add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0] |
| 1115 | assert mul_perp < add_perp |
| 1116 | |
| 1117 | llm.collective_rpc("reload_weights", kwargs={"weights_path": add_model}) |
| 1118 | mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0] |
| 1119 | add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0] |
| 1120 | assert add_perp < mul_perp |
| 1121 | |
| 1122 | |
| 1123 | def test_kv_scale_reload(vllm_runner): |
| 1124 | """Test reloading a checkpoint that contains k_scale/v_scale weights.""" |
| 1125 | if _fp8_reload_unsupported(): |
| 1126 | pytest.skip(reason="Requires FP8 support") |
| 1127 | |
| 1128 | model = "nm-testing/Llama-3.2-1B-Instruct-FP8-KV" |
| 1129 | |
| 1130 | # Load dummy weights, then reload real checkpoint |
| 1131 | with vllm_runner( |
| 1132 | model_name=model, |
| 1133 | load_format="dummy", |
| 1134 | enable_prefix_caching=False, |
| 1135 | max_model_len=16, |
| 1136 | max_num_seqs=1, |
| 1137 | ) as llm: |
| 1138 | llm.collective_rpc( |
| 1139 | "update_config", |
| 1140 | kwargs={"overrides": {"load_config": {"load_format": "auto"}}}, |
| 1141 | ) |
| 1142 | llm.collective_rpc("reload_weights", kwargs={"weights_path": model}) |
| 1143 | reloaded_perp = llm.generate_prompt_perplexity( |
| 1144 | ["The capital of France is the city of Paris"], |
| 1145 | mask=["The capital of France is"], |
| 1146 | )[0] |
| 1147 | |
| 1148 | assert reloaded_perp < 10 |
| 1149 | |
| 1150 | |
| 1151 | @pytest.mark.parametrize( |
| 1152 | "tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])] |
| 1153 | ) |
| 1154 | @pytest.mark.parametrize( |
| 1155 | "base_model,mul_model,add_model,quantization", |
| 1156 | [ |
| 1157 | pytest.param( |
| 1158 | "Qwen/Qwen3-0.6B", |
| 1159 | "inference-optimization/Qwen3-0.6B-debug-multiply", |
| 1160 | "inference-optimization/Qwen3-0.6B-debug-add", |
| 1161 | "fp8_per_tensor", |
| 1162 | ), |
| 1163 | pytest.param( |
| 1164 | "inference-optimization/DeepSeek-V3-debug-empty", |
| 1165 | "inference-optimization/DeepSeek-V3-debug-multiply", |
| 1166 | "inference-optimization/DeepSeek-V3-debug-add", |
| 1167 | "fp8_per_tensor", |
| 1168 | marks=[pytest.mark.slow_test], |
| 1169 | ), |
| 1170 | pytest.param( |
| 1171 | "Qwen/Qwen3-0.6B", |
| 1172 | "inference-optimization/Qwen3-0.6B-debug-multiply", |
| 1173 | "inference-optimization/Qwen3-0.6B-debug-add", |
| 1174 | "mxfp8", |
| 1175 | marks=[pytest.mark.slow_test], |
| 1176 | ), |
| 1177 | pytest.param( |
| 1178 | "inference-optimization/DeepSeek-V3-debug-empty", |
| 1179 | "inference-optimization/DeepSeek-V3-debug-multiply", |
| 1180 | "inference-optimization/DeepSeek-V3-debug-add", |
| 1181 | "mxfp8", |
| 1182 | marks=[ |
| 1183 | pytest.mark.slow_test, |
| 1184 | pytest.mark.xfail(reason="mxfp4 & mla is not supported yet"), |
| 1185 | ], |
| 1186 | ), |
| 1187 | ], |
| 1188 | ) |
| 1189 | def test_online_quantize_reload( |
| 1190 | base_model, mul_model, add_model, quantization, tp_size, vllm_runner |
| 1191 | ): |
| 1192 | if current_platform.device_count() < tp_size: |
| 1193 | pytest.skip(reason="Not enough GPU devices") |
| 1194 | |
| 1195 | if quantization == "fp8_per_tensor" and _fp8_reload_unsupported(): |
| 1196 | pytest.skip(reason="Requires FP8 support") |
| 1197 | |
| 1198 | with vllm_runner( |
| 1199 | model_name=base_model, |
| 1200 | quantization=quantization, |
| 1201 | tensor_parallel_size=tp_size, |
| 1202 | enable_expert_parallel=(tp_size > 1 and "DeepSeek" in base_model), |
| 1203 | enable_prefix_caching=False, |
| 1204 | max_model_len=16, |
| 1205 | max_num_seqs=1, |
| 1206 | ) as llm: |
| 1207 | llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_model}) |
| 1208 | mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0] |
| 1209 | add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0] |
| 1210 | assert mul_perp < add_perp |
| 1211 | |
| 1212 | llm.collective_rpc("reload_weights", kwargs={"weights_path": add_model}) |
| 1213 | mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0] |
| 1214 | add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0] |
| 1215 | assert add_perp < mul_perp |