Skip to content

Add safetensors export feature #3345

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.

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

Open
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

jonboh
Copy link
Contributor

@jonboh jonboh commented Jul 3, 2025

This is still a work in progress, I started out following antimora's indications in: #3260 (comment)

I did not find a clean way to walk the Module with the Serializer, and perform the serialization only with the Recorder. The problem is that once the Serializer is finished I'd need to reconstruct each TensorData back from its bytes, dtype and shape fields as the serializer will walk up to the basic types.

The alternative I found was to use the Serializer to get a mapping from ParamId to the TensorName, and use a ModuleVisitor to link to the TensorData. I think this is a bit cleaner, but the API is pretty different from the rest of the serializations, so I'm not sure, if it would be better to do it the other way, even if the logic ends up a bit more convoluted.
Any comments are appreciated :)

@@ -49,6 +49,7 @@ thiserror = { workspace = true, optional = true }
tracing-core = { workspace = true }
tracing-subscriber = { workspace = true }
zip = { workspace = true, optional = true }
safetensors = { workspace = true }
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should be optional = true and enabled via feature flag.

@antimora
Copy link
Collaborator

antimora commented Jul 3, 2025

I think I know what the issue you are facing is regarding the serializer going all the way down. It's okay to stop if you come across a TensorData struct type and handle serialization differently.

My recommendation:

  1. Don't use ModuleVisitor (you can get away with a serde serializer).
  2. The only intermediate data structure should be a hash type of full path name + tensor safetensor view.

@antimora antimora changed the title add safetensors export feature Add safetensors export feature Jul 12, 2025
@jonboh jonboh marked this pull request as ready for review July 13, 2025 12:01
@jonboh jonboh requested a review from antimora July 13, 2025 12:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants