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] safetensors integration (#25969)
Signed-off-by: Angel Li <liangel@meta.com>
liangel-02 committed
Oct 7, 2025 at 19:12 UTC
b32260ab853d22069ad0e002aeae00218e086701
5 files changed
+60
tests/quantization/test_torchao.py
+17
@@ -216,5 +216,22 @@ def test_reload_weights():
216
# print("-" * 60)
217
218
219
+@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
220
+@pytest.mark.skip(
221
+ reason="since torchao nightly is only compatible with torch nightly"
222
+ "currently https://github.com/pytorch/ao/issues/2919, we'll have to skip "
223
+ "torchao tests that requires newer versions (0.14.0.dev+) for now"
224
+)
225
+def test_opt_125m_float8_weight_only_safetensors_model_loading_with_params(vllm_runner):
226
+ torch._dynamo.reset()
227
+ model_name = (
228
+ "torchao-testing/opt-125m-Float8WeightOnlyConfig-v2-0.14.0.dev-safetensors"
229
+ )
230
+ with vllm_runner(model_name=model_name, dtype="bfloat16") as llm:
231
+ output = llm.generate_greedy(["The capital of France is"], max_tokens=32)
232
+
233
+ assert output
234
+
235
+
236
if __name__ == "__main__":
237
pytest.main([__file__])
vllm/config/load.py
+4
@@ -59,6 +59,10 @@ class LoadConfig:
59
This is recommended for models on network filesystems (e.g., Lustre, NFS)
60
as it avoids inefficient random reads, significantly speeding up model
61
initialization. However, it uses more CPU RAM.
62
+ - "torchao": Weights are loaded in upfront and then reconstructed
63
+ into torchao tensor subclasses. This is used when the checkpoint
64
+ was quantized using torchao and saved using safetensors.
65
+ Needs torchao >= 0.14.0
66
"""
67
model_loader_extra_config: Union[dict, TensorizerConfig] = field(
68
default_factory=dict
vllm/model_executor/layers/quantization/torchao.py
+15
@@ -1,10 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import importlib
4
import json
5
+from importlib.util import find_spec
6
from typing import Any, Optional
7
8
import torch
9
import torch.nn.functional as F
10
+from packaging import version
11
from torch.nn.parameter import Parameter
12
13
from vllm.logger import init_logger
@@ -23,6 +26,18 @@ from vllm.model_executor.utils import set_weight_attrs
26
logger = init_logger(__name__)
27
28
29
+def torchao_version_at_least(torchao_version: str) -> bool:
30
+ if find_spec("torchao"):
31
+ try:
32
+ if version.parse(importlib.metadata.version("torchao")) >= version.parse(
33
+ torchao_version
34
+ ):
35
+ return True
36
+ except (ImportError, version.InvalidVersion):
37
+ return False
38
+ return False
39
+
40
+
41
def should_skip(prefix: str, skip_modules: list[str]) -> bool:
42
"""
43
Robust skipping logic:
vllm/model_executor/model_loader/default_loader.py
+5
@@ -14,6 +14,7 @@ from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
14
from vllm.config import ModelConfig
15
from vllm.config.load import LoadConfig
16
from vllm.logger import init_logger
17
+from vllm.model_executor.layers.quantization.torchao import torchao_version_at_least
18
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
19
from vllm.model_executor.model_loader.weight_utils import (
20
download_safetensors_index_file_from_hf,
@@ -272,6 +273,10 @@ class DefaultModelLoader(BaseModelLoader):
273
)
274
275
def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
276
+ if model_config.quantization == "torchao" and torchao_version_at_least(
277
+ "0.14.0"
278
+ ):
279
+ self.load_config.safetensors_load_strategy = "torchao"
280
weights_to_load = {name for name, _ in model.named_parameters()}
281
282
# if we don't have `model.weight_metadata_and_attr_saved` defined and
vllm/model_executor/model_loader/weight_utils.py
+19
@@ -54,6 +54,8 @@ except ImportError:
54
SafeTensorsFileLoader = fastsafetensors.placeholder_attr("SafeTensorsFileLoader")
55
SingleGroup = fastsafetensors.placeholder_attr("SingleGroup")
56
57
+from vllm.model_executor.layers.quantization.torchao import torchao_version_at_least
58
+
59
logger = init_logger(__name__)
60
61
# use system-level temp directory for file locks, so that multiple users
@@ -602,6 +604,23 @@ def safetensors_weights_iterator(
604
with open(st_file, "rb") as f:
605
state_dict = load(f.read())
606
yield from state_dict.items()
607
+ elif safetensors_load_strategy == "torchao":
608
+ if not torchao_version_at_least("0.14.0"):
609
+ raise ValueError(
610
+ "Please use torchao version >= 0.14.0 \
611
+ to load torchao safetensors checkpoint"
612
+ )
613
+ from torchao.prototype.safetensors.safetensors_support import (
614
+ unflatten_tensor_state_dict,
615
+ )
616
+
617
+ with safe_open(st_file, framework="pt") as f:
618
+ state_dict = {}
619
+ for name in f.keys(): # noqa: SIM118
620
+ state_dict[name] = f.get_tensor(name)
621
+ metadata = f.metadata()
622
+ updated_state_dict = unflatten_tensor_state_dict(state_dict, metadata)
623
+ yield from updated_state_dict.items()
624
else:
625
with safe_open(st_file, framework="pt") as f:
626
for name in f.keys(): # noqa: SIM118