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] Add support for ModuleFqnToConfig using regex (#26001)
Signed-off-by: Jerry Zhang <jerryzh168@gmail.com>
Jerry Zhang committed
Oct 9, 2025 at 01:32 UTC
a83ff278d6e256baa5d89689dc95b39921a551d2
2 files changed
+38
-3
tests/quantization/test_torchao.py
+17
@@ -233,5 +233,22 @@ def test_opt_125m_float8_weight_only_safetensors_model_loading_with_params(vllm_
233
assert output
234
235
236
+@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
237
+@pytest.mark.skip(
238
+ reason="since torchao nightly is only compatible with torch nightly"
239
+ "currently https://github.com/pytorch/ao/issues/2919, we'll have to skip "
240
+ "torchao tests that requires newer versions (0.14.0.dev+) for now"
241
+)
242
+def test_opt_125m_module_fqn_to_config_regex_model(vllm_runner):
243
+ torch._dynamo.reset()
244
+ model_name = "torchao-testing/opt-125m-ModuleFqnToConfig-v1-regex-0.14.0.dev"
245
+ with vllm_runner(
246
+ model_name=model_name, dtype="bfloat16", pt_load_map_location="cuda:0"
247
+ ) as llm:
248
+ output = llm.generate_greedy(["The capital of France is"], max_tokens=32)
249
+
250
+ assert output
251
+
252
+
253
if __name__ == "__main__":
254
pytest.main([__file__])
vllm/model_executor/layers/quantization/torchao.py
+21
-3
@@ -5,6 +5,7 @@ import json
5
from importlib.util import find_spec
6
from typing import Any, Optional
7
8
+import regex as re
9
import torch
10
import torch.nn.functional as F
11
from packaging import version
@@ -192,9 +193,26 @@ class TorchAOConfig(QuantizationConfig):
193
module_fqn = prefix
194
if isinstance(self.torchao_config, ModuleFqnToConfig):
195
module_fqn_to_config = self.torchao_config.module_fqn_to_config
195
- c = module_fqn_to_config.get(module_fqn) or module_fqn_to_config.get(
196
- "_default", None
197
- )
196
+ c = None
197
+ if module_fqn in module_fqn_to_config:
198
+ assert not module_fqn.startswith("re:"), (
199
+ "module fqn should not start with"
200
+ "`re:`, which is used for specifying regex"
201
+ )
202
+ c = module_fqn_to_config[module_fqn]
203
+ else:
204
+ for maybe_module_fqn_pattern in module_fqn_to_config:
205
+ if not maybe_module_fqn_pattern.startswith("re:"):
206
+ continue
207
+ elif re.fullmatch(maybe_module_fqn_pattern[3:], module_fqn):
208
+ # we'll apply the config for first fully matched pattern
209
+ c = module_fqn_to_config[maybe_module_fqn_pattern]
210
+ break
211
+ else:
212
+ # fallback to use default if no module specific
213
+ # config is provided
214
+ c = module_fqn_to_config.get("_default", None)
215
+
216
if c is not None:
217
current_torchao_config = TorchAOConfig(
218
c, self.skip_modules, self.is_checkpoint_torchao_serialized