Skip to content

Add support for downcasting bf16 on TPUs - #523

Merged
muellerzr merged 7 commits into
mainfrom
downcast
Jul 20, 2022
Merged

Add support for downcasting bf16 on TPUs#523
muellerzr merged 7 commits into
mainfrom
downcast

Conversation

@muellerzr

Copy link
Copy Markdown
Contributor

This PR adds downcasting support for bf16 on TPUs. Per the XLA docs:

  • By default both torch.float and torch.double are torch.float on TPUs.
  • If XLA_USE_BF16 is set, then torch.float and torch.double are both bfloat16 on TPUs.
  • If XLA_DOWNCAST_BF16 is set, then torch.float is bfloat16 on TPUs and torch.double is float32 on TPUs.

@muellerzr muellerzr added enhancement New feature or request TPU Bug or feature on TPU platforms labels Jul 15, 2022
@muellerzr
muellerzr requested a review from sgugger July 15, 2022 09:52
@HuggingFaceDocBuilderDev

HuggingFaceDocBuilderDev commented Jul 15, 2022

Copy link
Copy Markdown

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

Comment thread src/accelerate/accelerator.py Outdated
Comment thread src/accelerate/accelerator.py Outdated
@muellerzr
muellerzr requested a review from sgugger July 19, 2022 17:34

@sgugger sgugger 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.

Much cleaner, thanks!

@muellerzr
muellerzr merged commit 158acdd into main Jul 20, 2022
@muellerzr
muellerzr deleted the downcast branch July 20, 2022 09:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request TPU Bug or feature on TPU platforms

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants