How do I get started with Jax on TPU VMs

#1

Post by yeesian » Sat, Mar 30, 2024, 3:35 AM UTC

How do I get started with Jax on TPU VMs
github.com

Re: How do I get started with Jax on TPU VMs

#2

Post by ein0p » Sat, Mar 30, 2024, 4:38 PM UTC

Hmm, 3% market share framework with barely any ecosystem and single vendor accelerators (Jax on TPU) vs a 60% market share framework with insanely rich ecosystem and ability to debug code on your own workstation (PyTorch on GPU)? In my informed opinion most people should use the latter unless they like wasting time on shiny things

Re: How do I get started with Jax on TPU VMs

#3

Post by drdirk » Sat, Mar 30, 2024, 4:43 PM UTC

ein0p wrote: post #2

Hmm, 3% market share framework with barely any ecosystem and single vendor accelerators (Jax on TPU) vs a 60% market share framework with insanely rich ecosystem and ability to debug code on your own workstation (PyTorch on GPU)? In my informed opinion most people should use the latter unless they like wasting time on shiny things

Jax uses the XLA compiler which is compatible with GPU and CPU.

Re: How do I get started with Jax on TPU VMs

#4

Post by _ntka » Sat, Mar 30, 2024, 4:51 PM UTC

ein0p wrote: post #2

Hmm, 3% market share framework with barely any ecosystem and single vendor accelerators (Jax on TPU) vs a 60% market share framework with insanely rich ecosystem and ability to debug code on your own workstation (PyTorch on GPU)? In my informed opinion most people should use the latter unless they like wasting time on shiny things

JAX is used by almost every large genAI player (Anthropic, Cohere, DeepMind, Midjourney, Character.ai, XAi, Apple, etc.). Its actual market share in foundation models development is something like 80%.

Also JAX is not just for TPU. It's mainly for GPU. It's usually 2-3x faster than torch on GPU: https://keras.io/getting_started/benchmarks/

Far more industry users of JAX use it on GPU compared to TPU.

Re: How do I get started with Jax on TPU VMs

#5

Post by kaycebasques » Sat, Mar 30, 2024, 5:09 PM UTC

_ntka wrote: post #4

ein0p wrote: post #2

Hmm, 3% market share framework with barely any ecosystem and single vendor accelerators (Jax on TPU) vs a 60% market share framework with insanely rich ecosystem and ability to debug code on your own workstation (PyTorch on GPU)? In my informed opinion most people should use the latter unless they like wasting time on shiny things

JAX is used by almost every large genAI player (Anthropic, Cohere, DeepMind, Midjourney, Character.ai, XAi, Apple, etc.). Its actual market share in foundation models development is something like 80%. Also JAX is not just for TPU. It's mainly for GPU. It's usually 2-3x faster than torch on GPU: https://keras.io/getting_started/benchmarks/ Far more industry users of JAX use it on GPU compared to TPU.

Are there any resources going into detail about why the big players prefer JAX? I've heard this before but have never seen explanations of why/how this happened.

Re: How do I get started with Jax on TPU VMs

#6

Post by ein0p » Sat, Mar 30, 2024, 5:17 PM UTC

_ntka wrote: post #4

ein0p wrote: post #2

Hmm, 3% market share framework with barely any ecosystem and single vendor accelerators (Jax on TPU) vs a 60% market share framework with insanely rich ecosystem and ability to debug code on your own workstation (PyTorch on GPU)? In my informed opinion most people should use the latter unless they like wasting time on shiny things

JAX is used by almost every large genAI player (Anthropic, Cohere, DeepMind, Midjourney, Character.ai, XAi, Apple, etc.). Its actual market share in foundation models development is something like 80%. Also JAX is not just for TPU. It's mainly for GPU. It's usually 2-3x faster than torch on GPU: https://keras.io/getting_started/benchmarks/ Far more industry users of JAX use it on GPU compared to TPU.

Are you on one of those (usually small) teams? No? Then it’s probably not a good choice for you.

Re: How do I get started with Jax on TPU VMs

#7

Post by ein0p » Sat, Mar 30, 2024, 5:20 PM UTC

drdirk wrote: post #3

ein0p wrote: post #2

Hmm, 3% market share framework with barely any ecosystem and single vendor accelerators (Jax on TPU) vs a 60% market share framework with insanely rich ecosystem and ability to debug code on your own workstation (PyTorch on GPU)? In my informed opinion most people should use the latter unless they like wasting time on shiny things

Jax uses the XLA compiler which is compatible with GPU and CPU.

Yes, but its performance on GPU leaves much to be desired, and 20 times as much research comes out on PyTorch. Would you rather just build on that or laboriously port and debug the models and their weights, losses, dataset readers, training regimes etc etc?

Re: How do I get started with Jax on TPU VMs

#8

Post by lern_too_spel » Sat, Mar 30, 2024, 5:33 PM UTC

ein0p wrote: post #2

Hmm, 3% market share framework with barely any ecosystem and single vendor accelerators (Jax on TPU) vs a 60% market share framework with insanely rich ecosystem and ability to debug code on your own workstation (PyTorch on GPU)? In my informed opinion most people should use the latter unless they like wasting time on shiny things

The parts of your comment that have any truth in them could have been said of PyTorch when it came out. People wasting time on shiny things is how we get better tools.

Re: How do I get started with Jax on TPU VMs

#9

Post by _ntka » Sat, Mar 30, 2024, 5:45 PM UTC

kaycebasques wrote: post #5

_ntka wrote: post #4

Earlier quoted context omitted.

JAX is used by almost every large genAI player (Anthropic, Cohere, DeepMind, Midjourney, Character.ai, XAi, Apple, etc.). Its actual market share in foundation models development is something like 80%. Also JAX is not just for TPU. It's mainly for GPU. It's usually 2-3x faster than torch on GPU: https://keras.io/getting_started/benchmarks/ Far more industry users of JAX use it on GPU compared to TPU.

Are there any resources going into detail about why the big players prefer JAX? I've heard this before but have never seen explanations of why/how this happened.

It's all about cost and performance. If you can train a foundation model 2x faster with JAX on the same hardware, you are effectively slashing your training costs by 2x, which is significant for a multi-million dollar training run.

Re: How do I get started with Jax on TPU VMs

#10

Post by _ntka » Sat, Mar 30, 2024, 5:48 PM UTC

ein0p wrote: post #6

_ntka wrote: post #4

Earlier quoted context omitted.

JAX is used by almost every large genAI player (Anthropic, Cohere, DeepMind, Midjourney, Character.ai, XAi, Apple, etc.). Its actual market share in foundation models development is something like 80%. Also JAX is not just for TPU. It's mainly for GPU. It's usually 2-3x faster than torch on GPU: https://keras.io/getting_started/benchmarks/ Far more industry users of JAX use it on GPU compared to TPU.

Are you on one of those (usually small) teams? No? Then it’s probably not a good choice for you.

Or alternatively, do you want faster training runs (and thus lower training costs)? Then JAX is a good choice for you.