search

pytorch / xla

Enabling PyTorch on XLA Devices (e.g. Google TPU)
https://pytorch.org/xla
Other
2.46k stars 467 forks source link

Add an example of using GKE with torchrun #7589

Closed will-cromar closed 3 months ago

will-cromar commented 3 months ago

I couldn't get access to a TPU GKE cluster today to test an example, so just putting a GPU example. The important part (torchrun) is identical.

Thanks @zpcore for letting me borrow your GPU cluster.