Skip to content

Fix loading with only 1 device or distributed config - #46197

Merged
Cyrilvallez merged 4 commits into
mainfrom
fix-device-map
May 26, 2026
Merged

Cyrilvallez merged 4 commits into
mainfrom
fix-device-map

Conversation

@Cyrilvallez

@Cyrilvallez Cyrilvallez commented May 25, 2026 •

Copy link
Copy Markdown
Member

What does this PR do?

Supersedes #46116. Some params will never go on the correct device if the device_map is not set correctly, even with distributed calls (e.g. all non-persistent buffers). This has changed in #45028 but should not have (setting the device only in core_model_loading.py to load the param is not enough, as non-persistent buffers are not part of the state_dict).
This is because we want to put directly everything on the correct device, INCLUDING non persistent buffers.

@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

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.

@Cyrilvallez Cyrilvallez changed the title Fix loading Fix loading with only 1 device or distributed config May 25, 2026
Comment on lines -4910 to +4915
for tied_param in self.all_tied_weights_keys.keys():
for tied_param in getattr(self, "all_tied_weights_keys", {}).keys():

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Unrelated but this was there for a very good reason, i.e. remote code

@github-actions

Copy link
Copy Markdown
Contributor

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=46197&sha=796937

@ArthurZucker ArthurZucker left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

let's add a remote code test!

@Cyrilvallez
Cyrilvallez merged commit 4a2e365 into main May 26, 2026
93 of 96 checks passed
@Cyrilvallez
Cyrilvallez deleted the fix-device-map branch May 26, 2026 02:12
kaixuanliu pushed a commit to kaixuanliu/transformers that referenced this pull request May 28, 2026
yuchenxie4645 pushed a commit to yuchenxie4645/transformers that referenced this pull request May 28, 2026
kashif pushed a commit to kashif/transformers that referenced this pull request Jun 1, 2026
khushali9 pushed a commit to khushali9/transformers that referenced this pull request Jun 8, 2026
stevhliu pushed a commit to stevhliu/transformers that referenced this pull request Jul 30, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants