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 15, 09:45 UTC
Add support for loading torchao models with `AOPerModuleConfig` (#17826)
Signed-off-by: Jerry Zhang <jerryzh168@gmail.com>
Jerry Zhang committed
May 14, 2025 at 16:24 UTC
7974736740aee195e454eb989cd9112577bc8652
2 files changed
+37
-10
tests/quantization/test_torchao.py
+15
-3
@@ -31,9 +31,6 @@ def test_pre_quantized_model(vllm_runner):
31
])
32
def test_opt_125m_int4wo_model_loading_with_params(vllm_runner,
33
pt_load_map_location):
34
- """
35
- Test loading roberta-base model with no lm_head.
36
- """
34
torch._dynamo.reset()
35
model_name = "jerryzh168/opt-125m-int4wo"
36
with vllm_runner(model_name=model_name,
@@ -47,5 +44,20 @@ def test_opt_125m_int4wo_model_loading_with_params(vllm_runner,
44
print(output)
45
46
47
+@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
48
+def test_opt_125m_int4wo_model_per_module_quant(vllm_runner):
49
+ torch._dynamo.reset()
50
+ model_name = "jerryzh168/opt-125m-int4wo-per-module"
51
+ with vllm_runner(model_name=model_name,
52
+ quantization="torchao",
53
+ dtype="bfloat16",
54
+ pt_load_map_location="cuda:0") as llm:
55
+ output = llm.generate_greedy(["The capital of France is"],
56
+ max_tokens=32)
57
+
58
+ assert output
59
+ print(output)
60
+
61
+
62
if __name__ == "__main__":
63
pytest.main([__file__])
vllm/model_executor/layers/quantization/torchao.py
+22
-7
@@ -5,10 +5,11 @@ import torch
5
import torch.nn.functional as F
6
from torch.nn.parameter import Parameter
7
8
-from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase
8
+from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
9
+ UnquantizedLinearMethod)
10
from vllm.model_executor.layers.quantization import QuantizationMethods
11
from vllm.model_executor.layers.quantization.base_config import (
11
- QuantizationConfig)
12
+ QuantizationConfig, QuantizeMethodBase)
13
from vllm.model_executor.utils import set_weight_attrs
14
15
@@ -55,10 +56,24 @@ class TorchAOConfig(QuantizationConfig):
56
return cls(ao_config)
57
58
def get_quant_method(self, layer: torch.nn.Module,
58
- prefix: str) -> Optional["TorchAOLinearMethod"]:
59
- if isinstance(layer, LinearBase):
60
- return TorchAOLinearMethod(self)
61
- return None
59
+ prefix: str) -> Optional["QuantizeMethodBase"]:
60
+ if not isinstance(layer, LinearBase):
61
+ return None
62
+
63
+ from torchao.quantization import AOPerModuleConfig
64
+
65
+ module_fqn = prefix
66
+ if isinstance(self.torchao_config, AOPerModuleConfig):
67
+ module_fqn_to_config = self.torchao_config.module_fqn_to_config
68
+ c = module_fqn_to_config.get(
69
+ module_fqn) or module_fqn_to_config.get("_default", None)
70
+ if c is not None:
71
+ current_torchao_config = TorchAOConfig(c)
72
+ return TorchAOLinearMethod(current_torchao_config)
73
+ else:
74
+ return UnquantizedLinearMethod()
75
+
76
+ return TorchAOLinearMethod(self)
77
78
def get_scaled_act_names(self) -> list[str]:
79
return []
@@ -75,7 +90,7 @@ def torchao_quantize_param_data(param: torch.Tensor,
90
"""
91
from torchao.core.config import AOBaseConfig
92
from torchao.quantization import quantize_
78
- assert isinstance(torchao_config, AOBaseConfig)
93
+ assert isinstance(torchao_config, AOBaseConfig), f"{torchao_config}"
94
dummy_linear = torch.nn.Linear(param.shape[1], param.shape[0], bias=False)
95
dummy_linear.weight = param
96
quantize_(dummy_linear, torchao_config)