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 8, 03:45 UTC
[Quantization] - Added uses_meta_device_weights to quant config (#34645)
Signed-off-by: Josephasafg <ajgard7@gmail.com>
Asaf Joseph Gardin committed
Feb 18, 2026 at 09:43 UTC
1faa8cb73cad1346224284002c136fb060dc89c9
3 files changed
+21
-8
vllm/model_executor/layers/quantization/base_config.py
+5
@@ -18,6 +18,11 @@ else:
18
class QuantizeMethodBase(ABC):
19
"""Base class for different quantized methods."""
20
21
+ # Whether this method creates weights on meta device for online quantization.
22
+ # When True, weights are created on meta device and quantized layer-wise
23
+ # in process_weights_after_loading, reducing peak memory during loading.
24
+ uses_meta_device: bool = False
25
+
26
@abstractmethod
27
def create_weights(
28
self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs
vllm/model_executor/layers/quantization/fp8.py
+4
@@ -527,6 +527,8 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod):
527
"""Online version of Fp8LinearMethod, loads the fp16/bf16 checkpoint
528
and quantized the weights during loading."""
529
530
+ uses_meta_device: bool = True
531
+
532
def create_weights(
533
self,
534
layer: torch.nn.Module,
@@ -1039,6 +1041,8 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
1041
quant_config: The quantization config.
1042
"""
1043
1044
+ uses_meta_device: bool = True
1045
+
1046
def __init__(self, quant_config: Fp8Config, layer: torch.nn.Module):
1047
super().__init__(quant_config, layer)
1048
assert not quant_config.is_checkpoint_fp8_serialized
vllm/model_executor/model_loader/weight_utils.py
+12
-8
@@ -1092,16 +1092,20 @@ def initialize_dummy_weights(
1092
is fixed, the random values generated by this function only depends on
1093
the parameter's number of elements and its data type.
1094
"""
1095
- # TODO(future PR): make the check below more generic as more online
1096
- # quant backends are added
1097
- is_fp8_py_quant = model_config.quantization == "fp8"
1095
+
1096
+ # Check if any module uses online quantization with meta device weights.
1097
+ # If so, we'll skip initializing params on meta device since they'll be
1098
+ # handled in `process_weights_after_loading`.
1099
+ def uses_meta_device(module: torch.nn.Module) -> bool:
1100
+ quant_method = getattr(module, "quant_method", None)
1101
+ return getattr(quant_method, "uses_meta_device", False)
1102
+
1103
+ has_online_quant = any(uses_meta_device(m) for m in model.modules())
1104
1105
for param in model.state_dict().values():
1100
- if is_fp8_py_quant and param.device == torch.device("meta"):
1101
- # for fp8.py's online quantization, dummy weight init will happen
1102
- # in `process_weights_after_loading`.
1103
- # TODO(future PR): consider refactoring dummy model init to compose
1104
- # better with online quantization
1106
+ if has_online_quant and param.device == torch.device("meta"):
1107
+ # For online quantization, weights are created on meta device and
1108
+ # dummy weight init will happen in `process_weights_after_loading`.
1109
continue
1110
1111
initialize_single_dummy_weight(param, low, high, seed)