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 12, 22:15 UTC
[torchao] fix safetensors for sharding (#28169)
Signed-off-by: Angel Li <liangel@meta.com>
liangel-02 committed
Nov 19, 2025 at 19:39 UTC
1d642872a27f1c6bedf28669642928cc7eec6532
3 files changed
+23
-11
tests/quantization/test_torchao.py
+4
-5
@@ -225,13 +225,12 @@ def test_reload_weights():
225
@pytest.mark.skip(
226
reason="since torchao nightly is only compatible with torch nightly"
227
"currently https://github.com/pytorch/ao/issues/2919, we'll have to skip "
228
- "torchao tests that requires newer versions (0.14.0.dev+) for now"
228
+ "torchao tests that requires newer versions (0.15.0.dev+) for now"
229
)
230
-def test_opt_125m_float8_weight_only_safetensors_model_loading_with_params(vllm_runner):
230
+def test_safetensors_model_loading_with_params(vllm_runner):
231
torch._dynamo.reset()
232
- model_name = (
233
- "torchao-testing/opt-125m-Float8WeightOnlyConfig-v2-0.14.0.dev-safetensors"
234
- )
232
+ # using this model to test safetensors loading with file sharding
233
+ model_name = "torchao-testing/Qwen3-8B-INT4-0.15.0dev-safetensors"
234
with vllm_runner(model_name=model_name, dtype="bfloat16") as llm:
235
output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
236
vllm/model_executor/model_loader/default_loader.py
+1
-1
@@ -279,7 +279,7 @@ class DefaultModelLoader(BaseModelLoader):
279
if (
280
hasattr(quant_config, "is_checkpoint_torchao_serialized")
281
and quant_config.is_checkpoint_torchao_serialized
282
- and torchao_version_at_least("0.14.0")
282
+ and torchao_version_at_least("0.15.0")
283
):
284
self.load_config.safetensors_load_strategy = "torchao"
285
vllm/model_executor/model_loader/weight_utils.py
+18
-5
@@ -595,6 +595,9 @@ def safetensors_weights_iterator(
595
if safetensors_load_strategy == "eager":
596
loading_desc += " (eager)"
597
598
+ state_dict = {}
599
+ leftover_state_dict: dict[str, torch.Tensor] = {}
600
+
601
for st_file in tqdm(
602
hf_weights_files,
603
desc=loading_desc,
@@ -606,9 +609,11 @@ def safetensors_weights_iterator(
609
state_dict = load(f.read())
610
yield from state_dict.items()
611
elif safetensors_load_strategy == "torchao":
609
- if not torchao_version_at_least("0.14.0"):
612
+ # we can't load flattened torchao tensor subclasses directly into the model
613
+ # instead we reconstruct the subclasses here before returning
614
+ if not torchao_version_at_least("0.15.0"):
615
raise ValueError(
611
- "Please use torchao version >= 0.14.0 \
616
+ "Please use torchao version >= 0.15.0 \
617
to load torchao safetensors checkpoint"
618
)
619
from torchao.prototype.safetensors.safetensors_support import (
@@ -616,12 +621,20 @@ def safetensors_weights_iterator(
621
)
622
623
with safe_open(st_file, framework="pt") as f:
619
- state_dict = {}
624
for name in f.keys(): # noqa: SIM118
625
state_dict[name] = f.get_tensor(name)
626
+
627
+ # update with leftover tensor data from previous iteration, if any
628
+ state_dict.update(leftover_state_dict)
629
metadata = f.metadata()
623
- updated_state_dict = unflatten_tensor_state_dict(state_dict, metadata)
624
- yield from updated_state_dict.items()
630
+ # due to sharded checkpoints, we are not guaranteed that we have all
631
+ # tensor subclass data on one file
632
+ # state_dict has the leftover data from this step and we wait for
633
+ # missing information to be provided in a future iteration
634
+ unflattened_state_dict, leftover_state_dict = (
635
+ unflatten_tensor_state_dict(state_dict, metadata)
636
+ )
637
+ yield from unflattened_state_dict.items()
638
else:
639
with safe_open(st_file, framework="pt") as f:
640
for name in f.keys(): # noqa: SIM118