pytorch-lightning icon indicating copy to clipboard operation
pytorch-lightning copied to clipboard
verified publisher icon Lightning-AI

[TPU] Add Trainer support for PyTorch XLA FSDP

Open gkroiz opened this issue 2 years ago • 3 comments

Implements Pytorch XLA FSDP for TPUs. This PR is based on https://github.com/Lightning-AI/lightning/pull/17421, but focuses only on Trainer related changes. Use the XLAFSDPStrategy to use FSDP on TPUs.

Fixes https://github.com/Lightning-AI/lightning/issues/13209

gkroiz avatar Oct 08 '23 12:10 gkroiz

This issue was resolved in the 1.2 update.

heathercran avatar May 21 '25 00:05 heathercran