MPS Mixed-precision Autocast

[laclouis5]

Description & Motivation

Support for MPS autocasting has recently be added in PyTorch 2.5.0 here and there is an ongoing effort to implement gradient scaling here.

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 @Borda

[lantiga]

Thank you @laclouis5, we'll be happy to accept a PR when grad scaling is in. Would you like to give it a shot?

[laclouis5]

@lantiga Yes, sure! The draft pull request is here: #20531.