Skip to content
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

Serialization: take into account meta tensor when splitting the state_dict #2591

Merged
merged 9 commits into from
Oct 10, 2024
Merged

Serialization: take into account meta tensor when splitting the state_dict #2591

Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
ff504a3
Enable meta tensor serialization
SunMarc Oct 7, 2024
c5a9c12
getattr is better
SunMarc Oct 7, 2024
93afb29
style
SunMarc Oct 7, 2024
8ac1145
skip meta tensors
SunMarc Oct 7, 2024
d57f813
Merge branch 'main' into serialization-meta-device
Wauplin Oct 9, 2024
09ff14c
update doc
SunMarc Oct 9, 2024
cac3858
Merge branch 'serialization-meta-device' of github.com:huggingface/hu…
SunMarc Oct 9, 2024
ed2ba2c
Update src/huggingface_hub/serialization/_torch.py
SunMarc Oct 9, 2024
0651877
oups
SunMarc Oct 9, 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
11 changes: 7 additions & 4 deletions src/huggingface_hub/serialization/_torch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -368,18 +368,21 @@ def _get_unique_id(tensor: "torch.Tensor") -> Union[int, Tuple[Any, ...]]:
return unique_id


def get_torch_storage_id(tensor: "torch.Tensor") -> Tuple["torch.device", Union[int, Tuple[Any, ...]], int]:
def get_torch_storage_id(tensor: "torch.Tensor") -> Optional[Tuple["torch.device", Union[int, Tuple[Any, ...]], int]]:
"""
Return unique identifier to a tensor storage.

Multiple different tensors can share the same underlying storage. For
example, "meta" tensors all share the same storage, and thus their identifier will all be equal. This identifier is
Multiple different tensors can share the same underlying storage. This identifier is
guaranteed to be unique and constant for this tensor's storage during its lifetime. Two tensor storages with
non-overlapping lifetimes may have the same id.
In the case of meta tensors, we return None since we can't tell if they share the same storage.

Taken from https://github.com/huggingface/transformers/blob/1ecf5f7c982d761b4daaa96719d162c324187c64/src/transformers/pytorch_utils.py#L278.
"""
return tensor.device, _get_unique_id(tensor), get_torch_storage_size(tensor)
if tensor.device.type == "meta":
return None
else:
return tensor.device, _get_unique_id(tensor), get_torch_storage_size(tensor)
Comment on lines +382 to +385
Copy link
Contributor

Choose a reason for hiding this comment

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

I'm fine with the changes you've made 👍
Could you just update the docstring just above to reflect this change? Currently it says

Multiple different tensors can share the same underlying storage. For example, "meta" tensors all share the same storage, and thus their identifier will all be equal. This identifier is guaranteed to be unique and constant for this tensor's storage during its lifetime. Two tensor storages with non-overlapping lifetimes may have the same id.

which is not true anymore (will always return None)

And thanks for looking into it in the first place!

Copy link
Member Author

Choose a reason for hiding this comment

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

Done !



def get_torch_storage_size(tensor: "torch.Tensor") -> int:
Expand Down
Loading