hard★ EssentialNewGoogleGoogle DeepMind

You have trained on GPUs. What transfers to training on TPUs, and what do you have to relearn?

The roofline, the sharding arithmetic and the parallelism trade-offs transfer unchanged, and the numbers are in the same units. What changes is who writes the kernels, how shapes must behave, and the interconnect topology you shard against. The v6e and v7 numbers worked through against an H100.

Updated Sep 2026 · Grounded in real AI infrastructure interview loops and written to a senior-engineer editorial bar, with every number worked and every diagram hand-built.

more free answers with an account · no card

The concepts behind this question

Ranked by how closely each one overlaps this question's topic, so the first card is the thing to read if the answer above moved too fast.

Advanced

🧩 GPU & Accelerator Architecture🔒 Premium

TPU Architecture and Systolic ArraysA TPU is a matrix unit first and a processor second: a systolic array that streams activations through a grid of multiply-accumulate cells holding stationary weights, fed by a compiler rather than a warp scheduler, and scaled out over a dedicated inter-chip interconnect into pods of thousands. The roofline thinking transfers from GPUs unchanged; the workflow does not, which is what the Google loop tests.

Advanced

🔌 Networking & Storage🔒 Premium

Data Loading Pipelines for TrainingThe dataloader is the only part of a training job that runs on the CPU, the disk and the network at once, and it is the part most often found starving the GPUs. A pipeline that keeps 1,024 accelerators fed has to read sharded files sequentially, decode and tokenize in parallel workers, prefetch several batches ahead, pin memory for the PCIe copy, and do it deterministically enough to resume mid-epoch. The symptom of failure is a GPU at 30% utilization with nothing wrong on the GPU.

Foundational

💻 Coding for Infra

Consistent Hashing and ShardingSplitting work across N servers with a modulo of N moves almost everything when N changes, which for a cache means throwing away almost all of it. Consistent hashing places servers and keys on a ring so adding or removing one moves only its share, and virtual nodes fix the imbalance a small ring otherwise has. In LLM serving the same structure routes requests by prompt prefix so a conversation reaches the replica already holding its cache.

Core

🕸️ Distributed TrainingSign in

ZeRO and FSDPZeRO and FSDP keep data parallelism's simple programming model but shard the optimizer state, gradients and parameters across ranks, cutting per-GPU memory from 16 bytes per parameter toward 16/N. The price is 1.5x DDP's communication and a dependence on tokens per GPU that decides when sharding stops paying and tensor parallelism takes over.

Explore the full curriculum →

UP NEXT ON YOUR JOURNEY

Next in this trackExplain how a systolic array computes a matrix multiply. Why is it efficient, and why does it want big square tiles?Next in this trackWould you move a cost-sensitive training and serving fleet from H100 to Trainium2? What do you gain, and what do you have to plan for?Next in this trackThe MI300X has 192 GB per card. When does AMD actually win against an H100 for LLM serving, and what would you check before betting on it?

FEDITOR'S NOTE

Scored on the candidate applying the same quantitative reasoning to a TPU spec sheet, then naming precisely which habits break: dynamic shapes, hand-written kernels, NVLink-shaped parallelism assumptions.

DISCUSSION · 0

No comments yet — be the first to share your approach.