Skip to content

Support special mapping of dtypes when preparing device map#1179

Merged
sgugger merged 1 commit intomainfrom
infer_device_map_spec_dtype
Mar 13, 2023
Merged

Support special mapping of dtypes when preparing device map#1179
sgugger merged 1 commit intomainfrom
infer_device_map_spec_dtype

Conversation

@sgugger
Copy link
Copy Markdown
Collaborator

@sgugger sgugger commented Mar 10, 2023

This PR will go along with [PR to come in Transformers] that fixes big model inference for T5 models. T5 models now have some weights that are kept in float32, which interferes with the computation done by infer_auto_device_map. On my setup, it gets me OOM when trying to load any big T5 model.

To fix that, on the Accelerate side we now accept a new argument special_dtypes which maps weight names to dtypes to use (and for the others we use the dtype provided).

@HuggingFaceDocBuilderDev
Copy link
Copy Markdown

HuggingFaceDocBuilderDev commented Mar 10, 2023

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

Copy link
Copy Markdown
Contributor

@muellerzr muellerzr left a comment

Choose a reason for hiding this comment

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

Thanks!

Copy link
Copy Markdown
Contributor

@younesbelkada younesbelkada left a comment

Choose a reason for hiding this comment

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

Thanks a mile for this!

@sgugger sgugger merged commit d942708 into main Mar 13, 2023
@sgugger sgugger deleted the infer_device_map_spec_dtype branch March 13, 2023 14:48
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.

4 participants