3 ms·

Jax uses XLA (a IR that targets multiple platforms, including the Google TPU) as backend. Pytorch also has a XLA backend.

carbocation 5y ago
Underscoring your point, GCP has a tutorial page for training Pytorch models on TPU: https://cloud.google.com/tpu/docs/tutorials/pytorch-pod https://cloud.google.com/tpu/docs/tutorials/pytorch-pod
sillysaurusx 5y ago
A terrible one. It’s literally one of the worst experiences you can have. I say that without an ounce of bias. When tpu SSH first came out, I immediately went to my pytorch buddy (who became famous in the meantime by pioneering the CLIP AI art you’ve probably been seeing: https://mobile.twitter.com/rivershavewings https://mobile.twitter.com/rivershavewings) and said “Rivers, you have to try TPUs! They’re wonderful! Now you can run pytorch on them directly!” She was skeptical, because she’d had nothing but endless problems with the official pytorch TPU notebooks (which use a TPU in “remote” mode, aka it fires up an RPC server you attach to). “No no, trust me — I bet their rpc driver was just horrible. It can’t be true that the Pytorch XLA backend is terrible on TPUs. How could that be? Facebook makes awesome software, and they’ve worked on this for so long. Some intern probably wrote the rpc system or something. Anyway, you can SSH into mine! See, no RPC server! What’s your pubkey?” A few hours later, she reported that it froze in exactly the same place. I could hardly believe it. I was so upset that I dug really deeply to try to figure out what the heck the problem was. It was recompiling itself. Every. Inference. Every compile took 15 minutes. Their XLA backend is so far behind Jax that it’s not even a contest. Worse, I’ve since realized why they can’t fix the problem. JAX gives you precise control over which functions get JIT’ed (compiled into XLA), and when that happens. Pytorch doesn’t. There’s no @torch.jit decorator for your functions. That means they need to infer which parts to jit (the inference model), and which parts not to (the inference loop). The moment that magic fails, congrats; it’s now recompiling every iteration. (15 min per iteration is impressive, since usually it’s measured in “iterations per second”…) JAX even has a TPU compiler cache now, which persists across Python runs. It’s like MacOS in 2021 vs Windows in 1995. This is deeply frustrating to me, because there’s no reason for it — I don’t know if Facebook is intentionally kneecapping pytorch to keep people off of TPUs, or if it’s just run of the mill incompetence. But (like Windows 95) they’re going to discover they’re not the only game in town forever.
- bertr4nd 5y ago
  The pytorch programming model is just really hard to adapt to an XLA-like compiler. Imperative python code doesn't translate to an ML graph compiler particularly well; Jax's API is functional, so it's easier to translate to the XLA API. By contrast, torch/xla uses "lazy tensors" that record the computation graph and compile when needed. The trouble is, if the compute graph changes from run to run, you end up recompiling a lot. I guess in Jax you'd just only apply `jax.jit` to the parts where the compute graph is static? I'd be curious to see examples of how this works in practice. Fwiw, there's an offshoot of pytorch that is aiming to provide this sort of API (see https://github.com/pytorch/functorch https://github.com/pytorch/functorch and look at eager_compilation.py). (Disclaimer: I worked on this until quite recently.)

carbocation 5y ago

Underscoring your point, GCP has a tutorial page for training Pytorch models on TPU: https://cloud.google.com/tpu/docs/tutorials/pytorch-pod https://cloud.google.com/tpu/docs/tutorials/pytorch-pod

sillysaurusx 5y ago

A terrible one. It’s literally one of the worst experiences you can have. I say that without an ounce of bias. When tpu SSH first came out, I immediately went to my pytorch buddy (who became famous in the meantime by pioneering the CLIP AI art you’ve probably been seeing: https://mobile.twitter.com/rivershavewings https://mobile.twitter.com/rivershavewings) and said “Rivers, you have to try TPUs! They’re wonderful! Now you can run pytorch on them directly!” She was skeptical, because she’d had nothing but endless problems with the official pytorch TPU notebooks (which use a TPU in “remote” mode, aka it fires up an RPC server you attach to). “No no, trust me — I bet their rpc driver was just horrible. It can’t be true that the Pytorch XLA backend is terrible on TPUs. How could that be? Facebook makes awesome software, and they’ve worked on this for so long. Some intern probably wrote the rpc system or something. Anyway, you can SSH into mine! See, no RPC server! What’s your pubkey?” A few hours later, she reported that it froze in exactly the same place. I could hardly believe it. I was so upset that I dug really deeply to try to figure out what the heck the problem was. It was recompiling itself. Every. Inference. Every compile took 15 minutes. Their XLA backend is so far behind Jax that it’s not even a contest. Worse, I’ve since realized why they can’t fix the problem. JAX gives you precise control over which functions get JIT’ed (compiled into XLA), and when that happens. Pytorch doesn’t. There’s no @torch.jit decorator for your functions. That means they need to infer which parts to jit (the inference model), and which parts not to (the inference loop). The moment that magic fails, congrats; it’s now recompiling every iteration. (15 min per iteration is impressive, since usually it’s measured in “iterations per second”…) JAX even has a TPU compiler cache now, which persists across Python runs. It’s like MacOS in 2021 vs Windows in 1995. This is deeply frustrating to me, because there’s no reason for it — I don’t know if Facebook is intentionally kneecapping pytorch to keep people off of TPUs, or if it’s just run of the mill incompetence. But (like Windows 95) they’re going to discover they’re not the only game in town forever.

bertr4nd 5y ago
The pytorch programming model is just really hard to adapt to an XLA-like compiler. Imperative python code doesn't translate to an ML graph compiler particularly well; Jax's API is functional, so it's easier to translate to the XLA API. By contrast, torch/xla uses "lazy tensors" that record the computation graph and compile when needed. The trouble is, if the compute graph changes from run to run, you end up recompiling a lot. I guess in Jax you'd just only apply `jax.jit` to the parts where the compute graph is static? I'd be curious to see examples of how this works in practice. Fwiw, there's an offshoot of pytorch that is aiming to provide this sort of API (see https://github.com/pytorch/functorch https://github.com/pytorch/functorch and look at eager_compilation.py). (Disclaimer: I worked on this until quite recently.)