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 11, 17:30 UTC
[Quantization] FP8 Weight Reloading for Quantized RL Rollout (#28480)
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Kyle Sayers committed
Dec 9, 2025 at 16:54 UTC
fccd5325874321f34daef0a8ed4d1b15b26ca34e
7 files changed
+207
-87
tests/quantization/test_fp8.py
+88
@@ -10,10 +10,14 @@ import torch
10
11
from tests.quantization.utils import is_quant_method_supported
12
from vllm import _custom_ops as ops
13
+from vllm.model_executor.layers.fused_moe import FusedMoE
14
from vllm.model_executor.layers.quantization.fp8 import (
15
+ Fp8Config,
16
Fp8KVCacheMethod,
17
Fp8LinearMethod,
18
+ Fp8MoEMethod,
19
)
20
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
21
from vllm.platforms import current_platform
22
23
MODELS = [
@@ -261,3 +265,87 @@ def test_scaled_fp8_quant(dtype) -> None:
265
torch.narrow(y_nc_pad, 0, 0, x_nc.shape[0]), inv_scale_nc, dtype
266
),
267
)
268
+
269
+
270
+@pytest.mark.parametrize("method_cls", [Fp8LinearMethod, Fp8MoEMethod])
271
+# FP8 weight reloading does not support online quantization
272
+@pytest.mark.parametrize("is_checkpoint_fp8_serialized", [True]) # skip False
273
+@pytest.mark.parametrize("weight_block_size", [None, [1, 1]])
274
+# any postprocessing that is applied to the weights such as padding and repacking
275
+# (excluding device sharding) must also be applied to the reloaded weights
276
+#
277
+# this is the case for marlin as well as per-tensor Fp8MoEMethod
278
+@pytest.mark.parametrize("use_marlin", [False]) # skip True
279
+def test_fp8_reloading(
280
+ method_cls, is_checkpoint_fp8_serialized, weight_block_size, use_marlin, dist_init
281
+):
282
+ if is_checkpoint_fp8_serialized is False:
283
+ pytest.skip("FP8 weight reloading does not support online quantization")
284
+
285
+ if method_cls is Fp8MoEMethod and weight_block_size is None:
286
+ pytest.skip(
287
+ "FP8 Tensor weight reloading does not support fusing w13_weight_scale. "
288
+ "If this is your use case, consider using a restore function like #26327"
289
+ )
290
+
291
+ with torch.device("cuda:0"):
292
+ config = Fp8Config(
293
+ is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized,
294
+ weight_block_size=weight_block_size,
295
+ )
296
+
297
+ if method_cls is Fp8LinearMethod:
298
+ layer = torch.nn.Linear(1, 1)
299
+ method = method_cls(config)
300
+ method.create_weights(
301
+ layer=layer,
302
+ input_size_per_partition=1,
303
+ output_partition_sizes=[1],
304
+ input_size=1,
305
+ output_size=1,
306
+ params_dtype=torch.bfloat16,
307
+ weight_loader=default_weight_loader,
308
+ )
309
+
310
+ else:
311
+ layer = FusedMoE(
312
+ num_experts=1,
313
+ top_k=1,
314
+ hidden_size=1,
315
+ intermediate_size=1,
316
+ )
317
+ method = method_cls(config, layer)
318
+ method.create_weights(
319
+ layer=layer,
320
+ num_experts=1,
321
+ hidden_size=1,
322
+ intermediate_size_per_partition=1,
323
+ params_dtype=torch.bfloat16,
324
+ weight_loader=default_weight_loader,
325
+ )
326
+
327
+ method.use_marlin = use_marlin
328
+
329
+ # capture weights format during loading
330
+ original_metadata = [
331
+ (name, param.shape, getattr(param, "weight_loader", default_weight_loader))
332
+ for name, param in layer.named_parameters()
333
+ ]
334
+
335
+ # test loading
336
+ for name, shape, _ in original_metadata:
337
+ param = getattr(layer, name)
338
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
339
+ weight_loader(param, torch.zeros(shape)) # cannot use empty
340
+
341
+ method.process_weights_after_loading(layer)
342
+
343
+ # test reloading works after loading
344
+ # assuming that no reshaping occurred
345
+ for name, shape, original_weight_loader in original_metadata:
346
+ param = getattr(layer, name)
347
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
348
+ assert weight_loader is original_weight_loader
349
+ weight_loader(param, torch.zeros(shape)) # cannot use empty
350
+
351
+ method.process_weights_after_loading(layer)
vllm/model_executor/layers/quantization/fp8.py
+72
-79
@@ -94,7 +94,7 @@ from vllm.model_executor.parameter import (
94
ModelWeightParameter,
95
PerTensorScaleParameter,
96
)
97
-from vllm.model_executor.utils import set_weight_attrs
97
+from vllm.model_executor.utils import replace_parameter, set_weight_attrs
98
from vllm.platforms import current_platform
99
from vllm.scalar_type import scalar_types
100
from vllm.utils.deep_gemm import (
@@ -548,46 +548,50 @@ class Fp8LinearMethod(LinearMethodBase):
548
assert not self.act_q_static
549
size_k_first = False
550
551
- weight, weight_scale = process_fp8_weight_block_strategy(
551
+ weight, weight_scale_inv = process_fp8_weight_block_strategy(
552
layer.weight, layer.weight_scale_inv
553
)
554
- # Delete the weight_scale_inv parameter to avoid confusion
555
- # with the weight_scale parameter
556
- del layer.weight_scale_inv
554
+
555
+ # Update layer with new values
556
+ replace_parameter(layer, "weight", weight.data)
557
+ replace_parameter(layer, "weight_scale_inv", weight_scale_inv.data)
558
559
# If checkpoint not serialized fp8, quantize the weights.
559
- elif not self.quant_config.is_checkpoint_fp8_serialized:
560
- qweight, weight_scale = ops.scaled_fp8_quant(layer.weight, scale=None)
561
- weight = qweight.t()
560
+ else:
561
+ if not self.quant_config.is_checkpoint_fp8_serialized:
562
+ qweight, weight_scale = ops.scaled_fp8_quant(layer.weight, scale=None)
563
+ weight = qweight.t()
564
+
565
+ # If checkpoint is fp8 per-tensor, handle that there are N scales for N
566
+ # shards in a fused module
567
+ else:
568
+ weight = layer.weight
569
+ weight_scale = layer.weight_scale
570
+
571
+ # If using w8a8, torch._scaled_mm needs per tensor, so
572
+ # requantize the logical shards as a single weight.
573
+ if not self.use_marlin:
574
+ weight, weight_scale, input_scale = (
575
+ process_fp8_weight_tensor_strategy(
576
+ weight,
577
+ weight_scale,
578
+ layer.logical_widths,
579
+ getattr(layer, "input_scale", None),
580
+ )
581
+ )
582
+ if self.act_q_static:
583
+ assert input_scale is not None
584
+ input_scale = input_scale.max()
585
+ weight = weight.t()
586
563
- # If checkpoint is fp8 per-tensor, handle that there are N scales for N
564
- # shards in a fused module
587
+ # Update layer with new values.
588
+ replace_parameter(layer, "weight", weight.data)
589
+ replace_parameter(layer, "weight_scale", weight_scale.data)
590
+
591
+ if input_scale is not None:
592
+ replace_parameter(layer, "input_scale", input_scale)
593
else:
566
- weight = layer.weight
567
- weight_scale = layer.weight_scale
568
-
569
- # If using w8a8, torch._scaled_mm needs per tensor, so
570
- # requantize the logical shards as a single weight.
571
- if not self.use_marlin:
572
- weight, weight_scale, input_scale = process_fp8_weight_tensor_strategy(
573
- weight,
574
- weight_scale,
575
- layer.logical_widths,
576
- getattr(layer, "input_scale", None),
577
- )
578
- if self.act_q_static:
579
- assert input_scale is not None
580
- input_scale = input_scale.max()
581
- weight = weight.t()
582
-
583
- # Update layer with new values.
584
- layer.weight = Parameter(weight.data, requires_grad=False)
585
- layer.weight_scale = Parameter(weight_scale.data, requires_grad=False)
586
- layer.input_scale = (
587
- Parameter(input_scale, requires_grad=False)
588
- if input_scale is not None
589
- else None
590
- )
594
+ layer.input_scale = None
595
596
if self.use_marlin:
597
prepare_fp8_layer_for_marlin(
@@ -614,7 +618,7 @@ class Fp8LinearMethod(LinearMethodBase):
618
return self.w8a8_block_fp8_linear.apply(
619
input=x,
620
weight=layer.weight,
617
- weight_scale=layer.weight_scale,
621
+ weight_scale=layer.weight_scale_inv,
622
input_scale=layer.input_scale,
623
bias=bias,
624
)
@@ -643,10 +647,15 @@ class Fp8LinearMethod(LinearMethodBase):
647
return torch.nn.functional.linear(x, weight_bf16.t(), bias)
648
649
if self.use_marlin:
650
+ if self.block_quant:
651
+ weight_scale = layer.weight_scale_inv
652
+ else:
653
+ weight_scale = layer.weight_scale
654
+
655
return apply_fp8_marlin_linear(
656
input=x,
657
weight=layer.weight,
649
- weight_scale=layer.weight_scale,
658
+ weight_scale=weight_scale,
659
workspace=layer.workspace,
660
size_n=layer.output_size_per_partition,
661
size_k=layer.input_size_per_partition,
@@ -660,7 +669,7 @@ class Fp8LinearMethod(LinearMethodBase):
669
return self.w8a8_block_fp8_linear.apply(
670
input=x,
671
weight=layer.weight,
663
- weight_scale=layer.weight_scale,
672
+ weight_scale=layer.weight_scale_inv,
673
input_scale=layer.input_scale,
674
bias=bias,
675
)
@@ -937,22 +946,18 @@ class Fp8MoEMethod(FusedMoEMethodBase):
946
w2_weight_scale_inv = layer.w2_weight_scale_inv
947
948
# torch.compile() cannot use Parameter subclasses.
940
- layer.w13_weight = Parameter(w13_weight, requires_grad=False)
941
- layer.w13_weight_scale_inv = Parameter(
942
- w13_weight_scale_inv, requires_grad=False
943
- )
944
- layer.w2_weight = Parameter(w2_weight, requires_grad=False)
945
- layer.w2_weight_scale_inv = Parameter(
946
- w2_weight_scale_inv, requires_grad=False
947
- )
949
+ replace_parameter(layer, "w13_weight", w13_weight)
950
+ replace_parameter(layer, "w13_weight_scale_inv", w13_weight_scale_inv)
951
+ replace_parameter(layer, "w2_weight", w2_weight)
952
+ replace_parameter(layer, "w2_weight_scale_inv", w2_weight_scale_inv)
953
if self.rocm_aiter_moe_enabled:
954
# reshaping weights is required for aiter moe kernel.
955
shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
956
layer.w13_weight.data, layer.w2_weight.data
957
)
958
954
- layer.w13_weight = torch.nn.Parameter(shuffled_w13, requires_grad=False)
955
- layer.w2_weight = torch.nn.Parameter(shuffled_w2, requires_grad=False)
959
+ replace_parameter(layer, "w13_weight", shuffled_w13)
960
+ replace_parameter(layer, "w2_weight", shuffled_w2)
961
962
# DeepGemm scales need to be transposed and aligned. We try to do
963
# it ahead of time for performance reasons.
@@ -990,13 +995,14 @@ class Fp8MoEMethod(FusedMoEMethodBase):
995
996
# Re-initialize w13_scale because we directly quantize
997
# merged w13 weights and generate a single scaling factor.
993
- layer.w13_weight_scale = torch.nn.Parameter(
998
+ replace_parameter(
999
+ layer,
1000
+ "w13_weight_scale",
1001
torch.ones(
1002
layer.local_num_experts,
1003
dtype=torch.float32,
1004
device=w13_weight.device,
1005
),
999
- requires_grad=False,
1006
)
1007
for expert in range(layer.local_num_experts):
1008
w13_weight[expert, :, :], layer.w13_weight_scale[expert] = (
@@ -1005,16 +1011,17 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1011
w2_weight[expert, :, :], layer.w2_weight_scale[expert] = (
1012
ops.scaled_fp8_quant(layer.w2_weight.data[expert, :, :])
1013
)
1008
- layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
1009
- layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
1014
+ replace_parameter(layer, "w13_weight", w13_weight)
1015
+ replace_parameter(layer, "w2_weight", w2_weight)
1016
+
1017
if self.rocm_aiter_moe_enabled:
1018
# reshaping weights is required for aiter moe kernel.
1019
shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
1020
layer.w13_weight, layer.w2_weight
1021
)
1022
1016
- layer.w13_weight = torch.nn.Parameter(shuffled_w13, requires_grad=False)
1017
- layer.w2_weight = torch.nn.Parameter(shuffled_w2, requires_grad=False)
1023
+ replace_parameter(layer, "w13_weight", shuffled_w13)
1024
+ replace_parameter(layer, "w2_weight", shuffled_w2)
1025
# If checkpoint is fp8, we need to handle that the
1026
# MoE kernels require single activation scale and single weight
1027
# scale for w13 per expert.
@@ -1035,12 +1042,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1042
"fp8 MoE layer. Using the maximum across experts "
1043
"for each layer."
1044
)
1038
- layer.w13_input_scale = torch.nn.Parameter(
1039
- layer.w13_input_scale.max(), requires_grad=False
1040
- )
1041
- layer.w2_input_scale = torch.nn.Parameter(
1042
- layer.w2_input_scale.max(), requires_grad=False
1043
- )
1045
+ replace_parameter(layer, "w13_input_scale", layer.w13_input_scale.max())
1046
+ replace_parameter(layer, "w2_input_scale", layer.w2_input_scale.max())
1047
if current_platform.is_fp8_fnuz():
1048
# Normalize the weights and scales
1049
w13_weight, w13_weight_scale, w13_input_scale = (
@@ -1054,22 +1057,14 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1057
)
1058
)
1059
# Reset the parameter
1057
- layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
1058
- layer.w13_weight_scale = torch.nn.Parameter(
1059
- w13_weight_scale, requires_grad=False
1060
- )
1060
+ replace_parameter(layer, "w13_weight", w13_weight)
1061
+ replace_parameter(layer, "w13_weight_scale", w13_weight_scale)
1062
if w13_input_scale is not None:
1062
- layer.w13_input_scale = torch.nn.Parameter(
1063
- w13_input_scale, requires_grad=False
1064
- )
1065
- layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
1066
- layer.w2_weight_scale = torch.nn.Parameter(
1067
- w2_weight_scale, requires_grad=False
1068
- )
1063
+ replace_parameter(layer, "w13_input_scale", w13_input_scale)
1064
+ replace_parameter(layer, "w2_weight", w2_weight)
1065
+ replace_parameter(layer, "w2_weight_scale", w2_weight_scale)
1066
if w2_input_scale is not None:
1070
- layer.w2_input_scale = torch.nn.Parameter(
1071
- w2_input_scale, requires_grad=False
1072
- )
1067
+ replace_parameter(layer, "w2_input_scale", w2_input_scale)
1068
1069
# Fp8 moe kernel needs single weight scale for w13 per expert.
1070
# We take the max then dequant and requant each expert.
@@ -1093,12 +1088,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1088
layer.w13_weight, layer.w2_weight
1089
)
1090
1096
- layer.w13_weight = torch.nn.Parameter(shuffled_w13, requires_grad=False)
1097
- layer.w2_weight = torch.nn.Parameter(shuffled_w2, requires_grad=False)
1091
+ replace_parameter(layer, "w13_weight", shuffled_w13)
1092
+ replace_parameter(layer, "w2_weight", shuffled_w2)
1093
1099
- layer.w13_weight_scale = torch.nn.Parameter(
1100
- max_w13_scales, requires_grad=False
1101
- )
1094
+ replace_parameter(layer, "w13_weight_scale", max_w13_scales)
1095
1096
if self.flashinfer_moe_backend is not None:
1097
# NOTE: weights have to be swapped since the activation is
vllm/model_executor/layers/quantization/kv_cache.py
+7
@@ -45,6 +45,13 @@ class BaseKVCacheMethod(QuantizeMethodBase):
45
raise RuntimeError(f"{self.__class__.__name__}.apply should not be called.")
46
47
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
48
+ # skip if there are no weights to process (for example, weight reloading)
49
+ if not hasattr(layer, "q_scale"):
50
+ assert not hasattr(layer, "k_scale")
51
+ assert not hasattr(layer, "v_scale")
52
+ assert not hasattr(layer, "prob_scale")
53
+ return
54
+
55
# If the kv-cache dtype is auto, we enforce the k/v_scale to be 1.0
56
# regardless whether the kv-scale is available in the checkpoint.
57
# No need to process kv scales after loading if we are going to
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+4
-3
@@ -27,6 +27,7 @@ from vllm.model_executor.parameter import (
27
ChannelQuantScaleParameter,
28
PerTensorScaleParameter,
29
)
30
+from vllm.model_executor.utils import replace_parameter
31
from vllm.platforms import current_platform
32
from vllm.triton_utils import tl, triton
33
from vllm.utils.deep_gemm import (
@@ -1404,12 +1405,12 @@ def maybe_post_process_fp8_weight_block(layer: torch.nn.Module):
1405
if should_use_deepgemm:
1406
dg_weight, dg_weight_scale = deepgemm_post_process_fp8_weight_block(
1407
wq=layer.weight.data,
1407
- ws=layer.weight_scale.data,
1408
+ ws=layer.weight_scale_inv.data,
1409
quant_block_shape=tuple(layer.weight_block_size),
1410
use_e8m0=is_deep_gemm_e8m0_used(),
1411
)
1411
- layer.weight = torch.nn.Parameter(dg_weight, requires_grad=False)
1412
- layer.weight_scale = torch.nn.Parameter(dg_weight_scale, requires_grad=False)
1412
+ replace_parameter(layer, "weight", dg_weight)
1413
+ replace_parameter(layer, "weight_scale_inv", dg_weight_scale)
1414
1415
1416
def expert_weight_is_col_major(x: torch.Tensor) -> bool:
vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
+7
-4
@@ -14,6 +14,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
14
marlin_quant_input,
15
should_use_atomic_add_reduce,
16
)
17
+from vllm.model_executor.utils import replace_parameter
18
from vllm.platforms import current_platform
19
from vllm.scalar_type import scalar_types
20
@@ -130,7 +131,7 @@ def prepare_fp8_layer_for_marlin(
131
size_n=part_size_n,
132
num_bits=8,
133
)
133
- layer.weight = torch.nn.Parameter(marlin_qweight, requires_grad=False)
134
+ replace_parameter(layer, "weight", marlin_qweight)
135
136
# WEIGHT SCALES
137
# Permute scales
@@ -138,7 +139,6 @@ def prepare_fp8_layer_for_marlin(
139
scales = layer.weight_scale.to(layer.orig_dtype)
140
elif "weight_scale_inv" in dir(layer):
141
scales = layer.weight_scale_inv.to(layer.orig_dtype)
141
- del layer.weight_scale_inv
142
143
group_size = -1 if weight_block_size is None else weight_block_size[1]
144
@@ -177,12 +177,15 @@ def prepare_fp8_layer_for_marlin(
177
)
178
if input_dtype != torch.float8_e4m3fn:
179
marlin_scales = fp8_fused_exponent_bias_into_scales(marlin_scales)
180
- layer.weight_scale = torch.nn.Parameter(marlin_scales, requires_grad=False)
180
+ if hasattr(layer, "weight_scale"):
181
+ replace_parameter(layer, "weight_scale", marlin_scales)
182
+ elif hasattr(layer, "weight_scale_inv"):
183
+ replace_parameter(layer, "weight_scale_inv", marlin_scales)
184
185
if hasattr(layer, "bias") and layer.bias is not None:
186
assert layer.bias.shape == (part_size_n,)
187
bias = marlin_permute_bias(layer.bias)
185
- layer.bias = torch.nn.Parameter(bias, requires_grad=False)
188
+ replace_parameter(layer, "bias", bias)
189
190
191
def prepare_moe_fp8_layer_for_marlin(
vllm/model_executor/layers/quantization/utils/w8a8_utils.py
+4
-1
@@ -118,8 +118,11 @@ def requantize_with_max_scale(
118
# from disk in this case. Skip requantization in this case (since)
119
# we already are quantized with the single scale.
120
# * Sample Model: nm-testing/Phi-3-mini-128k-instruct-FP8
121
+ #
122
+ # Extra note: upon weight reloading weight_scale.ndim == 0
123
unfused_module_in_checkpoint = (
122
- weight_scale[-1] > torch.finfo(torch.float8_e4m3fn).min
124
+ weight_scale.ndim != 0
125
+ and weight_scale[-1] > torch.finfo(torch.float8_e4m3fn).min
126
)
127
128
# If unfused checkpoint, need requanize with the single scale.
vllm/model_executor/utils.py
+25
@@ -50,6 +50,31 @@ def set_weight_attrs(
50
setattr(weight, key, value)
51
52
53
+def replace_parameter(layer: torch.nn.Module, param_name: str, new_data: torch.Tensor):
54
+ """
55
+ Replace a parameter of a layer while maintaining the ability to reload the weight.
56
+ Called within implementations of the `process_weights_after_loading` method.
57
+
58
+ This function should not be called on weights which are tied/shared
59
+
60
+ Args:
61
+ layer: Layer containing parameter to replace
62
+ param_name: Name of parameter to replace
63
+ new_data: New data of the new parameter
64
+ """
65
+ # should not be used on a tied/shared param
66
+ if isinstance(new_data, torch.nn.Parameter):
67
+ new_data = new_data.data
68
+ new_param = torch.nn.Parameter(new_data, requires_grad=False)
69
+
70
+ old_param: torch.nn.Parameter | None = getattr(layer, param_name, None)
71
+ if old_param is not None and hasattr(old_param, "weight_loader"):
72
+ weight_loader = old_param.weight_loader
73
+ set_weight_attrs(new_param, {"weight_loader": weight_loader})
74
+
75
+ setattr(layer, param_name, new_param)
76
+
77
+
78
def get_packed_modules_mapping(model: torch.nn.Module) -> dict[str, list[str]]:
79
parent_map = getattr(model, "packed_modules_mapping", None)
80
parent_map = copy.deepcopy(parent_map) if parent_map is not None else {}