Why We Started with JAX but Moved to PyTorch

When we started training foundation models at Liquid AI in 2023, we made what seemed like an obvious choice: build our entire ML infrastructure on JAX. Three months in, however, we switched to PyTorch. Here, we describe why and what we have learned.
Spoiler: it has nothing to do with the code.

Our initial decision to use JAX was a no-brainer. I had spent the latter part of my journey in academia heavily invested in the framework. My last main contributor paper before founding Liquid AI was Gigastep, a multi-agent RL (MARL) benchmark published at NeurIPS 2023. We built a suite of MARL environments entirely in JAX, supporting everything from agent observations in the form of rendered images and state vector observations, to heterogeneous agents types (think one strong agent versus five weak ones) and asymmetric teams and objectives.

The entire system was jax.jit, vmap, and pmap compatible, allowing us to vectorize and parallelize on GPUs efficiently, e.g., scaling batched rollout of a million agents simultaneously on a GPU, thus running meaningful MARL experiments in less than two minutes on a single RTX 3090.

But personal experience of me and the team was just one factor. JAX offered three compelling technical advantages for what we were building:

1. Built-in support for scans (linear and associative): Critical for accelerating S5-style recurrences that were central to our architecture experiments

2. Superior compilation promises: jax.jit compilation theoretically reduced overhead for our custom operations—something PyTorch wouldn’t match until torch.compile came along much later

3. Accessible model sharding: Coming from academia where data parallelism was the maximum with out GPU budget, JAX’s pjit made model sharding surprisingly approachable, with xmap providing lower-level distributed primitives when needed

Armed with these advantages and an AWS cluster of NVIDIA A100 40GB GPUs, we dove in.

Initially, everything worked brilliantly. We iterated rapidly on model architectures, training parameters, data pipelines, and evaluation metrics. Model sharding on a single node with 8 A100s worked like a charm. The developer experience was smooth, and our training throughput met expectations.

Then we scaled to multiple nodes.

When we started training on more than a single node, the numbers didn’t make sense. Training on two nodes was only half as fast as a single node. We needed four nodes just to break even, and only at eight nodes did we exceed single-node throughput.

Something was fundamentally broken.

Our first instinct was to blame ourselves. Maybe we implemented pjit sharding incorrectly? We spent some time experimenting with different configurations. When that failed, we still thought we were not using pjit correctly and went lower level, implementing manual sharding with xmap to eliminate any automatic optimization issues.

Same result.

Finally, we used xmap‘s low-level primitives to implement bandwidth benchmarks for collective operations (all_gather, all_reduce, etc). The results were stark: within a single node, bandwidth was excellent. Across nodes? It dropped by 100x.

The root cause became clear when we mapped out the communication stack:

• JAX: Google’s ML framework, optimized for TPUs as first-class citizens

• NCCL (NVIDIA Collective Communications Library): NVIDIA’s collective ops library, designed primarily for their own ConnectX/Mellanox/Quantum interconnect

• EFA (Elastic Fabric Adapter): AWS’s custom interconnect solution

Our JAX installation used a CUDA backend that relied on NCCL for collective communication, which should have used AWS’s EFA for cross-node communication. But the NCCL version shipped with our JAX installation did not seem to be working with the installed EFA.

Three vendors, three competing solutions, zero compatibility.

Interestingly, according to SemiAnalysis, OpenAI prefers not to use EFA with their AWS capacity but the NVIDIA ConnectX reference fabric. We might not be the only ones having a not-so-great experience with EFA.

The kicker? We were not directly using AWS but subrenting our cluster from a provider that was itself renting from AWS. We had no bare metal access, no sudo privileges, no docker, and no ability to install anything or change any meaningful configurations.

Our available option to solve the issue was to try a million different environment variables, eg. enabling stdout logging for NCCL and changing adapter settings, but we had no luck.

Moving to a different cluster was also not an option in the middle of the 2023 GPU shortage and the preferable contract we had in place.

In short, our hands were pretty tied.

While troubleshooting, we implemented the same bandwidth test in PyTorch. It worked perfectly, achieving full cross-node bandwidth, no configuration needed.

This was in the middle of 2023, just a few months after starting Liquid AI. Our codebase was written by only two people and still manageable, maybe a few thousand lines of core training code. We faced a choice: spend potentially valuable days trying to fix the JAX-NCCL-EFA interaction without proper access thus limited hope, or port everything to PyTorch.

We estimated a few days for the port and decided to try it.

The migration revealed something interesting about modern ML frameworks, they are more similar than we thought. In less than a day we had ported most of our training stack, and on day 2 we brought over the remaining components.

The Easy Wins:

• Model definitions translated almost 1:1 (mostly syntactic changes)

• Training loops remained conceptually identical minus the jax.jit of the train_step and eval_step

• Data loaders ported trivially (just wrap it in a torch.utils.data.Dataset)

• FSDP and Deepspeed made sharding just as accessible as pjit

• Activation checkpointing was actually more mature in PyTorch

The Pleasant Surprises:

• PyTorch eager mode matched JAX JIT performance on our workloads on single node runs (disappointing for JAX)

• Always running in eager mode made debugging much easier

• Debugging distributed multi-node training became trivial with:
  if global_rank==0: breakpoint(); barrier()
  This line alone saved us tons of error investigation time and headache.

• Implementing associative scan for our modified S5 operators took just a few hours, achieving comparable speed and memory usage as in JAX

• Access to a bigger ecosystem, including Flash Attention, novel optimizers, more research implementations, MegatronLM, Deepspeed, and the vast Github collection of Phil Wang aka “lucidrains”

The Minor Quirks:

• Weight initializer distributions have different names and defaults between the two frameworks, thus requiring a bit of extra care

• FSDP didn’t work so well with anything “non-standard”, eg. parameter optimizer groups for multiple learning ratesneeded careful handling

• Memory leak potential with PyTorch graphs required attention, eg. OOMs in the middle of a training epoch

This experience taught us invaluable lessons that guide our engineering decisions today:

1. Always benchmark your infrastructure first. Before writing a single line of model code, test your cluster’s bandwidth, especially for multi-node setups.

2. Framework lock-in is overstated. The similarities between modern ML frameworks mean that early migration is feasible if you haven’t accumulated years of technical debt. This is even more true in the age of coding agents.

3. Ecosystem matters as much as elegance. JAX’s functional programming paradigm is beautiful, but PyTorch’s ecosystem breadth and industry adoption provide practical advantages that compound over time.

4. Control your stack where it matters. Not having sudo access or at least some container engine available seems like a small constraint until it becomes the bottleneck for your entire operation.

We would have enjoyed building multi-node training stack for our foundation model research regardless of choosing JAX or PyTorch. JAX is an excellent framework, and in a different context (say, with TPUs or proper NCCL-compatible interconnect fabric, we might still be using it today).

What this experience reinforced is that at Liquid AI, we are building something interesting with a small, highly capable team. We make technical decisions based on what works, not what is currently trendy. We balance theoretical elegance with practical constraints. And when something does not work, we pivot fast.

We are Hiring

This kind of problem solving, balancing cutting-edge ML research with practical engineering constraints, is what we do every day at Liquid AI. If yo are an ML engineer or researcher who gets excited about building novel architectures, solving distributed systems challenges, and making pragmatic technical decisions that enable breakthrough AI capabilities, consider to join us in building the next generation of AI systems.