Use safetensors by default for `PyTorchModelHubMixin` #2033

bmuskalla · 2024-02-19T22:01:14Z

Switches PyTorchModelHubMixin to use safetensors.

To do discussed:

Should PyTorchModelHubMixin stay backward compatible and keep reading the pickle format via _from_pretrained? Not being sure if the model we're loading is safe is certainly a concern. Do you think the old way should continue working but issue a warning?

Fixes #1989

Wauplin

Hi @bmuskalla thanks for your contribution! 🔥 That will help starting the discussion :) So in my opinion we should:

when saving: start to save new files as .safetensors
when loading:
1. check if the local folder or remote repository contains a .safetensors file => if it's the case, load it
2. check if the local folder or remote repository contains a pytorch_model.bin file => if it's the case, load it
3. otherwise => raise exception

I would not update directly the PYTORCH_WEIGHTS_NAME constant since other libraries/users might use it in their workflow. What you can do is create a new PYTORCH_SAFE_WEIGHTS_NAME constant as done in transformers.

When transformers did the change, they introduced a new parameter safe_serialization: bool, first set to False (with a warning?) and then a few releases after set to True by default. Goal being to make the transition as smooth as possible.
EDIT: given the lower usage of this class (compared to transformers), we can skip the safe_serialization: bool parameter (e.g. no need to add it, let's make safetensors the default).

Pinging @LysandreJik who handled this process in transformers if I remember correctly (or at least kept an eye on it 🤗).

Wauplin · 2024-02-20T09:16:48Z

src/huggingface_hub/hub_mixin.py

+from safetensors import safe_open
+from safetensors.torch import save_file


This should not be in base imports as they should be optional for the users. What you should do is import them only in a if is_safetensors_available() statement below, as done for torch (see L14). Since huggingface_hub is a collection of many helpers used in various situations, we want to limit the number of required dependencies.

Thanks for the feedback, done in 859e230

julien-c

thanks for working on this @bmuskalla!

src/huggingface_hub/constants.py

tests/test_hubmixin.py

julien-c · 2024-02-20T09:20:01Z

What you can do is create a new PYTORCH_SAFE_WEIGHTS_NAME constant as done in transformers

or drop the PYTORCH entirely as those weights are not pytorch-specific

bmuskalla · 2024-02-23T09:27:33Z

check if the local folder or remote repository contains a .safetensors file => if it's the case, load it

check if the local folder or remote repository contains a pytorch_model.bin file => if it's the case, load it

otherwise => raise exception

@Wauplin Good call, I've implemented the fallback for now. We can still look into whether we should issue a warning later down the road.

thanks for working on this
drop the PYTORCH entirely as those weights are not pytorch-specific

@julien-c My pleasure. I've reused the constants for safetensors that were already present, no PYTORCH in the name anymore.

Wauplin

Thanks for iterating on this PR @bmuskalla! Looks good to me logic-wise. Thanks for taking care of the tests as well. Left a few comments mainly for styling matters but other than that we should be close to merging it :)

Wauplin · 2024-02-23T11:00:57Z

setup.py

-extras["torch"] = [
-    "torch",
-]
+extras["torch"] = ["torch", "safetensors"]


Suggested change

extras["torch"] = ["torch", "safetensors"]

extras["torch"] = [

"safetensors",

"torch",

]

(nit)

I'd prefer that as well but make style puts it on a single line ;)

You need to add a trailing comma "," to the last line otherwise ruff will fold it indeed. Made the change in 8ca8550.

src/huggingface_hub/utils/_runtime.py

tests/test_hub_mixin_pytorch.py

src/huggingface_hub/hub_mixin.py

Wauplin · 2024-02-23T11:13:52Z

tests/test_hub_mixin_pytorch.py

+        DummyModel().save_pretrained(self.cache_dir, config=TOKEN)
+        return self.cache_dir / "model.safetensors"
+
+    @patch.object(DummyModel, "_hf_hub_download")


Suggested change

@patch.object(DummyModel, "_hf_hub_download")

@patch("huggingface_hub.hf_hub_download")

Mocking like this should work and avoid the alias

Ha, didn't get this to work with the existing imports. My lack of python mock experience is exposed ;) Switching to import huggingface_hub makes it work. If you can enlighten me if there is a way to use @patch with a fqn while using from .file_download import hf_hub_download, more than happy to update the PR.

No worries! :D I've pushed a commit (abd0493) to make the @patch work with from .file_download import hf_hub_download. No big deal anyway, I just prefer this syntax for consistency with the rest of the codebase.

Today I learned - thanks for taking care of that

tests/test_hub_mixin_pytorch.py

Wauplin

Great! Thanks for the PR @bmuskalla! Hope that's fine with you but I've pushed 2 commits to arrange the last comments (see above). We should now be good to merge the PR as soon as the CI is green. So that it'll be shipped in the coming release! 🚀

EDIT: looks like we have issues in the CI but unrelated with this PR 😞

HuggingFaceDocBuilderDev · 2024-02-26T11:58:16Z

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.

Wauplin · 2024-02-26T12:42:55Z

Failing tests are unrelated so I'm merging this PR. Very nice contribution here 🔥

NielsRogge · 2024-03-01T10:47:18Z

Thanks a lot for your contribution! 🙌

FYI there was a comment by Meta authors as PyTorch was used previously: facebookresearch/hiera#26 (comment). Perhaps we can add an explicit flag to allow users to still do this, as in the Transformers library: https://github.com/huggingface/transformers/blob/15f8296a9b493eaa0770557fe2e931677fb62e2f/src/transformers/modeling_utils.py#L2182. It can then default to True.

Wauplin · 2024-03-01T13:35:45Z

@NielsRogge I don't think it's worth adding back support for saving to pytorch_model.bin. Saving to safetensors should be the default anyway in the future. For the record, the newest version always save as safetensors but can load from both .bin and .safetensors files. The only problem is if the model owner uploaded with the newest version of huggingface_hub -e.g. as safetensors- but the users still have an outdated version of huggingface_hub -e.g. not loading from safetensors-. I think it's fine given only newly uploaded models are impacted.

NielsRogge · 2024-03-01T15:04:18Z

Ok @Wauplin, sounds good to me!

bmuskalla added 2 commits February 19, 2024 12:44

Use safetensors in PyTorchModelHubMixin

228115f

Merge branch 'main' into 1989_safePyTorchMixin

5b4bef8

bmuskalla mentioned this pull request Feb 19, 2024

Use safetensors by default for PyTorchModelHubMixin class #1989

Closed

Wauplin reviewed Feb 20, 2024

View reviewed changes

julien-c reviewed Feb 20, 2024

View reviewed changes

src/huggingface_hub/constants.py Outdated Show resolved Hide resolved

tests/test_hubmixin.py Outdated Show resolved Hide resolved

bmuskalla added 7 commits February 22, 2024 13:03

Use safetensors constant

81f0647

Fix default download location and add test

8d4c050

Fallback to pickle model

9beca2c

make style

569f74e

Load safetensors dynamically

859e230

Merge branch 'main' into 1989_safePyTorchMixin

122aa5c

Migrate to new API

9f17725

Wauplin reviewed Feb 23, 2024

View reviewed changes

bmuskalla and others added 5 commits February 23, 2024 13:04

Add types

9fa71dd

improve assertion for safetensor header

8182031

Use @patch without the delegate

08152c5

make style

8ca8550

patch hf_hub_download

abd0493

Wauplin approved these changes Feb 26, 2024

View reviewed changes

Wauplin merged commit 46b38c2 into huggingface:main Feb 26, 2024
11 of 14 checks passed

NielsRogge mentioned this pull request Mar 11, 2024

KeyError: torch.complex64 when attempting to save PyTorch model huggingface/safetensors#450

Closed

2 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Use safetensors by default for `PyTorchModelHubMixin` #2033

Use safetensors by default for `PyTorchModelHubMixin` #2033

bmuskalla commented Feb 19, 2024 •

edited

Wauplin left a comment •

edited

Wauplin Feb 20, 2024

bmuskalla Feb 23, 2024

julien-c left a comment

julien-c commented Feb 20, 2024 •

edited

bmuskalla commented Feb 23, 2024

Wauplin left a comment

Wauplin Feb 23, 2024

bmuskalla Feb 23, 2024

Wauplin Feb 26, 2024

Wauplin Feb 23, 2024

bmuskalla Feb 23, 2024

Wauplin Feb 26, 2024

bmuskalla Feb 26, 2024

Wauplin left a comment •

edited

HuggingFaceDocBuilderDev commented Feb 26, 2024

Wauplin commented Feb 26, 2024

NielsRogge commented Mar 1, 2024

Wauplin commented Mar 1, 2024

NielsRogge commented Mar 1, 2024

		from safetensors import safe_open
		from safetensors.torch import save_file

	@patch.object(DummyModel, "_hf_hub_download")
	@patch("huggingface_hub.hf_hub_download")

Use safetensors by default for PyTorchModelHubMixin #2033

Use safetensors by default for PyTorchModelHubMixin #2033

Conversation

bmuskalla commented Feb 19, 2024 • edited

Wauplin left a comment • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

julien-c left a comment

Choose a reason for hiding this comment

julien-c commented Feb 20, 2024 • edited

bmuskalla commented Feb 23, 2024

Wauplin left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Wauplin left a comment • edited

Choose a reason for hiding this comment

HuggingFaceDocBuilderDev commented Feb 26, 2024

Wauplin commented Feb 26, 2024

NielsRogge commented Mar 1, 2024

Wauplin commented Mar 1, 2024

NielsRogge commented Mar 1, 2024

Use safetensors by default for `PyTorchModelHubMixin` #2033

Use safetensors by default for `PyTorchModelHubMixin` #2033

bmuskalla commented Feb 19, 2024 •

edited

Wauplin left a comment •

edited

julien-c commented Feb 20, 2024 •

edited

Wauplin left a comment •

edited