FRESH

Hacker News

TorchTPU: Running PyTorch Natively on TPUs at Google Scale

137 points by mji

by yu3zhou4

0 subcomment

Adding a support for new hardware to PyTorch is actually quite convenient. I did that with WebGPU using the same PrivateUse1 mechanism TorchTPU used. Every hardware has its own slot and identifier, and when you want to add a support for a new one without merging it into PyTorch, PrivateUse1 works essentially like plug-in slot
https://github.com/jmaczan/torch-webgpu

by in-silico

0 subcomment

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

by immanuwell

0 subcomment

pitch basically boils down to 'just change one line and it works' which sounds too good to be true, but if they actually pull it off at 100k-chip scale, that's genuinely a big deal

by Reubend

2 subcomments

Sounds good, but my main question is: is this a fork, or a new backend they're building in (like MPS)?

by noracists

0 subcomment

by MASNeo

1 subcomments

0 subcomment

by yujunjie

0 subcomment

by crimebrasil

0 subcomment