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.

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

fix(examples): Move non-persistent buffers to a valid device #91

Merged
merged 1 commit into from
Feb 27, 2024

Conversation

Eta0
Copy link
Contributor

@Eta0 Eta0 commented Feb 27, 2024

Non-persistent buffer devices

During deserialization using load_into_module, only the buffers loaded from the deserializer are overwritten on the target module, and any extra buffers on the target module remain unmodified. This can lead to an error if there are existing buffers on the target module that are on a different device from the deserializer, since the model will end up scattered across multiple devices.

This primarily arises as an issue when serializing with include_non_persistent_buffers=False, since non-persistent buffers may be created on any device, and will not be overwritten by the deserializer. In particular, this led to an error when running the code in deserialize.py (which does not expect buffers from anything but the deserializer) on a model serialized using examples/hf_serialization.py (which uses include_non_persistent_buffers=False).

This change adds an extra section to deserialize.py that moves remaining non-persistent buffers to a device matching the deserializer after regular deserialization finishes. It is not included in the speed measurement because non-persistent buffers are explicitly not any business of the deserializer—only of external code—and they can be included and measured by setting include_non_persistent_buffers=True if it is desired to force the deserializer to handle them instead.

@Eta0 Eta0 added the bug Something isn't working label Feb 27, 2024
@Eta0 Eta0 requested review from wbrown and harubaru February 27, 2024 00:54
@Eta0 Eta0 self-assigned this Feb 27, 2024
@Eta0 Eta0 linked an issue Feb 27, 2024 that may be closed by this pull request
@wbrown wbrown merged commit 0b769c7 into main Feb 27, 2024
6 of 7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Fail to deserialize model: Expected all tensors to be on the same device
2 participants