Lightning-AI / Lightning-AI/pytorch-lightning

MPS Mixed-precision Autocast

Open
#20,497 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

accelerator: mps feature
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.