Skip to content

Implemented https://github.com/huggingface/huggingface_hub/issues/2516 #2532

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

Merged
merged 4 commits into from
Sep 11, 2024
Merged

Implemented https://github.com/huggingface/huggingface_hub/issues/2516 #2532

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
1e4198b
implemented https://github.com/huggingface/huggingface_hub/issues/2516
yuxi-liu-wired Sep 11, 2024
70019ce
fix ruff error
yuxi-liu-wired Sep 11, 2024
04b4da2
Update setup.py
yuxi-liu-wired Sep 11, 2024
fa0aaa5
Merge branch 'main' into mixin_safetensor
yuxi-liu-wired Sep 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 14 additions & 11 deletions src/huggingface_hub/hub_mixin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,8 @@
import torch # type: ignore

if is_safetensors_available():
import packaging.version
import safetensors
from safetensors.torch import load_model as load_model_as_safetensor
from safetensors.torch import save_model as save_model_as_safetensor

Expand Down Expand Up @@ -827,17 +829,18 @@ def _load_as_pickle(cls, model: T, model_file: str, map_location: str, strict: b

@classmethod
def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
load_model_as_safetensor(model, model_file, strict=strict) # type: ignore [arg-type]
if map_location != "cpu":
# TODO: remove this once https://github.com/huggingface/safetensors/pull/449 is merged.
logger.warning(
"Loading model weights on other devices than 'cpu' is not supported natively."
" This means that the model is loaded on 'cpu' first and then copied to the device."
" This leads to a slower loading time."
" Support for loading directly on other devices is planned to be added in future releases."
" See https://github.com/huggingface/huggingface_hub/pull/2086 for more details."
)
model.to(map_location) # type: ignore [attr-defined]
if packaging.version.parse(safetensors.__version__) < packaging.version.parse("0.4.3"):
load_model_as_safetensor(model, model_file, strict=strict) # type: ignore [arg-type]
if map_location != "cpu":
logger.warning(
"Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors."
" This means that the model is loaded on 'cpu' first and then copied to the device."
" This leads to a slower loading time."
" Please update safetensors to version 0.4.3 or above for improved performance."
)
model.to(map_location) # type: ignore [attr-defined]
else:
safetensors.torch.load_model(model, model_file, strict=strict, device=map_location)
return model


Expand Down