+

[safetensors enablement] refactoring for huggingface integration #2936

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Merged

liangel-02 merged 1 commit into main from safetensors_refactoring

Sep 5, 2025

Contributor

liangel-02 commented Sep 4, 2025 •

edited

Loading

Summary

save_tensor_state_dict and load_tensor_state_dict was previously responsible for directly calling safetensors.save_file and savetensors.load_file.

we've now refactored that functionality so that these two functions (renamed to unflatten_tensor_state_dict and flatten_tensor_state_dict respectively) only flattens/unflattens tensor subclass state dict so that it is compatible with huggingface interface.

Test Plan
python test/prototype/safetensors/test_safetensors_support.py

updated test file to include utility function that calls savetensors.load_file and handles specific loading logic.

pytorch-bot bot commented Sep 4, 2025 •

edited

Loading

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 8708caf with merge base f35ae41 ():
💚 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 requested a review from jerryzh168

September 4, 2025 00:10

liangel-02 added topic: improvement module: rocm ciflow/rocm ciflow/4xh100 labels

liangel-02 marked this pull request as ready for review

September 4, 2025 00:11

jerryzh168 reviewed

View reviewed changes

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

jerryzh168 reviewed

View reviewed changes

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

jerryzh168 reviewed

View reviewed changes

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

jerryzh168 reviewed

View reviewed changes

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

liangel-02 force-pushed the safetensors_refactoring branch 3 times, most recently from cb0854f to 8d6f745 Compare

September 4, 2025 00:51

jerryzh168 reviewed

View reviewed changes

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

jerryzh168 reviewed

View reviewed changes

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

jerryzh168 reviewed

View reviewed changes

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

liangel-02 force-pushed the safetensors_refactoring branch 2 times, most recently from e3a6d23 to afb5270 Compare

September 4, 2025 15:16

jerryzh168 reviewed

View reviewed changes

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

jerryzh168 reviewed

View reviewed changes

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

jerryzh168 reviewed

View reviewed changes

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

jerryzh168 reviewed

View reviewed changes

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

jerryzh168 reviewed

View reviewed changes

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

liangel-02 force-pushed the safetensors_refactoring branch from afb5270 to bdab1d2 Compare

September 4, 2025 17:36

liangel-02 changed the title ~~refactoring functions for huggingface integration~~ [safetensors enablement] refactoring for huggingface integration

jerryzh168 reviewed

View reviewed changes

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

jerryzh168 approved these changes

View reviewed changes

liangel-02 force-pushed the safetensors_refactoring branch from bdab1d2 to 60bb577 Compare

September 4, 2025 17:56

liangel-02 force-pushed the safetensors_refactoring branch from 60bb577 to f4951c5 Compare

September 4, 2025 18:01

jerryzh168 reviewed

View reviewed changes

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


          refactoring functions for huggingface integration

8708caf

liangel-02 force-pushed the safetensors_refactoring branch from f4951c5 to 8708caf Compare

September 4, 2025 18:48

liangel-02 merged commit 7b81460 into main

20 of 21 checks passed

liangel-02 deleted the safetensors_refactoring branch

September 5, 2025 13:57

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 module: rocm topic: improvement

点击这是indexloc提供的php浏览器服务，不要输入任何密码和下载