safetensors support #2881

Merged

liangel-02 merged 21 commits intomainfrom

Aug 29, 2025

Contributor

liangel-02 commented Aug 26, 2025

Context

Currently we serialize and distribute torchao quantized model with pytorch native APIs, specifically torch.save(model.state_dict())and model.load_state_dict(state_dict, assign=True)

Summary

This PR builds out the functionality save_tensor_subclass_dict and load_tensor_subclass_dict using save_file and load_file from safetensors library for FP8 based off of the script here.

Test Plan

Outputs match after saving/loading model.state_dict()

pytorch-bot bot commented Aug 26, 2025 •

edited

Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2881

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 1 Pending

As of commit c4e9165 with merge base 2f78cfe ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

meta-cla bot added the CLA Signed label

liangel-02 added module: not user facing device: rocm ciflow/rocm ciflow/4xh100 labels

liangel-02 requested a review from jerryzh168

August 26, 2025 18:11

liangel-02 marked this pull request as ready for review

August 26, 2025 18:11

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensor.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensor.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensor.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensors_support.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensors_support.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensors_support.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensors_support.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensors_support.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensors_support.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensors_support.py Outdated Show resolved Hide resolved

liangel-02 changed the title ~~safetensor support~~ safetensors support

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensors_support.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensors_support.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensors_support.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensors_support.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensors_support.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/QuantizeTensorToFloat8KwargsJSON.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/QuantizeTensorToFloat8KwargsJSON.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensors_support.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensors_JSONEncoders.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensors_JSONEncoders.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/quantization/safetensors_JSONEncoders.py Show resolved Hide resolved

liangel-02 added 15 commits

August 28, 2025 13:19


          remove binary files

7252a33


          delete debug

54cd56a


          comment

499da7e


          refactor

dfe0e13


          remove binary

a707c91


          ruff

75e420d


          addressing comments

756e918


          format

26b96a9


          remove

54e433e


          address comments

abce5eb


          clean up code

4d92717


          fix comments

4f91503


          format

6397ed6


          tempfile

a9d960a


          logging

9824c96

liangel-02 force-pushed the safetensor branch from 8d89521 to bb82dff Compare

August 28, 2025 20:20

jerryzh168 reviewed

View reviewed changes

torchao/prototype/safetensors/safetensors_serialization.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/safetensors/safetensors_serialization.py Outdated Show resolved Hide resolved

jerryzh168 requested a review from mikaylagawarecki

August 28, 2025 20:34

jerryzh168 reviewed

View reviewed changes

torchao/prototype/safetensors/safetensors_serialization.py

Comment on lines +9 to +11

+              from torchao.quantization import Float8Tensor
+              from torchao.quantization.quantize_.common import KernelPreference
+              from torchao.quantization.quantize_.workflows import QuantizeTensorToFloat8Kwargs

Contributor

jerryzh168 Aug 28, 2025

looks like some imports are not used?

Contributor Author

liangel-02 Aug 28, 2025

these are used to create the ALLOWED_CLASSES dict

liangel-02 force-pushed the safetensor branch from bb82dff to 6c78c00 Compare

August 28, 2025 21:28

mikaylagawarecki approved these changes

View reviewed changes

mikaylagawarecki left a comment

Thanks for restricting to ALLOWED_CLASSES

jerryzh168 reviewed

View reviewed changes

torchao/prototype/safetensors/safetensors_serialization.py Outdated Show resolved Hide resolved

jerryzh168 reviewed

View reviewed changes

torchao/prototype/safetensors/safetensors_serialization.py Show resolved Hide resolved


          addressing comments

c4e9165

liangel-02 force-pushed the safetensor branch from 6c78c00 to c4e9165 Compare

August 28, 2025 22:20

liangel-02 merged commit 4236656 into main

20 checks passed

liangel-02 deleted the safetensor branch

August 29, 2025 01:10

liangel-02 mentioned this pull request

[torchao safetensors] integrate torchao safetensors support with transformers huggingface/transformers#40735

Merged

ksivaman mentioned this pull request

[PyTorch] Fix tests for 🤗 integration NVIDIA/TransformerEngine#2239

Merged

13 tasks

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/rocm ciflow/4xh100 CLA Signed device: rocm module: not user facing