TorchTPU: Running PyTorch Natively on TPUs at Google Scale

(developers.googleblog.com)

81 points | by mji 6 hours ago

6 comments

  • noracists 28 minutes ago
    Very excited for this.
  • in-silico 4 hours ago
    This is great to see.

    I did trained some research models using the existing PyTorch/XLA on TPUs, and it was a mess of undocumented behavior and bugs (silently hanging after 8 hours of training!).

    If anyone is trying to use PyTorch on TPU before TorchTPU is released, you can check out the training pipeline that I ended up building to support my research: https://github.com/aklein4/easy-torch-tpu

  • Reubend 3 hours ago
    Sounds good, but my main question is: is this a fork, or a new backend they're building in (like MPS)?
  • yujunjie 21 minutes ago
    [dead]
  • crimebrasil 2 hours ago
    [dead]