@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 648edcf729

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 21:00 UTC

[QeRL] Compose online quantization with quantized reloading (#38032)

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Kyle Sayers committed Mar 27, 2026 at 16:22 UTC 648edcf7291147fd3e76e9c525e3c5256e4725ee

10 files changed +184 -260

tests/model_executor/model_loader/test_reload.py

+57

@@ -148,3 +148,60 @@ def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner):

148 mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]

149 add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0]

150 assert add_perp < mul_perp

151 +

152 +

153 +@pytest.mark.parametrize("tp_size", [2])

154 +@pytest.mark.parametrize(

155 + "base_model,mul_model,add_model,quantization",

156 + [

157 + (

158 + "Qwen/Qwen3-0.6B",

159 + "inference-optimization/Qwen3-0.6B-debug-multiply",

160 + "inference-optimization/Qwen3-0.6B-debug-add",

161 + "fp8",

162 + ),

163 + (

164 + "inference-optimization/DeepSeek-V3-debug-empty",

165 + "inference-optimization/DeepSeek-V3-debug-multiply",

166 + "inference-optimization/DeepSeek-V3-debug-add",

167 + "fp8",

168 + ),

169 + (

170 + "Qwen/Qwen3-0.6B",

171 + "inference-optimization/Qwen3-0.6B-debug-multiply",

172 + "inference-optimization/Qwen3-0.6B-debug-add",

173 + "mxfp8",

174 + ),

175 + # ( TODO: support mxfp4 & mla

176 + # "inference-optimization/DeepSeek-V3-debug-empty",

177 + # "inference-optimization/DeepSeek-V3-debug-multiply",

178 + # "inference-optimization/DeepSeek-V3-debug-add",

179 + # "mxfp8",

180 + # ),

181 + ],

182 +)

183 +def test_online_quantize_reload(

184 + base_model, mul_model, add_model, quantization, tp_size, vllm_runner

185 +):

186 + if cuda_device_count_stateless() < tp_size:

187 + pytest.skip(reason="Not enough CUDA devices")

188 +

189 + if quantization == "fp8" and not current_platform.supports_fp8():

190 + pytest.skip(reason="Requires FP8 support")

191 +

192 + with vllm_runner(

193 + model_name=base_model,

194 + quantization=quantization,

195 + tensor_parallel_size=tp_size,

196 + enable_expert_parallel=(tp_size > 1 and "DeepSeek" in base_model),

197 + enable_prefix_caching=False,

198 + ) as llm:

199 + llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_model})

200 + mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]

201 + add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0]

202 + assert mul_perp < add_perp

203 +

204 + llm.collective_rpc("reload_weights", kwargs={"weights_path": add_model})

205 + mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]

206 + add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0]

207 + assert add_perp < mul_perp

vllm/model_executor/layers/quantization/fp8.py

+27 -202

@@ -73,7 +73,9 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

73 cutlass_fp8_supported,

74 normalize_e4m3fn_to_e4m3fnuz,

75 )

76 -from vllm.model_executor.model_loader.weight_utils import initialize_single_dummy_weight

76 +from vllm.model_executor.model_loader.reload.layerwise import (

77 + initialize_online_processing,

78 +)

79 from vllm.model_executor.parameter import (

80 BlockQuantScaleParameter,

81 ModelWeightParameter,

@@ -496,8 +498,8 @@ class Fp8LinearMethod(LinearMethodBase):

498

499

500 class Fp8OnlineLinearMethod(Fp8LinearMethod):

499 - """Online version of Fp8LinearMethod, loads the fp16/bf16 checkpoint

500 - and quantized the weights during loading."""

501 + """Online version of Fp8LinearMethod which loads a full precision checkpoint

502 + and quantizes weights during loading."""

503

504 uses_meta_device: bool = True

505

@@ -519,84 +521,25 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod):

521 layer.orig_dtype = params_dtype

522 layer.weight_block_size = None

523

522 - # WEIGHT

523 - def patched_weight_loader(param, loaded_weight, *args, **kwargs):

524 - # track how many elements we have updated

525 - if not hasattr(layer, "_loaded_numel"):

526 - layer._loaded_numel = 0

527 -

528 - # when the first `loaded_weight` is about to be

529 - # loaded to `param`, materialize `param` just-in-time

530 - weight = ModelWeightParameter(

531 - data=torch.empty_like(layer.weight, device=layer._load_device),

532 - input_dim=1,

533 - output_dim=0,

534 - weight_loader=patched_weight_loader,

535 - )

536 - _copy_missing_attrs(layer.weight, weight)

537 - layer.register_parameter("weight", weight)

538 - del layer._load_device

539 -

540 - # refresh the reference to `param` to reflect just-in-time

541 - # materialization

542 - param = layer.weight

543 -

544 - # load the current weight chunk

545 - copy_numel_counter = CopyNumelCounter()

546 - with copy_numel_counter:

547 - res = weight_loader(param, loaded_weight, *args, **kwargs) # type: ignore[misc]

548 - layer._loaded_numel += copy_numel_counter.copied_numel

549 -

550 - # if we have loaded all of the elements, call

551 - # process_weights_after_loading

552 - target_loaded_numel = layer.weight.numel()

553 - if layer._loaded_numel == target_loaded_numel:

554 - self.process_weights_after_loading(layer)

555 -

556 - # Prevent the usual `process_weights_after_loading` call from doing

557 - # anything

558 - layer._already_called_process_weights_after_loading = True

559 -

560 - # Note that we keep `layer._loaded_numel` around just in case

561 - # there is logic added to vllm in the future which calls a

562 - # weight loader twice - we do not want to re-initialize in

563 - # that case.

564 -

565 - return res

566 -

524 weight = ModelWeightParameter(

525 data=torch.empty(

526 output_size_per_partition,

527 input_size_per_partition,

571 - # materialized just-in-time in `patched_weight_loader`

572 - device="meta",

528 + device="meta", # materialized and processed during loading

529 dtype=params_dtype,

530 ),

531 input_dim=1,

532 output_dim=0,

577 - weight_loader=patched_weight_loader,

533 + weight_loader=weight_loader,

534 )

579 - # stash the correct device for `patched_weight_loader`

580 - layer._load_device = torch.get_default_device()

535 layer.register_parameter("weight", weight)

536

537 + initialize_online_processing(layer)

538 +

539 def process_weights_after_loading(self, layer: Module) -> None:

540 if getattr(layer, "_already_called_process_weights_after_loading", False):

541 return

542

587 - # deferred initialization of randomly initialized weights for the

588 - # `--load_format dummy` feature

589 - if layer.weight.device == torch.device("meta"):

590 - weight = ModelWeightParameter(

591 - data=torch.empty_like(layer.weight, device=layer._load_device),

592 - input_dim=1,

593 - output_dim=0,

594 - weight_loader=layer.weight.weight_loader,

595 - )

596 - _copy_missing_attrs(layer.weight, weight)

597 - layer.register_parameter("weight", weight)

598 - initialize_single_dummy_weight(layer.weight)

599 -

543 # TODO(future): support block_quant in online quant path

544 assert not self.block_quant

545

@@ -845,9 +788,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):

788 )

789

790 def process_weights_after_loading(self, layer: Module) -> None:

848 - if getattr(layer, "_already_called_process_weights_after_loading", False):

849 - return

850 -

791 # Allow for accessing weights and scales in standard way.

792 w13 = layer.w13_weight

793 w2 = layer.w2_weight

@@ -892,9 +832,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):

832 layer, w13, w2, w13_scale, w2_scale, w13_input_scale, w2_input_scale

833 )

834

895 - # Prevent duplicate processing (e.g., during weight reload)

896 - layer._already_called_process_weights_after_loading = True

897 -

835 def maybe_make_prepare_finalize(

836 self,

837 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

@@ -1013,86 +950,12 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):

950 layer.orig_dtype = params_dtype

951 layer.weight_block_size = None

952

1016 - # We are doing online quantization, patch the weight loaded

1017 - # to call `process_weights_after_loading` in a streaming fashion

1018 - # as soon as the last weight chunk is loaded.

1019 - weight_loader = extra_weight_attrs["weight_loader"]

1020 - # create a new holder to prevent modifying behavior of any other

1021 - # objects which might depend on the old one

1022 - new_extra_weight_attrs = extra_weight_attrs

1023 -

1024 - def patched_weight_loader(param, loaded_weight, *args, **kwargs):

1025 - # add a counter to track how many elements we have updated

1026 - if not hasattr(layer, "_loaded_numel"):

1027 - layer._loaded_numel = 0

1028 -

1029 - # save the ids of original w13 and w2 so that we can

1030 - # distinguish which one `param` should map to further

1031 - # down in this file

1032 - layer._w13_weight_orig_id = id(layer.w13_weight)

1033 - layer._w2_weight_orig_id = id(layer.w2_weight)

1034 -

1035 - # when the first `loaded_weight` is about to be

1036 - # loaded to `param`, materialize `param` just-in-time

1037 -

1038 - w13_weight = torch.nn.Parameter(

1039 - torch.empty_like(layer.w13_weight, device=layer._load_device),

1040 - requires_grad=False,

1041 - )

1042 - set_weight_attrs(w13_weight, extra_weight_attrs)

1043 - _copy_missing_attrs(layer.w13_weight, w13_weight)

1044 - layer.register_parameter("w13_weight", w13_weight)

1045 -

1046 - w2_weight = torch.nn.Parameter(

1047 - torch.empty_like(layer.w2_weight, device=layer._load_device),

1048 - requires_grad=False,

1049 - )

1050 - set_weight_attrs(w2_weight, extra_weight_attrs)

1051 - _copy_missing_attrs(layer.w2_weight, w2_weight)

1052 - layer.register_parameter("w2_weight", w2_weight)

1053 - del layer._load_device

1054 -

1055 - # refresh the reference to `param` to reflect just-in-time

1056 - # materialization

1057 - if id(param) == layer._w13_weight_orig_id:

1058 - param = layer.w13_weight

1059 - elif id(param) == layer._w2_weight_orig_id:

1060 - param = layer.w2_weight

1061 -

1062 - # load the current weight chunk

1063 - copy_numel_counter = CopyNumelCounter()

1064 - with copy_numel_counter:

1065 - res = weight_loader(param, loaded_weight, *args, **kwargs) # type: ignore[misc]

1066 - layer._loaded_numel += copy_numel_counter.copied_numel

1067 -

1068 - # if we have loaded all of the elements, call

1069 - # process_weights_after_loading

1070 - target_loaded_numel = layer.w13_weight.numel() + layer.w2_weight.numel()

1071 - if layer._loaded_numel == target_loaded_numel:

1072 - self.process_weights_after_loading(layer)

1073 -

1074 - # Prevent the usual `process_weights_after_loading` call

1075 - # from doing anything

1076 - layer._already_called_process_weights_after_loading = True

1077 -

1078 - # Note that we keep `layer._loaded_numel`,

1079 - # `layer._w13_weight_orig_id` and `layer._w2_weight_orig_id`

1080 - # around because if EP is on, weight loaders for non-local

1081 - # experts will run but not actually copy any elements, and we

1082 - # need to not re-initialize in that case.

1083 -

1084 - return res

1085 -

1086 - new_extra_weight_attrs["weight_loader"] = patched_weight_loader

1087 - extra_weight_attrs = new_extra_weight_attrs

1088 -

953 # WEIGHTS

954 w13_weight = torch.nn.Parameter(

955 torch.empty(

956 num_experts,

957 2 * intermediate_size_per_partition,

958 hidden_size,

1095 - # materialized just-in-time in `patched_weight_loader`

959 device="meta",

960 dtype=params_dtype,

961 ),

@@ -1106,91 +969,53 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):

969 num_experts,

970 hidden_size,

971 intermediate_size_per_partition,

1109 - # materialized just-in-time in `patched_weight_loader`

1110 - device="meta",

972 + device="meta", # materialized and processed during loading

973 dtype=params_dtype,

974 ),

975 requires_grad=False,

976 )

977 layer.register_parameter("w2_weight", w2_weight)

978 set_weight_attrs(w2_weight, extra_weight_attrs)

1117 - # stash the correct device for `patched_weight_loader`

1118 - layer._load_device = torch.get_default_device()

979

980 # BIASES (for models like GPT-OSS that have biased MoE)

981 if self.moe.has_bias:

1122 - # Use the original weight_loader (not patched) for biases

1123 - orig_extra_weight_attrs = dict(extra_weight_attrs)

1124 - orig_extra_weight_attrs["weight_loader"] = weight_loader

982 w13_bias = torch.nn.Parameter(

983 torch.zeros(

984 num_experts,

985 2 * intermediate_size_per_partition,

986 + device="meta", # materialized and processed during loading

987 dtype=layer.orig_dtype,

988 ),

989 requires_grad=False,

990 )

991 layer.register_parameter("w13_bias", w13_bias)

1134 - set_weight_attrs(w13_bias, orig_extra_weight_attrs)

992 + set_weight_attrs(w13_bias, extra_weight_attrs)

993 +

994 w2_bias = torch.nn.Parameter(

1136 - torch.zeros(num_experts, hidden_size, dtype=layer.orig_dtype),

995 + torch.zeros(

996 + num_experts,

997 + hidden_size,

998 + device="meta", # materialized and processed during loading

999 + dtype=layer.orig_dtype,

1000 + ),

1001 requires_grad=False,

1002 )

1003 layer.register_parameter("w2_bias", w2_bias)

1140 - set_weight_attrs(w2_bias, orig_extra_weight_attrs)

1141 -

1142 - # WEIGHT_SCALES

1143 - # Allocate 2 scales for w1 and w3 respectively.

1144 - # They will be combined to a single scale after weight loading.

1145 - w13_weight_scale = torch.nn.Parameter(

1146 - torch.ones(num_experts, dtype=torch.float32), requires_grad=False

1147 - )

1148 - w2_weight_scale = torch.nn.Parameter(

1149 - torch.ones(num_experts, dtype=torch.float32), requires_grad=False

1150 - )

1151 - layer.register_parameter("w13_weight_scale", w13_weight_scale)

1152 - layer.register_parameter("w2_weight_scale", w2_weight_scale)

1153 - set_weight_attrs(w13_weight_scale, extra_weight_attrs)

1154 - set_weight_attrs(w2_weight_scale, extra_weight_attrs)

1004 + set_weight_attrs(w2_bias, extra_weight_attrs)

1005

1156 - layer.w13_input_scale = None

1157 - layer.w2_input_scale = None

1006 + initialize_online_processing(layer)

1007

1008 def process_weights_after_loading(self, layer: Module) -> None:

1009 if getattr(layer, "_already_called_process_weights_after_loading", False):

1010 return

1011

1163 - # deferred initialization of randomly initialized weights for the

1164 - # `--load_format dummy` feature

1165 - if layer.w13_weight.device == torch.device("meta"):

1166 - w13_weight = torch.nn.Parameter(

1167 - torch.empty_like(layer.w13_weight, device=layer._load_device),

1168 - requires_grad=False,

1169 - )

1170 - set_weight_attrs(

1171 - w13_weight, {"weight_loader": layer.w13_weight.weight_loader}

1172 - )

1173 - _copy_missing_attrs(layer.w13_weight, w13_weight)

1174 - layer.register_parameter("w13_weight", w13_weight)

1175 - initialize_single_dummy_weight(layer.w13_weight)

1176 - if layer.w2_weight.device == torch.device("meta"):

1177 - w2_weight = torch.nn.Parameter(

1178 - torch.empty_like(layer.w2_weight, device=layer._load_device),

1179 - requires_grad=False,

1180 - )

1181 - set_weight_attrs(

1182 - w2_weight, {"weight_loader": layer.w2_weight.weight_loader}

1183 - )

1184 - _copy_missing_attrs(layer.w2_weight, w2_weight)

1185 - layer.register_parameter("w2_weight", w2_weight)

1186 - initialize_single_dummy_weight(layer.w2_weight)

1187 -

1188 - # If checkpoint is fp16, quantize in place.

1012 fp8_dtype = current_platform.fp8_dtype()

1013 w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)

1014 w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)

1192 - w13_scale = layer.w13_weight_scale

1193 - w2_scale = layer.w2_weight_scale

1015 + w13_scale = torch.ones(layer.num_experts, dtype=torch.float32)

1016 + w2_scale = torch.ones(layer.num_experts, dtype=torch.float32)

1017 + layer.w13_input_scale = None

1018 + layer.w2_input_scale = None

1019

1020 for expert in range(layer.local_num_experts):

1021 w13[expert, :, :], w13_scale[expert] = ops.scaled_fp8_quant(

@@ -1207,8 +1032,8 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):

1032 w2,

1033 w13_scale,

1034 w2_scale,

1210 - layer.w13_input_scale,

1211 - layer.w2_input_scale,

1035 + w13_input_scale=layer.w13_input_scale,

1036 + w2_input_scale=layer.w2_input_scale,

1037 )

1038

1039 # Prevent duplicate processing (e.g., during weight reload)

vllm/model_executor/layers/quantization/mxfp8.py

@@ -337,6 +337,8 @@ class Mxfp8OnlineMoEMethod(Fp8OnlineMoEMethod):

337 w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)

338 w13_scale = layer.w13_weight_scale

339 w2_scale = layer.w2_weight_scale

340 + layer.w13_input_scale = None

341 + layer.w2_input_scale = None

342

343 w13, w13_scale = self._quantize_mxfp8_moe_weight(layer.w13_weight)

344 w2, w2_scale = self._quantize_mxfp8_moe_weight(layer.w2_weight)

vllm/model_executor/model_loader/base_loader.py

+19 -7

@@ -9,6 +9,7 @@ import vllm.envs as envs

9 from vllm.config import ModelConfig, VllmConfig

10 from vllm.config.load import LoadConfig

11 from vllm.logger import init_logger

12 +from vllm.model_executor.model_loader.reload import finalize_layerwise_processing

13 from vllm.model_executor.model_loader.utils import (

14 initialize_model,

15 process_weights_after_loading,

@@ -49,16 +50,13 @@ class BaseModelLoader(ABC):

50 device_config.device if load_config.device is None else load_config.device

51 )

52 target_device = torch.device(load_device)

52 - with set_default_torch_dtype(model_config.dtype):

53 - with target_device:

54 - model = initialize_model(

55 - vllm_config=vllm_config, model_config=model_config, prefix=prefix

56 - )

57 -

53 + with set_default_torch_dtype(model_config.dtype), target_device:

54 + model = initialize_model(

55 + vllm_config=vllm_config, model_config=model_config, prefix=prefix

56 + )

57 log_model_inspection(model)

59 logger.debug("Loading weights on %s ...", load_device)

61 - # Quantization does not happen in `load_weights` but after it

60 self.load_weights(model, model_config)

62 # Log peak GPU memory after loading weights. This is needed

@@ -71,6 +69,11 @@ class BaseModelLoader(ABC):

69 scope="local",

70 )

72 + # Process weights into kernel format. Note that when using online

73 + # quantization, weights are (typically) quantized as they are loaded.

74 + if _has_online_quant(model):

75 + finalize_layerwise_processing(model, model_config)

76 +

77 process_weights_after_loading(model, model_config, target_device)

79 return model.eval()

@@ -84,3 +87,12 @@ def log_model_inspection(model: nn.Module) -> None:

87 from vllm.model_inspection import format_model_inspection

89 logger.info("vLLM model structure:\n%s", format_model_inspection(model))

90 +

91 +

92 +def _has_online_quant(model: nn.Module):

93 + for module in model.modules():

94 + quant_method = getattr(module, "quant_method", None)

95 + if getattr(quant_method, "uses_meta_device", False):

96 + return True

97 +

98 + return False

vllm/model_executor/model_loader/dummy_loader.py

@@ -1,10 +1,13 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import torch

4 import torch.nn as nn

6 from vllm.config import ModelConfig

7 from vllm.config.load import LoadConfig

8 from vllm.model_executor.model_loader.base_loader import BaseModelLoader

9 +from vllm.model_executor.model_loader.reload.meta import materialize_meta_tensor

10 +from vllm.model_executor.model_loader.reload.utils import get_layer_tensors

11 from vllm.model_executor.model_loader.weight_utils import initialize_dummy_weights

@@ -23,6 +26,12 @@ class DummyModelLoader(BaseModelLoader):

26 pass # Nothing to download

28 def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:

29 + # materialize meta tensors as part of online quantization lifecycle

30 + for layer in model.modules():

31 + for name, param in get_layer_tensors(layer).items():

32 + if param.device == torch.device("meta"):

33 + setattr(layer, name, materialize_meta_tensor(param))

34 +

35 # NOTE(woosuk): For accurate performance evaluation, we assign

36 # random values to the weights.

37 initialize_dummy_weights(model, model_config)

vllm/model_executor/model_loader/reload/__init__.py

@@ -21,12 +21,14 @@ Limitations:

21 __all__ = [

22 "record_metadata_for_reloading",

23 "initialize_layerwise_reload",

24 + "finalize_layerwise_processing",

25 "finalize_layerwise_reload",

26 "set_torchao_reload_attrs",

27 "support_quantized_model_reload_from_hp_weights",

28 ]

30 from .layerwise import (

31 + finalize_layerwise_processing,

32 finalize_layerwise_reload,

33 initialize_layerwise_reload,

34 record_metadata_for_reloading,

vllm/model_executor/model_loader/reload/layerwise.py

+61 -29

@@ -28,6 +28,7 @@ __all__ = [

28 "get_layerwise_info",

29 "record_metadata_for_reloading",

30 "initialize_layerwise_reload",

31 + "finalize_layerwise_processing",

32 "finalize_layerwise_reload",

33 ]

@@ -89,7 +90,7 @@ def initialize_layerwise_reload(model: torch.nn.Module):

90 info = get_layerwise_info(layer)

92 # Skip if the layer has already been initialized

92 - if info.can_process():

93 + if info.can_load():

94 continue

96 # Save current tensors for later copying

@@ -98,15 +99,21 @@ def initialize_layerwise_reload(model: torch.nn.Module):

99 # Restore layer parameters/buffers onto meta device

100 restore_layer_on_meta(layer, info)

101

101 - # Track loading progress to determine when to process/copy

102 - info.load_numel = 0

103 - info.load_numel_total = get_layer_size(layer)

102 + initialize_online_processing(layer)

103

105 - # Wrap each parameter's weight loader

106 - # Note that nested wrapping will occur for shared tensors

107 - for name, tensor in get_layer_tensors(layer).items():

108 - if _get_weight_loader(tensor).__name__ != "online_process_loader":

109 - tensor.weight_loader = make_online_process_loader(layer, name)

104 +

105 +def initialize_online_processing(layer: torch.nn.Module):

106 + info = get_layerwise_info(layer)

107 +

108 + # Track loading progress to determine when to process/copy

109 + info.load_numel = 0

110 + info.load_numel_total = get_layer_size(layer)

111 +

112 + # Wrap each parameter's weight loader

113 + # Note that nested wrapping will occur for shared tensors

114 + for name, tensor in get_layer_tensors(layer).items():

115 + if _get_weight_loader(tensor).__name__ != "online_process_loader":

116 + tensor.weight_loader = make_online_process_loader(layer, name)

117

118

119 def make_online_process_loader(layer: torch.nn.Module, param_name: str) -> Callable:

@@ -118,7 +125,7 @@ def make_online_process_loader(layer: torch.nn.Module, param_name: str) -> Calla

125

126 @wraps(original_loader, assigned=("__doc__", "__annotations__"))

127 def online_process_loader(*args, **kwargs):

121 - if not info.can_process():

128 + if not info.can_load():

129 # Unfortunately, some qconfigs are set up to load the same weight

130 # multiple times. For example, CT_WNA16 loads `weight_shape` for

131 # each of the qkv partitions. This results in layers loading extra

@@ -140,7 +147,7 @@ def make_online_process_loader(layer: torch.nn.Module, param_name: str) -> Calla

147 bound_args = loader_signature.bind(*args, **kwargs)

148 bound_args.apply_defaults()

149

143 - # Cache loaded weights, track loading progress

150 + # Buffer loaded weights, track loading progress

151 info.loaded_weights.append((param_name, bound_args))

152 num_loaded, ret = get_numel_loaded(original_loader, bound_args)

153 info.load_numel += num_loaded

@@ -163,19 +170,26 @@ def make_online_process_loader(layer: torch.nn.Module, param_name: str) -> Calla

170 return online_process_loader

171

172

166 -def finalize_layerwise_reload(model: torch.nn.Module, model_config: ModelConfig):

173 +def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelConfig):

174 """

168 - Remove the outermost layer of weight loading wrappers.

175 + Apply processing to any layers which were not layerwise processed during loading.

176 + This includes attention layers and layers which have weight elements which are not

177 + loaded (due to padding).

178

179 This function should be applied after `initialize_layerwise_reload` is applied

180 unwrap the layerwise weight loaders.

181

173 - Also processes Attention/MLA layers, which must be processed after all other layers

182 + :param model: model to finalize processing for

183 + :param model_config: config needed for applying processing to attention layers

184 """

175 - model._do_torchao_reload = model._original_do_torchao_reload

185 + if hasattr(model, "_original_do_torchao_reload"):

186 + model._do_torchao_reload = model._original_do_torchao_reload

187

188 for layer in model.modules():

189 info = get_layerwise_info(layer)

190 + if not info.can_load():

191 + info.reset()

192 + continue

193

194 # Attention/MLA layers are processed after all other layers

195 if isinstance(layer, (Attention, MLAAttention)):

@@ -184,17 +198,29 @@ def finalize_layerwise_reload(model: torch.nn.Module, model_config: ModelConfig)

198 "Layerwise reloading of Q/K/V scale weights is not implemented yet"

199 )

200

201 + elif info.kernel_tensors is None:

202 + raise NotImplementedError(

203 + "Layerwise loading of Q/K/V scale weights is not implemented yet"

204 + )

205 +

206 else:

207 _place_kernel_tensors(layer, info)

208 layer.process_weights_after_loading(model_config.dtype)

209

191 - # No weights were loaded, place kernel tensors back

192 - elif info.can_process() and info.load_numel <= 0:

193 - _place_kernel_tensors(layer, info)

210 + # No weights were loaded

211 + elif info.load_numel <= 0:

212 + # first load but received no weights. This happens on dummy load

213 + if info.kernel_tensors is None:

214 + materialize_layer(layer)

215 +

216 + # reloading: place kernel tensors back as a fallback

217 + else:

218 + logger.warning("%s: Failed to load weights", layer.__class__.__name__)

219 + _place_kernel_tensors(layer, info)

220

221 # Process non-attention layers which did not load all elements. This can happen

222 # if the created weight has extra padding elements which are not loaded

197 - # Having too many of these delayed layers can lead to execess memory usage

223 + # Having too many of these delayed layers can lead to excess memory usage

224 # see Limitations(4)

225 elif info.load_numel > 0 and info.load_numel < info.load_numel_total: # type: ignore[operator]

226 logger.debug("%s: Delayed processing", layer.__class__.__name__)

@@ -203,20 +229,24 @@ def finalize_layerwise_reload(model: torch.nn.Module, model_config: ModelConfig)

229 info.reset()

230

231

232 +def finalize_layerwise_reload(*args, **kwargs):

233 + finalize_layerwise_processing(*args, **kwargs)

234 +

235 +

236 def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo):

237 """

208 - Finalize layer loading after all weights have been cached.

238 + Finalize layer loading after all weights have been buffered.

239

240 This function:

241 1. Materializes the layer onto the target device

212 - 2. Loads all cached weights

242 + 2. Loads all buffered weights

243 3. Runs quantization processing if applicable

244 4. Copies processed values back to original tensor storage

245 """

246 # Materialize layer tensors onto device

247 materialize_layer(layer)

248

219 - # Reset FP8 online quantization flag so process_weights_after_loading

249 + # Reset online quantization flag so process_weights_after_loading

250 # will run again during reload

251 if hasattr(layer, "_already_called_process_weights_after_loading"):

252 delattr(layer, "_already_called_process_weights_after_loading")

@@ -225,7 +255,7 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo):

255 for param in get_layer_tensors(layer).values():

256 param.weight_loader = _get_original_loader(param)

257

228 - # Load all cached weights into materialized layer (using original loaders)

258 + # Load all buffered weights into materialized layer (using original loaders)

259 for name, args in info.loaded_weights:

260 param = getattr(layer, name)

261 args.arguments["param"] = param

@@ -239,13 +269,14 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo):

269

270 # Copy processed values into original tensor storage (preserves cudagraph refs)

271 # this code is a no-op if not reloading (because kernel tensors is empty)

242 - parameters, buffers = info.kernel_tensors

243 - for name, param in parameters.items():

244 - param.data.copy_(getattr(layer, name))

245 - for name, buffer in buffers.items():

246 - buffer.data.copy_(getattr(layer, name))

272 + if info.kernel_tensors is not None:

273 + parameters, buffers = info.kernel_tensors

274 + for name, param in parameters.items():

275 + param.data.copy_(getattr(layer, name))

276 + for name, buffer in buffers.items():

277 + buffer.data.copy_(getattr(layer, name))

278

248 - _place_kernel_tensors(layer, info)

279 + _place_kernel_tensors(layer, info)

280

281 info.reset()

282 logger.debug("%s: Processed", layer.__class__.__name__)

@@ -268,6 +299,7 @@ def _place_kernel_tensors(layer: torch.nn.Module, info: LayerReloadingInfo):

299 for name in get_layer_tensors(layer):

300 delattr(layer, name)

301

302 + assert info.kernel_tensors is not None

303 parameters, buffers = info.kernel_tensors

304 for name, param in parameters.items():

305 layer.register_parameter(name, param)

vllm/model_executor/model_loader/reload/meta.py

+3 -4

@@ -104,7 +104,7 @@ def materialize_layer(layer: torch.nn.Module) -> None:

104 setattr(layer, name, materialize_meta_tensor(tensor))

105

106

107 -class MetaCopyCounter(TorchDispatchMode):

107 +class CopyCounter(TorchDispatchMode):

108 """

109 Tracks total number of elements modified with `copy_`.

110

@@ -122,7 +122,7 @@ class MetaCopyCounter(TorchDispatchMode):

122 if kwargs is None:

123 kwargs = {}

124

125 - if func is torch.ops.aten.copy_.default and args[0].device.type == "meta":

125 + if func is torch.ops.aten.copy_.default:

126 assert args[0].numel() == args[1].numel()

127 self.copied_numel += args[0].numel()

128

@@ -140,7 +140,6 @@ def get_numel_loaded(

140 :return: number of elements loaded by the weight loader, the return value of the

141 weight loader

142 """

143 - assert args.arguments["param"].device.type == "meta"

144 - with MetaCopyCounter() as counter:

143 + with CopyCounter() as counter:

144 return_value = weight_loader(*args.args, **args.kwargs)

145 return counter.copied_numel, return_value

vllm/model_executor/model_loader/reload/types.py

+3 -3

@@ -16,8 +16,8 @@ class LayerReloadingInfo:

16 # model format (meta), populated by `record_metadata_for_reloading`

17 restore_metadata: LayerTensors = field(default_factory=lambda: ({}, {}))

19 - # kernel format (device)

20 - kernel_tensors: LayerTensors = field(default_factory=lambda: ({}, {}))

19 + # kernel format (device), used to copy into when reloading only

20 + kernel_tensors: LayerTensors | None = None

22 # track how many restored elements are ready for loading

23 load_numel: int = 0

@@ -29,5 +29,5 @@ class LayerReloadingInfo:

29 def reset(self):

30 self.__init__(restore_metadata=self.restore_metadata) # type: ignore[misc]

32 - def can_process(self) -> bool:

32 + def can_load(self) -> bool:

33 return self.load_numel_total is not None

vllm/model_executor/model_loader/weight_utils.py

+1 -15

@@ -1323,25 +1323,11 @@ def initialize_dummy_weights(

1323 is fixed, the random values generated by this function only depends on

1324 the parameter's number of elements and its data type.

1325 """

1326 -

1327 - # Check if any module uses online quantization with meta device weights.

1328 - # If so, we'll skip initializing params on meta device since they'll be

1329 - # handled in `process_weights_after_loading`.

1330 - def uses_meta_device(module: torch.nn.Module) -> bool:

1331 - quant_method = getattr(module, "quant_method", None)

1332 - return getattr(quant_method, "uses_meta_device", False)

1333 -

1334 - has_online_quant = any(uses_meta_device(m) for m in model.modules())

1335 -

1326 for param in model.state_dict().values():

1337 - if has_online_quant and param.device == torch.device("meta"):

1338 - # For online quantization, weights are created on meta device and

1339 - # dummy weight init will happen in `process_weights_after_loading`.

1340 - continue

1341 -

1327 initialize_single_dummy_weight(param, low, high, seed)

1328

1329

1330 +@torch.no_grad()

1331 def initialize_single_dummy_weight(

1332 param: torch.Tensor,

1333 low: float = -1e-3,