KEMBAR78
[MPS] Add support for bf16 autocast by hvaara · Pull Request #139390 · pytorch/pytorch · GitHub
Skip to content

[MPS] Add support for bf16 autocast #139390

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

Jump to bottom
Closed
wants to merge 4 commits into from
Closed

[MPS] Add support for bf16 autocast #139390

wants to merge 4 commits into from

Conversation

@hvaara
Copy link
Contributor

@hvaara hvaara commented Oct 31, 2024 •
edited by pytorch-bot bot
Loading

This PR adds support for bf16 autocast. Most of the code and ideas are copied from #99272.

Most of the heavy lifting was done by AI.

Fixes #139386

cc @mcarilli @ptrblck @leslie-fang-intel @jgong5

@hvaara @kulinseth @skotapati
[MPS] Add support for bf16 autocast
3a5ee0f 
Co-authored-by: Kulin Seth <kulin_seth@apple.com>
Co-authored-by: Siddharth Kotapati <skotapati@apple.com>
@pytorch-bot
Copy link

pytorch-bot bot commented Oct 31, 2024 •
edited
Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/139390

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

✅ You can merge normally! (1 Unrelated Failure)

As of commit e716891 with merge base 3192bde (image):

FLAKY - The following job failed but was likely due to flakiness present on trunk:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

hvaara
hvaara commented Oct 31, 2024
View reviewed changes
This was referenced Oct 31, 2024
Closed
Closed
@cpuhrsch cpuhrsch requested review from albanD and kulinseth October 31, 2024 22:21
@cpuhrsch cpuhrsch added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Oct 31, 2024
malfet
malfet reviewed Nov 1, 2024
View reviewed changes
hvaara and others added 3 commits November 2, 2024 23:40
@hvaara @malfet
Use helper for version predicate
dc71f88 
Co-authored-by: Nikita Shulga <2453524+malfet@users.noreply.github.com>
@hvaara
Remove superfluous import
b1b15a8
@hvaara
Remove superfluous unit test
e716891
@hvaara hvaara requested a review from malfet November 2, 2024 22:47
@qqaatw qqaatw added release notes: mps Release notes category ciflow/trunk Trigger trunk jobs on your pull request labels Nov 8, 2024
@hvaara
Copy link
Contributor Author

hvaara commented Nov 14, 2024

@malfet is this good to go?

malfet
malfet approved these changes Nov 20, 2024
View reviewed changes
Copy link
Contributor

@malfet malfet left a comment

Choose a reason for hiding this comment

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

LGTM

@malfet
Copy link
Contributor

malfet commented Nov 20, 2024

@pytorchbot merge -i

@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged while ignoring the following 1 checks: pull / cuda12.1-py3.10-gcc9-sm75 / test (pr_time_benchmarks, 1, 1, linux.g4dn.metal.nvidia.gpu)

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

@hvaara hvaara deleted the extend-autocast-support branch November 20, 2024 22:13
This was referenced Nov 28, 2024
Closed
Closed
pobin6 pushed a commit to pobin6/pytorch that referenced this pull request Dec 5, 2024
@hvaara @kulinseth @malfet
[MPS] Add support for bf16 autocast (pytorch#139390)
b7afb0a 
This PR adds support for bf16 autocast. Most of the code and ideas are copied from pytorch#99272.

Most of the heavy lifting was done by AI.

Fixes pytorch#139386

Pull Request resolved: pytorch#139390
Approved by: https://github.com/malfet

Co-authored-by: Kulin Seth <kulin_seth@apple.com>
Co-authored-by: Nikita Shulga <2453524+malfet@users.noreply.github.com>
@hg0428
Copy link

hg0428 commented Feb 18, 2025

What happened here?
What's the progress?
How do I use autocast with MPS?

@hvaara
Copy link
Contributor Author

hvaara commented Feb 18, 2025

@hg0428 The PR was merged to main in bc69a19, and is available in stable since v2.6.0. PyTorch's merge process is handled oob which is why the PR is closed on GitHub.

For a primer on how to use autocast please see https://pytorch.org/docs/stable/amp.html and https://pytorch.org/docs/stable/notes/amp_examples.html. If it's still not clear, please ask on https://discuss.pytorch.org/.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@malfet malfet malfet approved these changes

@kulinseth kulinseth Awaiting requested review from kulinseth

@albanD albanD Awaiting requested review from albanD

Assignees

No one assigned

Labels

ciflow/trunk Trigger trunk jobs on your pull request Merged module: amp (automated mixed precision) autocast open source release notes: mps Release notes category triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[MPS] Extend autocast support to bf16

7 participants

@hvaara @malfet @pytorchmergebot @hg0428 @cpuhrsch @pytorchbot @qqaatw