Re:How do I get started with Jax on TPU VMs
by _ntka on Saturday March 30, 2024 @08:26PM
(#39880101)
No, I am saying, with JAX you train on G.P.U., with a G, and your training runs are >2x faster, so your training costs are 2x lower, which matters whether your training spend is $1k or $100M. You're not interested in that? That's ok, but most people are.
[ Parent
| on HN ]
Re:How do I get started with Jax on TPU VMs
by ein0p on Saturday March 30, 2024 @09:05PM
(#39880354)
Have you actually tried that or are you just regurgitating Google’s marketing? I’ve seen Jax perform _slower_ than PyTorch on practical GPU workloads on the exact same machine, and not by a little, by something like 20%. I too thought I’d be getting great performance and “saving money”, but reality turned out to be a bit more complicated than that - you have to benchmark and tune.