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, 07:15 UTC
Fix AOPerModuleConfig name changes (#18869)
Signed-off-by: Jerry Zhang <jerryzh168@gmail.com>
Jerry Zhang committed
Jun 5, 2025 at 21:51 UTC
c8134bea15826876e37694834ad87d9c4bdfb26b
3 files changed
+25
-5
.buildkite/test-pipeline.yaml
+3
@@ -424,6 +424,9 @@ steps:
424
- vllm/model_executor/layers/quantization
425
- tests/quantization
426
commands:
427
+ # temporary install here since we need nightly, will move to requirements/test.in
428
+ # after torchao 0.12 release
429
+ - pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu126
430
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization
431
432
- label: LM Eval Small Models # 53min
tests/quantization/test_torchao.py
+3
-3
@@ -13,7 +13,7 @@ TORCHAO_AVAILABLE = importlib.util.find_spec("torchao") is not None
13
14
@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
15
def test_pre_quantized_model(vllm_runner):
16
- with vllm_runner("drisspg/float8_dynamic_act_float8_weight-opt-125m",
16
+ with vllm_runner("drisspg/fp8-opt-125m",
17
quantization="torchao",
18
dtype="bfloat16",
19
enforce_eager=True) as llm:
@@ -30,10 +30,10 @@ def test_pre_quantized_model(vllm_runner):
30
"cuda:0",
31
# {"": "cuda"},
32
])
33
-def test_opt_125m_int4wo_model_loading_with_params(vllm_runner,
33
+def test_opt_125m_int8wo_model_loading_with_params(vllm_runner,
34
pt_load_map_location):
35
torch._dynamo.reset()
36
- model_name = "jerryzh168/opt-125m-int4wo"
36
+ model_name = "jerryzh168/opt-125m-int8wo-partial-quant"
37
with vllm_runner(model_name=model_name,
38
quantization="torchao",
39
dtype="bfloat16",
vllm/model_executor/layers/quantization/torchao.py
+19
-2
@@ -6,6 +6,7 @@ import torch
6
import torch.nn.functional as F
7
from torch.nn.parameter import Parameter
8
9
+from vllm.logger import init_logger
10
from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
11
UnquantizedLinearMethod)
12
from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -13,12 +14,28 @@ from vllm.model_executor.layers.quantization.base_config import (
14
QuantizationConfig, QuantizeMethodBase)
15
from vllm.model_executor.utils import set_weight_attrs
16
17
+logger = init_logger(__name__)
18
+
19
20
class TorchAOConfig(QuantizationConfig):
21
"""Config class for torchao."""
22
23
def __init__(self, torchao_config) -> None:
24
self.torchao_config = torchao_config
25
+ """
26
+ # TorchAO quantization relies on tensor subclasses. In order,
27
+ # to enable proper caching this needs standalone compile
28
+ if is_torch_equal_or_newer("2.8.0"):
29
+ os.environ["VLLM_TEST_STANDALONE_COMPILE"] = "1"
30
+ logger.info(
31
+ "Using TorchAO: Setting VLLM_TEST_STANDALONE_COMPILE=1")
32
+
33
+ # TODO: remove after the torch dependency is updated to 2.8
34
+ if is_torch_equal_or_newer(
35
+ "2.7.0") and not is_torch_equal_or_newer("2.8.0"):
36
+ os.environ["VLLM_DISABLE_COMPILE_CACHE"] = "1"
37
+ logger.info("Using TorchAO: Setting VLLM_DISABLE_COMPILE_CACHE=1")
38
+ """
39
40
def __repr__(self) -> str:
41
return f"TorchAOConfig({self.torchao_config})"
@@ -61,10 +78,10 @@ class TorchAOConfig(QuantizationConfig):
78
if not isinstance(layer, LinearBase):
79
return None
80
64
- from torchao.quantization import AOPerModuleConfig
81
+ from torchao.quantization import ModuleFqnToConfig
82
83
module_fqn = prefix
67
- if isinstance(self.torchao_config, AOPerModuleConfig):
84
+ if isinstance(self.torchao_config, ModuleFqnToConfig):
85
module_fqn_to_config = self.torchao_config.module_fqn_to_config
86
c = module_fqn_to_config.get(
87
module_fqn) or module_fqn_to_config.get("_default", None)