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 9, 02:30 UTC
fix memory for online fp8 quantization with streaming weight load (#31914)
Signed-off-by: vasiliy <vasiliy@fb.com>
Vasiliy Kuznetsov committed
Feb 2, 2026 at 14:17 UTC
0130223bd9900710a0d93e46a4255ec5d1a077a8
5 files changed
+282
-37
tests/quantization/test_fp8.py
+96
@@ -5,7 +5,10 @@
5
Run `pytest tests/quantization/test_fp8.py --forked`.
6
"""
7
8
+import logging
9
+
10
import pytest
11
+import regex as re
12
import torch
13
14
from tests.quantization.utils import is_quant_method_supported
@@ -195,6 +198,99 @@ def test_online_quantization(
198
print(outputs[0][1])
199
200
201
+@pytest.mark.skipif(
202
+ not is_quant_method_supported("fp8"),
203
+ reason="FP8 is not supported on this GPU type.",
204
+)
205
+def test_online_quant_peak_mem(
206
+ vllm_runner,
207
+ caplog_mp_spawn,
208
+ monkeypatch,
209
+) -> None:
210
+ # Note: `allenai/OLMoE-1B-7B-0125-Instruct` was selected because:
211
+ # 1. it covers both Linear and MoE paths
212
+ # 2. it is already used by other tests in CI, so adding it here
213
+ # does not increase disk space for CI runners
214
+ # I really wanted to use `ibm-granite/granite-3.0-1b-a400m-base`
215
+ # which I think is the smallest MoE model in vLLM (2.5 GiB bf16,
216
+ # 1.3 GiB fp8), but could not as adding one more model makes CI
217
+ # run out of disk space.
218
+ model_name = "allenai/OLMoE-1B-7B-0125-Instruct"
219
+
220
+ # Force spawn to ensure caplog_mp_spawn works consistently
221
+ # (it relies on VLLM_LOGGING_CONFIG_PATH which spawn reads but fork ignores)
222
+ monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
223
+
224
+ with (
225
+ caplog_mp_spawn(logging.DEBUG) as log_holder,
226
+ vllm_runner(
227
+ model_name,
228
+ quantization="fp8",
229
+ enforce_eager=True,
230
+ ) as llm,
231
+ ):
232
+ outputs = llm.generate_greedy(["The future of AI is"], max_tokens=4)
233
+ print(outputs[0][1])
234
+
235
+ log_text = log_holder.text
236
+
237
+ # Parse memory usage from captured logs
238
+ model_memory_gib = None
239
+ peak_memory_gib = None
240
+ for line in log_text.splitlines():
241
+ if model_memory_gib is None:
242
+ match = re.search(r"Model loading took ([\d.]+) GiB memory", line)
243
+ if match:
244
+ model_memory_gib = float(match.group(1))
245
+ if peak_memory_gib is None:
246
+ match = re.search(
247
+ r"Peak GPU memory after loading weights: ([\d.]+) GiB", line
248
+ )
249
+ if match:
250
+ peak_memory_gib = float(match.group(1))
251
+
252
+ assert model_memory_gib is not None, "Could not find model loading memory log"
253
+ assert peak_memory_gib is not None, "Could not find peak memory log"
254
+ print(f"GPU memory used after loading weights: {model_memory_gib} GiB")
255
+ print(f"Peak GPU memory usage while loading weights: {peak_memory_gib} GiB")
256
+
257
+ # model specific, allenai/OLMoE-1B-7B-0125-Instruct fp8 online quant
258
+ # uses 6.65 GiB for weight loading (bf16 checkpoint is ~12.89 GiB)
259
+ expected_model_memory_gib = 6.7
260
+
261
+ # for allenai/OLMoE-1B-7B-0125-Instruct the number we see today is 9.06
262
+ # GiB, which is 1.36x above model_memory_gib. A slightly higher number is
263
+ # expected as when we load and quantize weights in a streaming fashion we
264
+ # need to have individual weights in bf16 + fp8 alive at the same time.
265
+ expected_peak_memory_gib = expected_model_memory_gib * 1.4
266
+
267
+ assert model_memory_gib < expected_model_memory_gib, (
268
+ f"{model_memory_gib=} higher than {expected_model_memory_gib}"
269
+ )
270
+ assert peak_memory_gib < expected_peak_memory_gib, (
271
+ f"{peak_memory_gib=} higher than {expected_peak_memory_gib}"
272
+ )
273
+
274
+
275
+@pytest.mark.skipif(
276
+ not is_quant_method_supported("fp8"),
277
+ reason="FP8 is not supported on this GPU type.",
278
+)
279
+def test_online_quant_load_format_dummy(
280
+ vllm_runner,
281
+ monkeypatch,
282
+ caplog,
283
+) -> None:
284
+ with vllm_runner(
285
+ "ibm-granite/granite-3.0-1b-a400m-base",
286
+ quantization="fp8",
287
+ enforce_eager=True,
288
+ load_format="dummy",
289
+ ) as llm:
290
+ outputs = llm.generate_greedy(["The future of AI is"], max_tokens=4)
291
+ print(outputs[0][1])
292
+
293
+
294
@pytest.mark.skipif(
295
not is_quant_method_supported("fp8"),
296
reason="FP8 is not supported on this GPU type.",
vllm/model_executor/layers/quantization/fp8.py
+119
-4
@@ -86,6 +86,7 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
86
cutlass_fp8_supported,
87
normalize_e4m3fn_to_e4m3fnuz,
88
)
89
+from vllm.model_executor.model_loader.weight_utils import initialize_single_dummy_weight
90
from vllm.model_executor.parameter import (
91
BlockQuantScaleParameter,
92
ModelWeightParameter,
@@ -293,6 +294,16 @@ class CopyNumelCounter(TorchDispatchMode):
294
return out
295
296
297
+def _copy_missing_attrs(old: torch.Tensor, new: torch.Tensor) -> None:
298
+ """Copies any attrs present in `old` but not in `new` to `new`"""
299
+ new_attrs = set(dir(new))
300
+ attrs_to_set = {}
301
+ for attr in dir(old):
302
+ if attr not in new_attrs:
303
+ attrs_to_set[attr] = getattr(old, attr)
304
+ set_weight_attrs(new, attrs_to_set)
305
+
306
+
307
class Fp8LinearMethod(LinearMethodBase):
308
"""Linear method for FP8.
309
Supports loading FP8 checkpoints with static weight scale and
@@ -578,6 +589,22 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod):
589
if not hasattr(layer, "_loaded_numel"):
590
layer._loaded_numel = 0
591
592
+ # when the first `loaded_weight` is about to be
593
+ # loaded to `param`, materialize `param` just-in-time
594
+ weight = ModelWeightParameter(
595
+ data=torch.empty_like(layer.weight, device=layer._load_device),
596
+ input_dim=1,
597
+ output_dim=0,
598
+ weight_loader=patched_weight_loader,
599
+ )
600
+ _copy_missing_attrs(layer.weight, weight)
601
+ layer.register_parameter("weight", weight)
602
+ del layer._load_device
603
+
604
+ # refresh the reference to `param` to reflect just-in-time
605
+ # materialization
606
+ param = layer.weight
607
+
608
# load the current weight chunk
609
copy_numel_counter = CopyNumelCounter()
610
with copy_numel_counter:
@@ -590,30 +617,50 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod):
617
if layer._loaded_numel == target_loaded_numel:
618
self.process_weights_after_loading(layer)
619
593
- # Delete the bookkeeping
594
- del layer._loaded_numel
620
# Prevent the usual `process_weights_after_loading` call from doing
621
# anything
622
layer._already_called_process_weights_after_loading = True
623
624
+ # Note that we keep `layer._loaded_numel` around just in case
625
+ # there is logic added to vllm in the future which calls a
626
+ # weight loader twice - we do not want to re-initialize in
627
+ # that case.
628
+
629
return res
630
631
weight = ModelWeightParameter(
632
data=torch.empty(
633
output_size_per_partition,
634
input_size_per_partition,
635
+ # materialized just-in-time in `patched_weight_loader`
636
+ device="meta",
637
dtype=params_dtype,
638
),
639
input_dim=1,
640
output_dim=0,
641
weight_loader=patched_weight_loader,
642
)
643
+ # stash the correct device for `patched_weight_loader`
644
+ layer._load_device = torch.get_default_device()
645
layer.register_parameter("weight", weight)
646
647
def process_weights_after_loading(self, layer: Module) -> None:
648
if getattr(layer, "_already_called_process_weights_after_loading", False):
649
return
650
651
+ # deferred initialization of randomly initialized weights for the
652
+ # `--load_format dummy` feature
653
+ if layer.weight.device == torch.device("meta"):
654
+ weight = ModelWeightParameter(
655
+ data=torch.empty_like(layer.weight, device=layer._load_device),
656
+ input_dim=1,
657
+ output_dim=0,
658
+ weight_loader=layer.weight.weight_loader,
659
+ )
660
+ _copy_missing_attrs(layer.weight, weight)
661
+ layer.register_parameter("weight", weight)
662
+ initialize_single_dummy_weight(layer.weight)
663
+
664
# TODO(future): support block_quant in online quant path
665
assert not self.block_quant
666
@@ -1069,6 +1116,39 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
1116
if not hasattr(layer, "_loaded_numel"):
1117
layer._loaded_numel = 0
1118
1119
+ # save the ids of original w13 and w2 so that we can
1120
+ # distinguish which one `param` should map to further
1121
+ # down in this file
1122
+ layer._w13_weight_orig_id = id(layer.w13_weight)
1123
+ layer._w2_weight_orig_id = id(layer.w2_weight)
1124
+
1125
+ # when the first `loaded_weight` is about to be
1126
+ # loaded to `param`, materialize `param` just-in-time
1127
+
1128
+ w13_weight = torch.nn.Parameter(
1129
+ torch.empty_like(layer.w13_weight, device=layer._load_device),
1130
+ requires_grad=False,
1131
+ )
1132
+ set_weight_attrs(w13_weight, extra_weight_attrs)
1133
+ _copy_missing_attrs(layer.w13_weight, w13_weight)
1134
+ layer.register_parameter("w13_weight", w13_weight)
1135
+
1136
+ w2_weight = torch.nn.Parameter(
1137
+ torch.empty_like(layer.w2_weight, device=layer._load_device),
1138
+ requires_grad=False,
1139
+ )
1140
+ set_weight_attrs(w2_weight, extra_weight_attrs)
1141
+ _copy_missing_attrs(layer.w2_weight, w2_weight)
1142
+ layer.register_parameter("w2_weight", w2_weight)
1143
+ del layer._load_device
1144
+
1145
+ # refresh the reference to `param` to reflect just-in-time
1146
+ # materialization
1147
+ if id(param) == layer._w13_weight_orig_id:
1148
+ param = layer.w13_weight
1149
+ elif id(param) == layer._w2_weight_orig_id:
1150
+ param = layer.w2_weight
1151
+
1152
# load the current weight chunk
1153
copy_numel_counter = CopyNumelCounter()
1154
with copy_numel_counter:
@@ -1081,12 +1161,16 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
1161
if layer._loaded_numel == target_loaded_numel:
1162
self.process_weights_after_loading(layer)
1163
1084
- # Delete the bookkeeping
1085
- del layer._loaded_numel
1164
# Prevent the usual `process_weights_after_loading` call
1165
# from doing anything
1166
layer._already_called_process_weights_after_loading = True
1167
1168
+ # Note that we keep `layer._loaded_numel`,
1169
+ # `layer._w13_weight_orig_id` and `layer._w2_weight_orig_id`
1170
+ # around because if EP is on, weight loaders for non-local
1171
+ # experts will run but not actually copy any elements, and we
1172
+ # need to not re-initialize in that case.
1173
+
1174
return res
1175
1176
new_extra_weight_attrs["weight_loader"] = patched_weight_loader
@@ -1098,6 +1182,8 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
1182
num_experts,
1183
2 * intermediate_size_per_partition,
1184
hidden_size,
1185
+ # materialized just-in-time in `patched_weight_loader`
1186
+ device="meta",
1187
dtype=params_dtype,
1188
),
1189
requires_grad=False,
@@ -1110,12 +1196,16 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
1196
num_experts,
1197
hidden_size,
1198
intermediate_size_per_partition,
1199
+ # materialized just-in-time in `patched_weight_loader`
1200
+ device="meta",
1201
dtype=params_dtype,
1202
),
1203
requires_grad=False,
1204
)
1205
layer.register_parameter("w2_weight", w2_weight)
1206
set_weight_attrs(w2_weight, extra_weight_attrs)
1207
+ # stash the correct device for `patched_weight_loader`
1208
+ layer._load_device = torch.get_default_device()
1209
1210
# WEIGHT_SCALES
1211
# Allocate 2 scales for w1 and w3 respectively.
@@ -1138,6 +1228,31 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
1228
if getattr(layer, "_already_called_process_weights_after_loading", False):
1229
return
1230
1231
+ # deferred initialization of randomly initialized weights for the
1232
+ # `--load_format dummy` feature
1233
+ if layer.w13_weight.device == torch.device("meta"):
1234
+ w13_weight = torch.nn.Parameter(
1235
+ torch.empty_like(layer.w13_weight, device=layer._load_device),
1236
+ requires_grad=False,
1237
+ )
1238
+ set_weight_attrs(
1239
+ w13_weight, {"weight_loader": layer.w13_weight.weight_loader}
1240
+ )
1241
+ _copy_missing_attrs(layer.w13_weight, w13_weight)
1242
+ layer.register_parameter("w13_weight", w13_weight)
1243
+ initialize_single_dummy_weight(layer.w13_weight)
1244
+ if layer.w2_weight.device == torch.device("meta"):
1245
+ w2_weight = torch.nn.Parameter(
1246
+ torch.empty_like(layer.w2_weight, device=layer._load_device),
1247
+ requires_grad=False,
1248
+ )
1249
+ set_weight_attrs(
1250
+ w2_weight, {"weight_loader": layer.w2_weight.weight_loader}
1251
+ )
1252
+ _copy_missing_attrs(layer.w2_weight, w2_weight)
1253
+ layer.register_parameter("w2_weight", w2_weight)
1254
+ initialize_single_dummy_weight(layer.w2_weight)
1255
+
1256
# If checkpoint is fp16, quantize in place.
1257
fp8_dtype = current_platform.fp8_dtype()
1258
w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
vllm/model_executor/model_loader/base_loader.py
+13
@@ -13,6 +13,8 @@ from vllm.model_executor.model_loader.utils import (
13
initialize_model,
14
process_weights_after_loading,
15
)
16
+from vllm.platforms import current_platform
17
+from vllm.utils.mem_utils import format_gib
18
from vllm.utils.torch_utils import set_default_torch_dtype
19
20
logger = init_logger(__name__)
@@ -56,6 +58,17 @@ class BaseModelLoader(ABC):
58
logger.debug("Loading weights on %s ...", load_device)
59
# Quantization does not happen in `load_weights` but after it
60
self.load_weights(model, model_config)
61
+
62
+ # Log peak GPU memory after loading weights. This is needed
63
+ # to have test coverage on peak memory for online quantization.
64
+ if current_platform.is_cuda():
65
+ peak_memory = torch.cuda.max_memory_allocated()
66
+ logger.debug_once(
67
+ "Peak GPU memory after loading weights: %s GiB",
68
+ format_gib(peak_memory),
69
+ scope="local",
70
+ )
71
+
72
process_weights_after_loading(model, model_config, target_device)
73
74
return model.eval()
vllm/model_executor/model_loader/dummy_loader.py
+1
-1
@@ -25,4 +25,4 @@ class DummyModelLoader(BaseModelLoader):
25
def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
26
# NOTE(woosuk): For accurate performance evaluation, we assign
27
# random values to the weights.
28
- initialize_dummy_weights(model)
28
+ initialize_dummy_weights(model, model_config)
vllm/model_executor/model_loader/weight_utils.py
+53
-32
@@ -1059,6 +1059,7 @@ def composed_weight_loader(
1059
1060
def initialize_dummy_weights(
1061
model: torch.nn.Module,
1062
+ model_config: ModelConfig,
1063
low: float = -1e-3,
1064
high: float = 1e-3,
1065
seed: int = 1234,
@@ -1075,41 +1076,61 @@ def initialize_dummy_weights(
1076
is fixed, the random values generated by this function only depends on
1077
the parameter's number of elements and its data type.
1078
"""
1079
+ # TODO(future PR): make the check below more generic as more online
1080
+ # quant backends are added
1081
+ is_fp8_py_quant = model_config.quantization == "fp8"
1082
+
1083
for param in model.state_dict().values():
1079
- if torch.is_floating_point(param):
1080
- if current_platform.is_tpu():
1081
- generator = torch.Generator(device="cpu")
1082
- generator.manual_seed(seed)
1083
- # Note: The param.uniform_ function cannot be used in this
1084
- # context because it demands more TPU HBM than directly copying
1085
- # from a CPU tensor.
1086
- # Note: We avoid using torch.rank_like as it doesn't currently
1087
- # support the generator argument.
1088
- param.copy_(
1089
- (high - low)
1090
- * torch.rand(
1091
- param.shape,
1092
- generator=generator,
1093
- dtype=param.dtype,
1094
- layout=param.layout,
1095
- requires_grad=param.requires_grad,
1096
- device="cpu",
1097
- )
1098
- + low
1099
- )
1100
- torch._sync(param)
1101
- continue
1084
+ if is_fp8_py_quant and param.device == torch.device("meta"):
1085
+ # for fp8.py's online quantization, dummy weight init will happen
1086
+ # in `process_weights_after_loading`.
1087
+ # TODO(future PR): consider refactoring dummy model init to compose
1088
+ # better with online quantization
1089
+ continue
1090
1103
- generator = torch.Generator(device=param.data.device)
1091
+ initialize_single_dummy_weight(param, low, high, seed)
1092
+
1093
+
1094
+def initialize_single_dummy_weight(
1095
+ param: torch.Tensor,
1096
+ low: float = -1e-3,
1097
+ high: float = 1e-3,
1098
+ seed: int = 1234,
1099
+) -> None:
1100
+ if torch.is_floating_point(param):
1101
+ if current_platform.is_tpu():
1102
+ generator = torch.Generator(device="cpu")
1103
generator.manual_seed(seed)
1105
- if torch.finfo(param.data.dtype).bits < 16:
1106
- # uniform_ doesn't support < 16-bit datatypes (FP8)
1107
- dtype = param.data.dtype
1108
- tmp_param = param.data.to(torch.float16)
1109
- tmp_param = tmp_param.uniform_(low, high, generator=generator).to(dtype)
1110
- param.data.copy_(tmp_param)
1111
- else:
1112
- param.uniform_(low, high, generator=generator)
1104
+ # Note: The param.uniform_ function cannot be used in this
1105
+ # context because it demands more TPU HBM than directly copying
1106
+ # from a CPU tensor.
1107
+ # Note: We avoid using torch.rank_like as it doesn't currently
1108
+ # support the generator argument.
1109
+ param.copy_(
1110
+ (high - low)
1111
+ * torch.rand(
1112
+ param.shape,
1113
+ generator=generator,
1114
+ dtype=param.dtype,
1115
+ layout=param.layout,
1116
+ requires_grad=param.requires_grad,
1117
+ device="cpu",
1118
+ )
1119
+ + low
1120
+ )
1121
+ torch._sync(param)
1122
+ return
1123
+
1124
+ generator = torch.Generator(device=param.data.device)
1125
+ generator.manual_seed(seed)
1126
+ if torch.finfo(param.data.dtype).bits < 16:
1127
+ # uniform_ doesn't support < 16-bit datatypes (FP8)
1128
+ dtype = param.data.dtype
1129
+ tmp_param = param.data.to(torch.float16)
1130
+ tmp_param = tmp_param.uniform_(low, high, generator=generator).to(dtype)
1131
+ param.data.copy_(tmp_param)
1132
+ else:
1133
+ param.uniform_(low, high, generator=generator)
1134
1135
1136
def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None: