Lightning-AI / Lightning-AI/pytorch-lightning
MPS Mixed-precision Autocast
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Description & Motivation
Support for MPS autocasting has recently be added in PyTorch 2.5.0 [here](https://github.com/pytorch/pytorch/pull/99272) and there is an ongoing effort to implement gradient scaling [here](https://github.com/pytorch/pytorch/issues/142397).
PyTorch Lightning does not currently support mixed-precision on MPS device but it could be added in a near future when gradient scaling is finalized.
Is this feature considered? This would allow reducing memory usage and improving training time for some models.
### Pitch
Currently PyTorch Lightning falls back to FP32 when trying to use mixed-precision and issues a warning mentioning CUDA.
I think that considering adding a path for MPS mixed-precision would be great.
### Alternatives
Stick to FP32 training when using MPS device.
### Additional context
thanks for your work!
cc @lantiga @borda @justusschock
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reviewing the linked PyTorch MPS autocasting pull request and the gradient-scaling issue to confirm the supported APIs and their status. Then trace PyTorch Lightning's mixed-precision device handling and its CUDA warning path; done means MPS mixed-precision selects the supported path without falling back to FP32, with coverage for the relevant behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100