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 14, 05:30 UTC
[torchao] Support quantization configs using module swap (#21982)
Signed-off-by: Jerry Zhang <jerryzh168@gmail.com>
Jerry Zhang committed
Sep 10, 2025 at 23:53 UTC
2048c4e37909a42847cd2f51c7e0cf92e3b63466
3 files changed
+33
-7
.buildkite/test-pipeline.yaml
+4
@@ -507,6 +507,10 @@ steps:
507
commands:
508
# temporary install here since we need nightly, will move to requirements/test.in
509
# after torchao 0.12 release, and pin a working version of torchao nightly here
510
+
511
+ # since torchao nightly is only compatible with torch nightly currently
512
+ # https://github.com/pytorch/ao/issues/2919, we'll have to skip new torchao tests for now
513
+ # we can only upgrade after this is resolved
514
- pip install --pre torchao==0.13.0.dev20250814 --index-url https://download.pytorch.org/whl/nightly/cu128
515
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization
516
tests/quantization/test_torchao.py
+20
@@ -75,5 +75,25 @@ def test_qwenvl_int8wo_model_loading_with_params(vllm_runner):
75
print(output)
76
77
78
+@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
79
+@pytest.mark.skip(
80
+ reason="since torchao nightly is only compatible with torch nightly"
81
+ "currently https://github.com/pytorch/ao/issues/2919, we'll have to skip "
82
+ "torchao tests that requires newer versions (0.14.0.dev+) for now")
83
+def test_opt_125m_awq_int4wo_model_loading_with_params(vllm_runner):
84
+ torch._dynamo.reset()
85
+ model_name = ("torchao-testing/opt-125m-AWQConfig-Int4WeightOnlyConfig-v2"
86
+ "-0.14.0.dev")
87
+ with vllm_runner(model_name=model_name,
88
+ quantization="torchao",
89
+ dtype="bfloat16",
90
+ pt_load_map_location="cuda:0") as llm:
91
+ output = llm.generate_greedy(["The capital of France is"],
92
+ max_tokens=32)
93
+
94
+ assert output
95
+ print(output)
96
+
97
+
98
if __name__ == "__main__":
99
pytest.main([__file__])
vllm/model_executor/layers/quantization/torchao.py
+9
-7
@@ -152,18 +152,20 @@ def torchao_quantize_param_data(param: torch.Tensor,
152
from torchao.quantization import quantize_
153
154
assert isinstance(torchao_config, AOBaseConfig), f"{torchao_config}"
155
- """
156
- Avoid real weight allocation for faster load, since we will
155
+ """
156
+ Avoid real weight allocation for faster load, since we will
157
end up setting it to param.
158
"""
159
with torch.device("meta"):
160
- dummy_linear = torch.nn.Linear(param.shape[1],
161
- param.shape[0],
162
- bias=False)
160
+ # linear can't be top level module since quantize_ is inplace
161
+ # while some of our configs need to do module swap, and only non-top
162
+ # level modules support module swap
163
+ dummy_linear = torch.nn.Sequential(
164
+ torch.nn.Linear(param.shape[1], param.shape[0], bias=False))
165
164
- dummy_linear.weight = param
166
+ dummy_linear[0].weight = param
167
quantize_(dummy_linear, torchao_config)
166
- return dummy_linear.weight
168
+ return dummy_linear[0].weight
169
170
171
class TorchAOLinearMethod(LinearMethodBase):