PaperlayerView exact revision on arXiv

Source-linked AI summary

Accelerator Choice Is Not Enough: AlphaFold2 Inference on Cloud TPUs

Lorenzo Pazienza, Ihab El Bani

arXiv:2609.34818v2cs.DCcs.LGcs.PF

TL;DR

The paper asks how much AlphaFold2 inference performance depends on factors beyond accelerator choice, including compilation, chip mapping, and cloud-session stability. It measures the same workload across CPU, GPU, and TPU backends and examines JAX’s use of multi-chip devices. A single TPU chip outperformed the T4 in the August baseline, but later reruns did not reproduce the cloud baselines.

  • Problem

    How compilation, device mapping, and cloud-session stability affect AlphaFold2 performance beyond accelerator choice had received little systematic attention.

    Source: §1 · p. 1

  • Method

    The paper compares AlphaFold2 inference across CPU, GPU, and TPU backends and tests explicit JAX mappings for using multiple chips.

    Source: §13 · pp. 18–19

  • Results

    0.47 s per call on one TPU chip was 27.8× faster than the T4 in the August baseline.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Explicit pmap mapping can use the eight-chip slice, whereas default execution leaves seven chips idle and vmap does not recover them.

    Source: §13 · pp. 18–19

    The August baselines and headline ratios were not reproduced in later reruns, and the GPU baseline changed by roughly a factor of two.

    Source: §12 · p. 18

Abstract

from arXiv · show

AlphaFold2 is written in JAX, so the same inference code compiles and runs unchanged on CPUs, GPUs and Google Cloud TPUs. That portability makes the accelerator look like the main decision a user has to make. We show that it is not. Running one AlphaFold2 inference workload across a Colab CPU runtime, an NVIDIA T4 GPU and a dedicated eight-chip Cloud TPU v5e slice, we find a large hardware advantage for the TPU, 0.47 s per call in steady state on a single chip against 13.1 s on the T4 in the same measurement campaign, and three ways in which the software layer decides how much of it a user actually gets. The default execution path uses one chip of the eight, and at list prices the idle capacity makes the slice cost about as much per prediction as the GPU. Batching with JAX's vmap never exceeds single-query throughput, while mapping queries across chips with JAX's pmap gives eight chips 6.5-7.9x the throughput of one on a matched grid; automatic sharding leaves the per-chip footprint unchanged, consistent with replication, most plausibly because AlphaFold2 carries no sharding annotations. Our retained trace analysis of a first call at a new input shape reports about three quarters of the traced span in JAX tracing and compilation rather than execution. Reruns five weeks later reproduced neither cloud baseline, the GPU one off by roughly a factor of two, so the hardware ratio above is specific to one campaign.

1 Introduction

The paper examines how compilation, device mapping, and cloud-session stability shape AlphaFold2 inference performance beyond accelerator choice. It compares CPU, GPU, and TPU execution and explains how JAX/XLA and AlphaFold2’s computation structure affect the use of a multi-chip slice.

1 Introduction

  • The study measures backend performance, first-call compilation, and multi-chip use, addressing a gap in systematic attention to the hardware-to-user performance gap.The same AlphaFold2 inference code and input are run on a Colab CPU, NVIDIA T4, and dedicated eight-chip Cloud TPU v5e slice.Source: §1 · pp. 1–2
  • 0.47 s per call on one TPU chip was 27.8× faster than the T4 in the August baseline, though the cloud baselines did not reproduce weeks later.The TPU figure measures one chip of an eight-chip slice; the GPU rerun was roughly twice as fast, changing the TPU-over-GPU ratio.Source: §1 · p. 2
  • jax.vmap does not improve throughput, while jax.pmap maps independent queries across chips; automatic sharding leaves per-chip memory unchanged, consistent with replication.The paper also demonstrates an external ensemble mapped across chips with pmap, without changing AlphaFold2’s source.Source: §1 · p. 2

2.1 TPU slices and device placement

  • A TPU slice does not distribute work across its chips automatically: JAX/XLA or the programmer determines device mapping, while the chips communicate through an ICI fabric.The v5e slice uses a 2D torus inter-chip fabric, independent of the host network.Source: §2.1 · p. 3

2.2 The JAX/XLA compilation model

  • JAX traces new argument shapes into a jaxpr, lowers it to XLA HLO, and compiles it; repeated calls with the same signature reuse the binary, while changed shapes or dtypes trigger compilation.A persistent compilation cache can avoid XLA compilation across process restarts, but a fresh process still retraces Python to a jaxpr.Source: §2.2 · p. 3

2.3 AlphaFold2’s computational structure

  • AlphaFold2’s Evoformer and structure module are recycled, while its sequential ensemble loop offers no array axis to shard across chips.The paper spreads ensemble members across chips by explicitly mapping them, distinct from partitioning a single-member forward pass.Source: §2.3 · pp. 3–4

3 Methodology

The methodology compares AlphaFold2 inference on Colab CPU/GPU runtimes and a dedicated TPU slice, with controlled workload and timing procedures. It records provenance, variability, and measurement caveats that bound interpretation of the results.

3 Methodology

  • The study audits its measurement setup and records what the run artifacts do not contain.The repository includes a file-level provenance audit.Source: §3 · p. 4

3.1 Hardware

  • The TPU allocation is an eight-chip slice, but default single-query execution uses one chip while CPU and GPU runs use shared Colab runtimes.The TPU is billed as a full slice; Colab machines may vary between sessions, and shared tenancy is only a hypothesis for observed variation.Source: §3.1 · p. 4

3.2 Software

  • AlphaFold2’s public source is unmodified apart from an external JAX clip compatibility shim, but the August runs lack recorded source commits and locked environments.September reruns are tied to a repository commit and captured environment.Source: §3.2 · p. 4

3.3 Workload

  • Most runs use model_3, one ensemble member, no recycling, and float32; single-query tests use a fixed synthetic sequence, while several sweeps use rotated amino-acid alphabets.The two input families are not identical, including at 118 residues.Source: §3.3 · p. 5

3.4 Timing procedure

  • Each single-query process times initialization, a first predict call that compiles and executes, and a second call that reuses compilation as steady state.Timed predict calls include host-side confidence metrics and use block_until_ready; first-call profiler traces add finalization overhead, whose TPU magnitude is unknown.Source: §3.4 · p. 5

3.5 Repetitions and noise

  • Most configurations were measured in one process; repeated TPU measurements estimate noise only for that setup, while CPU and GPU reruns spanned later Colab sessions.The TPU noise characterization used three runs in the same pod; access to the TPU cluster later ended.Source: §3.5 · p. 6

3.6 Provenance

  • The original results and scripts were committed together, preventing a result-to-script revision match; August baselines and their headline ratios cannot be regenerated from the repository.The September CPU and GPU reruns are reproducible in the stronger sense of being tied to an explicit commit and captured environment.Source: §3.6 · p. 6

4 Single-Chip Hardware Comparison

The August baseline favors a single TPU v5e chip over the T4 and CPU for steady-state inference, but the comparison is qualified by unused slice capacity, cold-start overhead, and cross-session variability.

4.1 Main comparison

  • At steady state, the TPU is 451× faster than the CPU and 27.8× faster than the T4; the T4 is 16.2× faster than the CPU.Table 2 covers one 118-residue inference per backend, with untrained weights and float32; each ratio compares exactly two measured points.Source: §4.1 · p. 6, Table 2 · p. 6

4.2 TPU figures are single-chip figures

  • The eight-chip slice’s default single-query path uses one chip, while the other seven report 0 MB of device memory in use.The reported TPU comparisons are therefore against one v5e chip, not an eight-chip slice.Source: §4.2 · pp. 6–7

4.3 The cold-start regime, corrected for profiler overhead

  • Corrected cold/steady ratios are approximately 1.11× for CPU and 4.25× for GPU; the TPU’s raw 59.1× ratio remains unconfirmed and uncorrected.Profiler-trace finalization inflated the timed CPU and GPU first-predict calls; no separate TPU profiler-overhead log was available.Source: §4.3 · p. 7

4.4 Cross-session variability, on two backends

  • CPU steady state rose from 212.113 s (n=1, August) to 348.861 ± 4.690 s (n=5, 2026-09-15), while GPU fell from 13.086 s (n=1, August) to 6.567 ± 0.057 s (n=5, 2026-09-15).Neither backend reproduced its August baseline, and the repository records no verified cause for either gap.Source: §4.4 · pp. 7–8, Figure 2 · p. 8
  • The CPU shift cannot be assigned to host or JAX changes because neither was recorded for the August run.The cross-session gaps are consistent with cloud variability, but the evidence is thin and does not confirm a cause.Source: §4.4 · p. 8

5 Profiling the Cold-Start Regime

This section traces a first TPU call at a new input shape to locate its remaining cost, then tests how a persistent compilation cache changes cold-start timings. It relates the trace finding to prior work on AlphaFold2 compilation costs.

5.1 Tracing a first call

  • The trace captured one first-predict call for the single-query TPU configuration (model_3, 118 residues, num_recycle=0) in jax.profiler.trace.The trace was copied from the pod before reclamation, since pods were torn down within roughly a minute of job completion.Source: §5.1 · p. 8

5.2 The cache_miss finding

  • About 76% of the 16.56-s traced apply_fn span is JAX cache_miss self-time; the TPU stays idle while the host traces and compiles.The 76% is a share of the traced span, not the whole first-predict call: the trace spans 16.56 s, while separate first-predict timings are 27.43–28.80 s.Source: §5.2 · p. 9

5.3 Compilation cache

  • 66.48 s to 20.72 s: with a persistent cache, combined init_params and first-predict startup time falls 69% across process restarts at a fixed input shape.The first predict call alone drops from 28.80 s to 15.19 s (1.90×); steady state is unchanged.Source: Figure 3 · p. 9, §5.3 · p. 9

5.4 Relation to prior work

  • The cold-start finding is consistent with ParaFold’s report of JAX recompilation costs for AlphaFold2 on GPUs, rather than an independent discovery.ParaFold mitigates recompilation within a process by sorting input sequences by length [Zhong et al., 2022].Source: §5.4 · p. 10

6 Expressing Parallelism on a Multi-Chip TPU Slice

On the same eight-chip TPU slice, the JAX transformation determines whether inference uses one chip or distributes work across all eight. Explicit mapping distributes independent queries, while batching within one device and automatic sharding do not use the slice in the same way.

6.1 Default execution leaves most of the slice idle

  • Under the default path, only TPU_0 holds model data; the other seven chips report zero HBM use.The eight-chip slice exposes independent devices, but jax.jit or AlphaFold2's inference function does not distribute computation across them.Source: §6.1 · p. 10

6.2 Batching within one device: jax.vmap

  • At vmap batch sizes 1, 2, 4, and 8, throughput is 2.051, 1.928, 1.487, and 1.508 proteins/second, never exceeding batch size 1.The transformation vectorizes operations within one device; the measured degradation is non-monotonic.Source: §6.2 · pp. 10–11

6.3 Explicit device mapping: jax.pmap

  • At 100 residues, jax.pmap delivers a 6.53× eight-over-one-chip speedup (81.6% parallel efficiency), while the 6.92× single-query gain uses a different input family.Figure 4 contrasts vmap throughput with pmap across eight chips; the two speedups use different comparisons.Source: Figure 4 · p. 10, §6.3 · p. 11

6.4 Automatic sharding

  • Automatic mesh assignment retains 463 MB per chip, consistent with replication; absent sharding annotations are a plausible but unverified explanation, not the inactive ensemble loop.No compiled layout or partitioner decision was retained, and the run used num_ensemble = 1, so the loop body did not execute.Source: §6.4 · p. 11

6.5 An external ensemble with pmap + lax.pmean

  • The external ensemble runs one differently seeded single-member forward pass per chip, then averages predicted_lddt logits with pmean.All eight chips report nonzero HBM allocation; steady-state latency is 0.538 s, with no sequential num_ensemble=8 baseline for a speedup comparison.Source: §6.5 · p. 12
  • This ensemble demonstration leaves AlphaFold2's internal ensembling at one member and averages confidence logits, not predicted structures.Its eight members use a trivial single-sequence MSA and differ only by seed, so it demonstrates the mechanism rather than a scientifically meaningful ensemble.Source: §6.5 · p. 12

6.6 Summary

  • Across the four configurations, the hardware is unchanged; the JAX transformation determines whether work stays on one chip, spreads across queries, or distributes ensemble members.The external ensemble does not partition one model instance or repair automatic sharding.Source: §6.6 · pp. 12–13

7 Scaling with Chip Count and Sequence Length

A 16-point multi-query pmap grid measures throughput across chip counts and sequence lengths. Eight-chip speedup approaches linear scaling at longer lengths, while a descriptive power-law fit summarizes—but does not closely track—the grid at every length.

7.1 Measurement grid

  • The 16-point grid varies chips (1, 2, 4, 8) and sequence length (100, 250, 500, 1000); at length 100, throughput rises from 2.935 to 19.153 proteins/second across one to eight chips.Each point is a single steady-state measurement using the rotated-alphabet input; the grid uses one input family and timer.Source: §7.1 · p. 13, Figure 5 · p. 13

7.2 Measured speedup at each length

  • Eight-over-one-chip speedup rises from 6.53× at length 100 to 7.91× at length 1000, approaching ideal linear scaling at longer sequences.Fixed per-call overhead takes a larger share of short calls, leaving less of the eight-chip speedup to be realized.Source: §7.2 · p. 13

7.3 A descriptive power-law fit

  • The fitted law is throughput ≈4527.77 × chips^0.963 × length^-1.572 (R² = 0.981), with throughput falling faster than inverse sequence length.The fit systematically over-predicts at lengths 100 and 1000 and under-predicts at 250 and 500; with one measurement per cell, regression intervals would understate misspecification.Source: §7.3 · p. 14

8 Cost Analysis

This section estimates prediction costs from August on-demand rates and steady-state throughput, then compares whole-pod billing with hypothetical chip-count pricing. Keeping the pod busy matters more to unit cost than minimizing chip count.

8.1 Pricing basis

  • The pricing comparison uses August 2026 on-demand list rates, not actual bills; the TPU rate is $1.20 per chip-hour, or $9.60/hour for the full pod.The CPU estimate is the weakest input, and costs for smaller TPU slices are hypothetical because they failed to initialize on this setup.Source: §8.1 · p. 14

8.2 Cost per prediction, single call

  • At August list prices, the full-pod TPU costs $1.25/1k predictions, nearly matching the GPU despite 27.4× higher hourly cost.Seven idle chips offset the single chip’s higher throughput; the comparison is specific to the August GPU baseline.Source: §8.2 · p. 14, Table 3 · p. 14
  • Keeping all eight chips busy gives a whole-pod cost of about $0.181/1k predictions, within about 15% of a hypothetical single-chip rate.The comparison is not strictly like-for-like: multi-query throughput uses a rotated-alphabet input and times apply, while the single-chip baseline uses a fixed sequence and times predict.Source: §8.2 · p. 15

8.3 Cost as a function of chip count

  • At 100 residues, moving from 1 to 8 chips raises cost per prediction by 22.6%; the increase falls to about 1% at 1000 residues.The 1-to-8-chip penalty is 7% at 250 residues and 3% at 500; these are single-run grid measurements pricing hypothetical slices.Source: §8.3 · p. 15, Table 4 · p. 15

8.4 Practical implication

  • Choose chip count for latency and throughput needs, not unit-cost minimization; the dominant cost lever in this workload is whether rented chips stay busy.An eight-chip pod running a single-chip job wastes 87.5% of paid capacity; running eight independent proteins yields 6.92× the predictions without increasing hourly cost.Source: §8.4 · p. 15

9 Portability to AlphaFold3

The paper tests whether its AlphaFold2 TPU findings extend to AlphaFold3, using the newer model’s trained weights. The tested release lacks a TPU path, and outputs also differ across backends.

9 Portability to AlphaFold3

  • The AlphaFold3 test asks whether AlphaFold2’s TPU findings generalize, but uses real trained weights rather than the untrained weights used for AlphaFold2 elsewhere in the paper.The setup, timings, and per-sample score comparisons are reported in Appendix A.Source: §9 · p. 15

9.1 No TPU backend in the release we tested

  • The AlphaFold3 release tested rejects --jax_backend=tpu before model execution; its accepted backend values are cpu, gpu, and mps.Its documentation lists CPU and NVIDIA GPUs of compute capability 7.0 or greater as supported targets, not TPU.Source: §9.1 · pp. 15–16

9.2 Output differs across backends

  • Across backends, AlphaFold3’s per-sample ranking score differs by up to 32.33% on one sample, versus at most 1.67% across machines using the same backend.The proposed CUDA-related cause remains plausible but unconfirmed because the responsible operation was not isolated.Source: §9.2 · p. 16

9.3 Summary

  • The AlphaFold2 TPU-parallelism finding cannot be tested on this AlphaFold3 release; backend availability and output matching across sessions also vary in the tested setting.The output concern is tied to a documented GPU hardware edge case, but the authors do not confirm its cause.Source: §9.1 · p. 15, §9.3 · p. 16, §9.2 · p. 16

10 Related Work

Prior AlphaFold2 optimization work largely targets GPUs, while this paper examines end-to-end inference on a multi-chip TPU slice. Related TPU and parallelism studies address different tasks or provide design context rather than direct AlphaFold2 evidence.

10.1 AlphaFold2 acceleration

  • GPU-focused AlphaFold2 work covers inference and training, but the cited studies do not report AlphaFold2 inference across a multi-chip TPU slice using pmap, vmap, or automatic partitioning.ManyFold runs TPU training for a different, smaller model, while its AlphaFold-style inference comparison uses a GPU; MatrixFold studies ARM CPUs.Source: §10.1 · p. 16
  • JAXBench benchmarks AlphaFold2-derived TPU kernels rather than the end-to-end model, and the cited AlphaFold3 characterization excludes TPU.The broader AlphaFold2 optimization literature described here is GPU- and CPU-centric.Source: §10.1 · p. 16

10.2 Accelerator benchmarking methodology

  • Language-model TPU-versus-GPU comparisons provide methodological precedent, but the cited studies do not address protein folding.The hardware specifications rely on TPU architecture papers and Google’s v5e documentation because v5e has no dedicated peer-reviewed architecture paper in the cited account.Source: §10.2 · pp. 16–17

10.3 Compilation and parallelism frameworks

  • GSPMD’s annotation-driven design and DrJAX’s partitioner example are consistent with the hypothesis that unannotated AlphaFold2 modules leave automatic partitioning little to shard, but neither tests protein-folding models.The authors also read reported cloud-performance variability as consistent with reproducibility literature, without isolating a cause.Source: §10.3 · p. 17

11 Discussion

The discussion separates hardware capability from the performance exposed by a program’s execution path. It emphasizes both the cold-start boundary of the steady-state results and the idle capacity hidden by default single-query execution.

11 Discussion

  • Realized TPU performance depends on how parallelism is expressed, what profiler traces measure, and whether rented chips stay busy—not only on accelerator capability.The authors report that closing these gaps required JAX parallelism knowledge but no modification to AlphaFold2 source.Source: §11 · p. 17

11.1 Steady state versus cold start

  • The TPU’s raw cold-to-steady-state gap is 59.1×, uncorrected for profiler overhead; the measured steady-state advantage may not hold for workloads that never reach steady state.One-off calls at varying input shapes pay compilation costs repeatedly, while repeated use of one compiled shape amortizes them.Source: §11.1 · p. 17

11.2 The cost of the default path

  • The unmodified single-query driver uses one chip on an eight-chip slice without warning, leaving seven idle; the authors attribute this to documented JAX behavior, not a bug.The default configuration does not ask for a different execution path, and the resulting idle capacity is also billed.Source: §11.2 · p. 17

12 Limitations

The study’s limitations constrain how broadly its timing, scaling, and cost results can be reproduced or generalized. They include narrow input and hardware coverage, measurement and provenance gaps, and unpinned dependencies.

Scope of measurements

  • Most single-query timings use one 118-residue sequence, and the length sweep uses a different input family; all TPU measurements are on v5e only.The authors make no claim that findings hold on other TPU generations.Source: §12 · pp. 17–18
  • Randomly initialized weights make the results statements about compute rather than prediction quality, and the study did not test whether trained weights change timing.Random initialization does not affect the compiled graph’s shape.Source: §12 · p. 18

Measurement and provenance limits

  • CPU and GPU ran on provider-managed shared Colab machines, whereas the TPU used a dedicated Kubernetes Job, so measured session variability applies only to CPU and GPU.The authors do not treat that variability as evidence about TPU stability.Source: §12 · p. 18
  • The TPU-over-GPU and TPU-over-CPU speedups use August baselines that shifted on rerun, and the records do not establish that all three backends ran in one session.The original August scripts were not committed in the form run; September reruns cover CPU and GPU only, and no TPU rerun exists for a matched comparison.Source: §12 · p. 18
  • Most TPU results are single runs, the multi-chip and scaling-grid experiments lack repeat checks, and access ended before the experiments could be regenerated.The single-query configuration had three repeats and the auto-mesh experiment ran twice; neither covers pmap or the scaling grid.Source: §12 · p. 18
  • The TPU software stack was only partly pinned, and the AlphaFold2 commit was not recorded, so reruns may resolve different dependencies.The authors cannot separate any dependency effects from hardware and session effects.Source: §12 · p. 18

Scaling-law fit

  • The scaling-law fit deviates from measured grid values by up to 32% at tested extremes, with no confidence interval from the single measurement per cell.The global two-exponent power law is presented as a convenient summary, not a precise model.Source: §12 · p. 18

AlphaFold3 and cost limits

  • AlphaFold3’s five diffusion samples are single-process runs, so the CPU–GPU divergence is one comparison, and its proposed cause remains unconfirmed.The study also prices each backend from one steady-state throughput value and a single on-demand list-price snapshot, not distributions.Source: §12 · p. 18

13 Conclusion and Future Work

The conclusion argues that accelerator choice alone does not determine realized AlphaFold2 inference performance: software mapping, cold-start costs, baseline drift, and billing all matter. The authors identify further tests needed to assess how their findings generalize and provide the project materials for reproduction.

13 Conclusion and Future Work

  • Eight TPU chips deliver 6.53–7.91× one-chip throughput on a matched grid, but automatic sharding leaves per-chip footprint unchanged, consistent with replication; the proposed missing-annotation explanation is unconfirmed.An explicit pmap with jax.lax.pmean also maps an external ensemble across all eight chips without changing AlphaFold2 source.Source: §13 · pp. 18–19

Future work

  • Future tests include longer sequences, multimer inputs, later TPU generations, and whether FastFold’s Dynamic Axial Parallelism transfers to TPU and compares with pmap.The authors identify single sequence length and TPU generation as direct generalization gaps.Source: Future work · p. 19

Code and Data Availability

  • Benchmark scripts, Kubernetes Job definitions, notebooks, raw results, and a results catalogue are available in the project repository at the specified branch and commit.The catalogue records each value’s source and status, while methodology documentation describes setup provenance and missing records.Source: Code and Data Availability · p. 19

A AlphaFold3 Setup, Timings and Score Comparison

The appendix documents the AlphaFold3 setup and reports CPU–GPU timing and output-score differences. It also cautions that the observed divergence has a plausible, unconfirmed explanation.

A AlphaFold3 Setup, Timings and Score Comparison

A.1 Bringing AlphaFold3 up

  • AlphaFold3 required a native C++ build and a separate data-building step, while the T4 required an XLA flag to run; its released weights decompress to 1.15 GB.The C++ build ran inside a Kubernetes Job because root access was unavailable outside one; the weight-size comparison is against an approximate AlphaFold2 figure from notes.Source: §A.1 · p. 21

A.2 CPU and GPU results

  • 2401.38 s on the Colab CPU versus 86.21 s on the Colab T4 yields a 27.86× GPU-over-CPU speedup with compilation included on both sides.The authors do not compute an AlphaFold3-over-AlphaFold2 speed ratio because the models’ timing scopes and output counts differ.Source: §A.2 · p. 21

A.3 A reproducibility finding

  • Across backends, the best-ranked sample’s score falls from 0.41 to 0.33 and its disordered-residue fraction from 0.37 to 0.21, while the T4-only XLA flag’s effect was not isolated.CUDA Capability 7.x numerical issues are a plausible explanation for the divergence, but the authors did not instrument AlphaFold3 to confirm the cause.Source: §A.3 · pp. 21–22
Loading 2609.34818v2…