[torchao safetensors] integrate torchao safetensors support with transformers #40735
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
liangel-02
force-pushed
the
torchao_safetensors
branch
2 times, most recently
from
d60acfe to
392a504
Compare
jerryzh168
reviewed
Sep 5, 2025
jerryzh168
reviewed
Sep 5, 2025
jerryzh168
reviewed
Sep 5, 2025
jerryzh168
reviewed
Sep 5, 2025
jerryzh168
reviewed
Sep 5, 2025
jerryzh168
reviewed
Sep 5, 2025
jerryzh168
reviewed
Sep 5, 2025
|
cc @MekkCyber for quantization |
liangel-02
force-pushed
the
torchao_safetensors
branch
7 times, most recently
from
7246421 to
6a26d01
Compare
liangel-02
force-pushed
the
torchao_safetensors
branch
from
6a26d01 to
40cb596
Compare
|
[For maintainers] Suggested jobs to run (before merge) run-slow: torchao_integration |
jerryzh168
reviewed
Sep 8, 2025
| @@ -399,7 +401,7 @@ def test_autoquant(self): | |||
|
|
|||
| check_autoquantized(self, quantized_model.model.layers[0].self_attn.v_proj) | |||
|
|
|||
| EXPECTED_OUTPUT = "What are we having for dinner?\n\nJane: (sighs)" | |||
| EXPECTED_OUTPUT = "What are we having for dinner?\n\nJessica: (smiling)" | |||
There was a problem hiding this comment.
should this be reverted?
There was a problem hiding this comment.
i double checked that this fails on main as well, so i just added the correction
jerryzh168
reviewed
Sep 8, 2025
|
looks good, please update PR summary to align with recent code changes as well cc @SunMarc @MekkCyber please check if the API changes make sense |
jerryzh168
reviewed
Sep 8, 2025
|
Would be very nice to have this propagated to |
SunMarc
reviewed
Sep 9, 2025
Comment on lines
+4023
to
+4027
| if hf_quantizer.quantization_config.quant_method is QuantizationMethod.TORCHAO: | ||
| state_dict, metadata = hf_quantizer.get_state_dict_and_metadata(self, safe_serialization) | ||
| else: | ||
| state_dict = hf_quantizer.get_state_dict(self) | ||
| metadata["format"] = "pt" |
There was a problem hiding this comment.
<aybe we remove get_state_dict completely and only keep get_state_dict_and_metadata. I think this will be more clearer.
SunMarc
reviewed
Sep 9, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Context
Currently, we need to use
safe_serialization=Falsewhile saving models as shown here. This PR enables safetensors support for torchao so that users can now save and load checkpoints using safetensors. Currently, only Float8Tensor is supported (Float8DynamicActivationFloat8WeightConfig,Float8WeightOnlyConfig) but allowing other subclasses should involve minimal code changes.Summary
Changes to transformers code includes:
TorchAoHfQuantizer, we provideget_state_dictandupdate_state_dict_with_metadatathat flattens/unflattens a model state dict with tensor subclasses by calling functionality built out in this PR.modeling_utils.py, we make appropriate changes to support propagating the metadata from tensor subclasses. We also add logic similar tohqqandbnbto directly load ontocpurather thanmeta.Test Plan
Modified unit test to allow safe serialization. Run using
python tests/quantization/torchao_integration/test_torchao.pyReference https://huggingface.co/torchao-testing/opt-125m-Float8WeightOnlyConfig-v2-0.14.0.dev-safetensors for an example of a serialized model and test script