Study: TPU speed gains for AlphaFold2 vanish without proper batching
JAX's default path leaves seven of eight TPU chips idle. That erases the hardware's cost edge over a GPU.
The facts
- Researchers Lorenzo Pazienza and Ihab El Bani ran AlphaFold2 inference on a Colab CPU, an NVIDIA T4 GPU, and an 8-chip Cloud TPU v5e slice
- In steady state, a single TPU chip took 0.47s per call versus 13.1s on the T4 GPU in the same measurement campaign
- The default execution path uses only one of the eight TPU chips, so at list prices the idle capacity makes the slice cost about as much per prediction as the GPU
- Batching with JAX's vmap never exceeded single-query throughput; mapping queries across chips with pmap gave 6.5-7.9x the throughput of one chip on a matched grid
- JAX's automatic sharding left per-chip footprint unchanged, consistent with replication rather than sharding, likely because AlphaFold2 has no sharding annotations
- About three-quarters of a traced first-call span went to JAX tracing and compilation rather than execution
- Reruns five weeks later reproduced neither cloud baseline, with the GPU result off by roughly a factor of two, so the paper says the hardware ratio is specific to one measurement campaign
Picking the fastest chip for AlphaFold2 inference tells you almost nothing about what you'll actually pay or get. That is the argument of a new paper by Lorenzo Pazienza and Ihab El Bani.
The pair ran the same JAX-based protein structure model, unchanged, on three targets: a Colab CPU, an NVIDIA T4 GPU, and a Google Cloud TPU v5e slice with eight chips. Per chip, the TPU did beat the GPU by a wide margin. In steady state it took 0.47 seconds per call, against 13.1 seconds on the T4, in the same measurement campaign.
But that headline number obscures the part that matters for anyone paying the bill. The default execution path on the eight-chip slice uses exactly one chip. The other seven sit idle while the bill covers all eight. At list prices, the authors say, that idle capacity makes the TPU slice cost about as much per prediction as the GPU. The hardware advantage disappears once the software layer is accounted for.
The paper also tests the obvious fixes. Batching with JAX's vmap never beat single-query throughput at all. Mapping queries across chips with pmap did work, giving 6.5 to 7.9 times the throughput of a single chip on a matched grid. Automatic sharding, by contrast, left the per-chip memory footprint unchanged. That is consistent with the framework simply replicating the model across chips rather than splitting it, which the authors attribute to AlphaFold2 carrying no sharding annotations. A retained trace of a first call at a new input shape found about three-quarters of that span spent in JAX tracing and compilation, not execution — a fixed cost that recurs whenever the input shape changes.
Reproducibility took a further hit. Rerunning the same benchmarks five weeks later matched neither cloud baseline. The GPU figure was off by roughly a factor of two. The authors are explicit: their 0.47-versus-13.1-second ratio is specific to one measurement campaign, not a durable hardware fact.
Why it matters
Anyone provisioning a multi-chip TPU slice for a JAX model without sharding annotations should assume they are paying for seven idle chips. That holds unless they explicitly parallelise with pmap. Re-benchmark before trusting a vendor's or a paper's per-chip number as a cost figure.
Primary source
arXiv cs.LG — arxiv.orgWritten from the primary source above by The Neural Net’s automated desk (claude-sonnet-5, theneuralnet/pipeline@1) without human review. How this works. Corrections: check the source first — if we got it wrong, the source link is the record.