Skip to content

Adding use_safetensors argument to give more control to users #2123

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

Merged
merged 12 commits into from
Mar 16, 2023
Merged

Adding use_safetensors argument to give more control to users #2123

merged 12 commits into from
Mar 16, 2023

Conversation

Narsil
Copy link
Contributor

@Narsil Narsil commented Jan 26, 2023

about which weights they use.

@HuggingFaceDocBuilderDev
Copy link

HuggingFaceDocBuilderDev commented Jan 26, 2023

The documentation is not available anymore as the PR was closed or merged.

@patrickvonplaten
Copy link
Contributor

I understand the need for the PR, but I wonder whether we could use the soon-to-be-added "variant" kwarg for this instead to not add too many kwargs.

The idea with "variant" is summarized here: #1764

I think it might be better to handle all of this with the "variant" kwarg instead. E.g. from_pretrained(variant="safetensors") then safetensors has to be installed and the files have to be there or an error is thrown.

We need to also align this with transformers cc @sgugger @LysandreJik .

@sgugger
Copy link
Contributor

sgugger commented Jan 26, 2023

For Transformers, safetensors is the default if available, there is no special flag for it and we are not planning on adding one. We just look at model.safetensors before anything else.

@@ -350,6 +350,11 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
also tries to not use more than 1x model size in CPU memory (including peak memory) while loading the
model. This is only supported when torch version >= 1.9.0. If you are using an older version of torch,
setting this argument to `True` will raise an error.
use_safetensors (`bool`, *optional*, defaults to `None`):
Copy link
Contributor

Choose a reason for hiding this comment

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

Suggested change
use_safetensors (`bool`, *optional*, defaults to `None`):
use_safetensors (`bool`, *optional*):

Copy link
Contributor

@patrickvonplaten patrickvonplaten left a comment

Choose a reason for hiding this comment

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

Actually upon section reflection, I think this makes sense.

Can we also add tests though?
And also would maybe save_loading suit better (since we have safe_serialization ?)

@@ -350,6 +350,11 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
also tries to not use more than 1x model size in CPU memory (including peak memory) while loading the
model. This is only supported when torch version >= 1.9.0. If you are using an older version of torch,
setting this argument to `True` will raise an error.
use_safetensors (`bool`, *optional*, defaults to `None`):
If set to `True`, the pipeline will forcibly load the models using `safetensors` weights. If set to
Copy link
Contributor

Choose a reason for hiding this comment

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

Suggested change
If set to `True`, the pipeline will forcibly load the models using `safetensors` weights. If set to
If set to `True`, the pipeline will forcibly load the model from `safetensors` weights. If set to

@@ -350,6 +350,11 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
also tries to not use more than 1x model size in CPU memory (including peak memory) while loading the
model. This is only supported when torch version >= 1.9.0. If you are using an older version of torch,
setting this argument to `True` will raise an error.
use_safetensors (`bool`, *optional*, defaults to `None`):
If set to `True`, the pipeline will forcibly load the models using `safetensors` weights. If set to
`None` (the default). The pipeline will load using `safetensors` if the safetensors weights are
Copy link
Contributor

Choose a reason for hiding this comment

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

Suggested change
`None` (the default). The pipeline will load using `safetensors` if the safetensors weights are
`None` (the default). The pipeline will load using `safetensors` if safetensors weights are

use_safetensors (`bool`, *optional*, defaults to `None`):
If set to `True`, the pipeline will forcibly load the models using `safetensors` weights. If set to
`None` (the default). The pipeline will load using `safetensors` if the safetensors weights are
actually available *and* you have the library installed. If the to `False` the pipeline will *not* use
Copy link
Contributor

Choose a reason for hiding this comment

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

Suggested change
actually available *and* you have the library installed. If the to `False` the pipeline will *not* use
available *and* if `safetensors` is installed. If set to `False` the pipeline will *not* use

If set to `True`, the pipeline will forcibly load the models using `safetensors` weights. If set to
`None` (the default). The pipeline will load using `safetensors` if the safetensors weights are
actually available *and* you have the library installed. If the to `False` the pipeline will *not* use
`safetensors` at all.
Copy link
Contributor

Choose a reason for hiding this comment

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

Suggested change
`safetensors` at all.
`safetensors`.

@@ -400,6 +400,11 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
setting this argument to `True` will raise an error.
return_cached_folder (`bool`, *optional*, defaults to `False`):
If set to `True`, path to downloaded cached folder will be returned in addition to loaded pipeline.
use_safetensors (`bool`, *optional*, defaults to `None`):
Copy link
Contributor

Choose a reason for hiding this comment

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

Suggested change
use_safetensors (`bool`, *optional*, defaults to `None`):
use_safetensors (`bool`, *optional*):

@@ -400,6 +400,11 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
setting this argument to `True` will raise an error.
return_cached_folder (`bool`, *optional*, defaults to `False`):
If set to `True`, path to downloaded cached folder will be returned in addition to loaded pipeline.
use_safetensors (`bool`, *optional*, defaults to `None`):
If set to `True`, the pipeline will forcibly load the models using `safetensors` weights. If set to
`None` (the default). The pipeline will load using `safetensors` if the safetensors weights are
Copy link
Contributor

Choose a reason for hiding this comment

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

Suggested change
`None` (the default). The pipeline will load using `safetensors` if the safetensors weights are
`None` (the default), the pipeline will be loaded from `safetensors` if the safetensors weights are

use_safetensors (`bool`, *optional*, defaults to `None`):
If set to `True`, the pipeline will forcibly load the models using `safetensors` weights. If set to
`None` (the default). The pipeline will load using `safetensors` if the safetensors weights are
actually available *and* you have the library installed. If the to `False` the pipeline will *not* use
Copy link
Contributor

Choose a reason for hiding this comment

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

Suggested change
actually available *and* you have the library installed. If the to `False` the pipeline will *not* use
available *and* if `safetensors` is installed. If the to `False` the pipeline will *not* use

If set to `True`, the pipeline will forcibly load the models using `safetensors` weights. If set to
`None` (the default). The pipeline will load using `safetensors` if the safetensors weights are
actually available *and* you have the library installed. If the to `False` the pipeline will *not* use
`safetensors` at all.
Copy link
Contributor

Choose a reason for hiding this comment

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

Suggested change
`safetensors` at all.
`safetensors`.

@@ -505,13 +511,13 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P

user_agent = http_user_agent(user_agent)

if is_safetensors_available():
if use_safetensors in {None, True}:
Copy link
Contributor

@patrickvonplaten patrickvonplaten Jan 27, 2023

Choose a reason for hiding this comment

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

Suggested change
if use_safetensors in {None, True}:
if use_safetensors is not False:

Copy link
Contributor

Choose a reason for hiding this comment

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

actually I would like to avoid the extra call to the Hub if safetensors are not installed. Can we make sure that we still check for whether safetensors are installed?

info = model_info(
pretrained_model_name_or_path,
use_auth_token=use_auth_token,
revision=revision,
)
if is_safetensors_compatible(info):
if use_safetensors is True or is_safetensors_compatible(info):
Copy link
Contributor

Choose a reason for hiding this comment

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

Suggested change
if use_safetensors is True or is_safetensors_compatible(info):
if use_safetensors or is_safetensors_compatible(info):

@github-actions
Copy link
Contributor

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

@github-actions github-actions bot added the stale Issues that haven't received updates label Feb 26, 2023
@Narsil
Copy link
Contributor Author

Narsil commented Feb 27, 2023

Unstale

@patrickvonplaten
Copy link
Contributor

@Narsil, I think we should merge such a PR soon, but the logic has been changed quite a bit since: #2305

Can we maybe rebase this one?

@@ -463,7 +469,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P

model = load_flax_checkpoint_in_pytorch_model(model, model_file)
else:
if is_safetensors_available():
if use_safetensors in {None, True}:
Copy link
Contributor

Choose a reason for hiding this comment

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

Suggested change
if use_safetensors in {None, True}:
if use_safetensors is not False:

@patrickvonplaten
Copy link
Contributor

Still very much interested in merging this soon, but before merging this we probably need the same functionality in transformers - otherwise it doesn't make too much sense.

@patrickvonplaten
Copy link
Contributor

huggingface/transformers#22083 is in - think we can continue with this PR :-) cc @Narsil

@Narsil
Copy link
Contributor Author

Narsil commented Mar 15, 2023

M1 failure not linked to this PR is it ?

@patrickvonplaten
Copy link
Contributor

Thanks a lot @Narsil - looks good to me!

Comment on lines +395 to +398
use_safetensors (`bool`, *optional* ):
If set to `True`, the pipeline will forcibly load the models from `safetensors` weights. If set to
`None` (the default). The pipeline will load using `safetensors` if safetensors weights are available
*and* if `safetensors` is installed. If the to `False` the pipeline will *not* use `safetensors`.
Copy link
Member

Choose a reason for hiding this comment

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

Cool!

Copy link
Member

@sayakpaul sayakpaul left a comment

Choose a reason for hiding this comment

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

Thanks!

Cool test cases.

Comment on lines 154 to 155
if (is_safetensors_available() and weight_name is None) or (
if (use_safetensors is not False and weight_name is None) or (
Copy link
Contributor

Choose a reason for hiding this comment

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

hmm, is not False is pretty confusing. With this and the above line,

        use_safetensors = kwargs.pop("use_safetensors", None if is_safetensors_available() else False)

I'm a bit confused about what the default behavior of the function is supposed to be

Copy link
Contributor

Choose a reason for hiding this comment

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

Ah ok I see the comment later on in the PR. Could we change the code (here and in the other places where we take use_safetensors as an argument with a default) to approximately

use_safetensors = kwargs.pop("use_safetensors", None)

if use_safetensors is None:
    use_safetensors = is_safetensors_available()

if use_safetensors and not is_safetensors_available():
  raise ValueError("`use_safetensors`=True but safetensors is not installed. Please install safetensors with `pip install safetenstors")

# in the rest of the function just use `use_safetensors` as a boolean

The rule of thumb here is that values which can be booleans or None can be hard to read/code with, especially when used in conditionals

Copy link
Contributor Author

Choose a reason for hiding this comment

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

But that's not it here.

If safetensors is installed, but the weights do not exist on the remote, we don't want to fail, we want to use seamlessly the PT weights.

That's why it gets more complex.

False -> Use only PT (easy)
None -> Use safetensors if both available on system and remote (which we don't know when entering the function and without checking the remote)
True -> USe only safetensors (easy)

Copy link
Contributor

Choose a reason for hiding this comment

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

Ah ok, I think the same logic still mostly applies

We can move the error check first and add another boolean for the case when we are allowed to fallback to non-safetensors

use_safetensors = kwargs.pop("use_safetensors", None)


if use_safetensors and not is_safetensors_available():
  raise ValueError("`use_safetensors`=True but safetensors is not installed. Please install safetensors with `pip install safetenstors")

if use_safetensors is None:
    use_safetensors = is_safetensors_available()
    allow_pickle = True


# in the rest of the function just use `use_safetensors` as a boolean

# somewhere else
if use_safetensors and not_safetensors_compatible(...):
     if not allow_pickle:
         raise ValueError("... pickle not allowed ...")
     # do the pickle loading    

Is this correct or am I still missing something?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

2 booleans work, Not sure it's easier to read ( I to avoid variable that create nonsensical combos, for instance allow_pickle=True and use_safetensors=False doesn't make sense)

False
None
True

Seems easier to reason about, but indeed creates some unpythoness.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I'll move to 2 booleans if you think it's clearer

Copy link
Contributor

Choose a reason for hiding this comment

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

t'would be much appreciated :)

Comment on lines 172 to 175
except EnvironmentError:
except EnvironmentError as e:
if use_safetensors is True:
raise e
Copy link
Contributor

Choose a reason for hiding this comment

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

It's potentially not the job of this PR because EnvironmentError was here before, but EnvironmentError is a very broad error e.g. anything io related. Can we add some docs on why it's ok to catch it here if we're using safe tensors?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I can add some comment, but I don't know why it's EnvironmentError, I'm guessing it could be IOError since it's a file missing error, but maybe we also want to catch other types of "missing error".

Copy link
Contributor

@williamberman williamberman Mar 16, 2023

Choose a reason for hiding this comment

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

If it's an IOError because a file is missing, should we just explicitly check for the file not being present? cc @patrickvonplaten I think you committed this, do you have context?

Comment on lines 520 to 521
with self.assertRaises(EnvironmentError):
new_model.load_attn_procs(tmpdirname, use_safetensors=True)
Copy link
Contributor

@williamberman williamberman Mar 16, 2023

Choose a reason for hiding this comment

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

Super nit: Similar as above, EnvironmentError is really broad. Should we be checking against a specific error class or error message for this test case that safetensors could not be used?

@williamberman
Copy link
Contributor

Love this and think it makes a lot of sense :) Just a few nits on errors and how we handle default cases

@Narsil
Copy link
Contributor Author

Narsil commented Mar 16, 2023

@williamberman adressed all your comments I think.

Copy link
Contributor

@williamberman williamberman left a comment

Choose a reason for hiding this comment

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

love it!

Comment on lines -172 to +185
except EnvironmentError:
except IOError as e:
if not allow_pickle:
raise e
Copy link
Contributor

Choose a reason for hiding this comment

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

@patrickvonplaten before you merge can you double check this error change is ok?

@patrickvonplaten
Copy link
Contributor

patrickvonplaten commented Mar 16, 2023

Tbh, I thought it was better before, also because the code was more in line how we've merged it to transformes: huggingface/transformers#22083 and also because I don't think we should introduce new mental concepts like allow_pickle that the user has to understand in order to read the code (many people have no idea what pickle is and we're never talking about it in the code).

Ok to merge now

@patrickvonplaten patrickvonplaten merged commit d9227cf into huggingface:main Mar 16, 2023
@Narsil Narsil deleted the add_use_safetensors branch March 16, 2023 15:00
w4ffl35 pushed a commit to w4ffl35/diffusers that referenced this pull request Apr 14, 2023
…ingface#2123)

* Adding `use_safetensors` argument to give more control to users

about which weights they use.

* Doc style.

* Rebased (not functional).

* Rebased and functional with tests.

* Style.

* Apply suggestions from code review

* Style.

* Addressing comments.

* Update tests/test_pipelines.py

Co-authored-by: Will Berman <[email protected]>

* Black ???

---------

Co-authored-by: Patrick von Platen <[email protected]>
Co-authored-by: Will Berman <[email protected]>
yoonseokjin pushed a commit to yoonseokjin/diffusers that referenced this pull request Dec 25, 2023
…ingface#2123)

* Adding `use_safetensors` argument to give more control to users

about which weights they use.

* Doc style.

* Rebased (not functional).

* Rebased and functional with tests.

* Style.

* Apply suggestions from code review

* Style.

* Addressing comments.

* Update tests/test_pipelines.py

Co-authored-by: Will Berman <[email protected]>

* Black ???

---------

Co-authored-by: Patrick von Platen <[email protected]>
Co-authored-by: Will Berman <[email protected]>
AmericanPresidentJimmyCarter pushed a commit to AmericanPresidentJimmyCarter/diffusers that referenced this pull request Apr 26, 2024
…ingface#2123)

* Adding `use_safetensors` argument to give more control to users

about which weights they use.

* Doc style.

* Rebased (not functional).

* Rebased and functional with tests.

* Style.

* Apply suggestions from code review

* Style.

* Addressing comments.

* Update tests/test_pipelines.py

Co-authored-by: Will Berman <[email protected]>

* Black ???

---------

Co-authored-by: Patrick von Platen <[email protected]>
Co-authored-by: Will Berman <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
stale Issues that haven't received updates
Projects
None yet
Development

Successfully merging this pull request may close these issues.

6 participants