Handle shared layers in save_torch_state_dict + add save_torch_model
#2373
Conversation
This file contains 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
Wauplin
requested review from
Narsil,
muellerzr,
amyeroberts,
sayakpaul,
LysandreJik and
SunMarc
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
sayakpaul
reviewed
Jul 5, 2024
| The main helper of the `serialization` module takes a state dictionary as input (e.g. a mapping between layer names and related tensors), splits it into several shards while creating a proper index in the process and save everything to disk. At the moment, only `torch` tensors are supported. Under the hood, it delegates the logic to split the state dictionary to [`split_torch_state_dict_into_shards`]. | ||
| The main helper of the `serialization` module takes a torch `nn.Module` as input and saves it to disk. It handles the logic to save shared tensors (see [safetensors explanation](https://huggingface.co/docs/safetensors/torch_shared_tensors)) as well as logic to split the state dictionary into shards, using [`split_torch_state_dict_into_shards`] under the hood. At the moment, only `torch` framework is supported. | ||
|
|
||
| If you want to save a state dictionary (e.g. a mapping between layer names and related tensors) instead of a `nn.Module`, you can use [`save_torch_state_dict`] which provides the same features. This is useful for example if you want to apply custom logic to the state dict before saving it. |
There was a problem hiding this comment.
I see the point of mentioning this but also I think for the Torch community, it's fairly standard practice to ship the model classes and their state dictionaries (i.e., the parameters) separately unlike TensorFlow/Keras, for example.
sayakpaul
reviewed
Jul 5, 2024
| ) | ||
|
|
||
|
|
||
| def get_tensor_size(tensor: "tf.Tensor") -> int: | ||
| def get_tf_storage_size(tensor: "tf.Tensor") -> int: |
There was a problem hiding this comment.
Do we have all the equivalent torch methods for TensorFlow? Or is that not necessary?
There was a problem hiding this comment.
Not yet no. Let's build for torch first and then expand to TF after if needed. For now for TF we have the logic to split a state dict into shards but nothing to save to disk.
sayakpaul
reviewed
Jul 5, 2024
Comment on lines
+249
to
+251
| "metadata": {**state_dict_split.metadata, **metadata}, | ||
| "weight_map": state_dict_split.tensor_to_filename, | ||
| } |
There was a problem hiding this comment.
Should there be any sanity check on the additional metadata if not already done?
There was a problem hiding this comment.
metadata is at the discretion of the frameworks that will use it (transformers/diffusers/accelerate). In practice, I don't think it'll be much used. In any case, we can't really do sanity check since we are supposed to accept anything that is jsonable.
sayakpaul
reviewed
Jul 5, 2024
LysandreJik
approved these changes
Jul 11, 2024
|
Thanks for the reviews! Let's ship it yes 😄 |
Wauplin
added a commit
that referenced
this pull request
Jul 11, 2024
* Use extended path on Windows when downloading to local dir Change the path of the local dir to an extended path by prepending "\\?\" to the absolute path, when the absolute path is longer than 255 characters on Windows. Also fixed a small typo. * Use extended path on Windows when downloading to local dir Change the path of the local dir to an extended path by prepending "\\?\" to the absolute path, when the absolute path is longer than 255 characters on Windows. Also fixed a small typo. * Move path handling to `get_local_download_paths()` for robustness On Windows we check the length of `lock_path` and if it is longer than 255 characters we prepend the `\\?\` prefix to all paths if it does not already exist. We only need to check the length of `lock_path` because it is guaranteed to be the longest path. * `safetensors[torch]` (#2371) * Fix token=False not respected in file download (#2386) * Fix token=False not respected in file download * lint * Handle shared layers in `save_torch_state_dict` + add `save_torch_model` (#2373) * Handle shared layers in save_torch_state_dict + save_torch_model + some helpers * fix pytest rerun * more reruns * Support `expand` parameter in `xxx_info` and `list_xxxs` (model/dataset/Space) (#2333) * First draft to support `expand` parameter for models * add expand support for dataset * add expand support for Space * Use extended path on Windows when downloading to local dir Change the path of the local dir to an extended path by prepending "\\?\" to the absolute path, when the absolute path is longer than 255 characters on Windows. Also fixed a small typo. * Move path handling to `get_local_download_paths()` for robustness On Windows we check the length of `lock_path` and if it is longer than 255 characters we prepend the `\\?\` prefix to all paths if it does not already exist. We only need to check the length of `lock_path` because it is guaranteed to be the longest path. * Use extended path on Windows when downloading to local dir Change the path of the local dir to an extended path by prepending "\\?\" to the absolute path, when the absolute path is longer than 255 characters on Windows. Also fixed a small typo. * Removed old path handling * Reorder path check; add tests * Skip test if opn Windows The test now shows up a `skipped` if executed on a non-Windows machine Co-authored-by: Lucain <[email protected]> * Fix indentation for test_local_folder.py * Fix code style --------- Co-authored-by: Lucain <[email protected]> Co-authored-by: Quentin Gallouédec <[email protected]> Co-authored-by: Lucain <[email protected]>
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.
Partially resolve #2065.
Follow-up PR after #2314.
In #2314, we introduce
save_torch_state_dict. This new PR:save_torch_modelto directly save a torchnn.Moduleget_tf_storage_size/get_torch_storage_sizeand make them public + documentedA last follow-up PR should had
load_torch_state_dict/load_torch_modelhelpers as well to correctly reload those files, including the shared layers.I'm pinging transformers/accelerate/diffusers cores maintainers for visibility as well. Feel free to comment if someone should be done differently.