← Back to Table of Contents

Chapter 27 — ASICs & Accelerators

“NVIDIA dominates, but the AI hardware landscape is diversifying. Specialized silicon — TPUs, LPUs, and custom accelerators — are reshaping what’s possible.”

The AI Accelerator Landscape

Google TPUs

TPUs (Tensor Processing Units) are Google’s in-house accelerators, designed from the ground up for neural network workloads.

Architecture

TPUs use systolic arrays — a grid of multiply-accumulate units that data flows through rhythmically, performing matmuls with maximal data reuse and minimal memory access:

TPU Generations

TPUs are deployed in pods — tightly interconnected clusters connected via Inter-Chip Interconnect (ICI). A TPU v5p pod has 8,960 chips. Google used 16,384 TPU v5p chips to train Gemini.

TPU Programming Model

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# JAX: NumPy-like API that compiles to XLA for TPUs
import jax
import jax.numpy as jnp
from jax import pmap  # parallel map across devices

# Automatically shard computation across TPU cores
@pmap
def train_step(params, batch):
    def loss_fn(p):
        logits = model_forward(p, batch['input_ids'])
        return cross_entropy(logits, batch['labels'])
    
    grads = jax.grad(loss_fn)(params)
    grads = jax.lax.pmean(grads, axis_name='batch')  # all-reduce
    params = jax.tree.map(lambda p, g: p - lr * g, params, grads)
    return params

# With jax.sharding for FSDP-style training
from jax.sharding import NamedSharding, PartitionSpec as P, Mesh
mesh = Mesh(jax.devices(), ('dp', 'tp'))

AMD GPUs

AMD’s MI300X directly competes with NVIDIA’s H100/H200:

AMD’s advantage: 2.4× more HBM at competitive bandwidth. This matters for inference where memory capacity determines the largest model you can serve. The software gap (ROCm vs CUDA) remains the primary barrier to adoption, but PyTorch and JAX now support ROCm.

Intel Gaudi

Intel Gaudi accelerators target training and inference with a focus on cost efficiency:

Gaudi uses standard Ethernet (RoCE) for inter-chip communication rather than proprietary interconnects, reducing networking costs for training clusters.

Groq LPU

The Language Processing Unit (LPU) by Groq takes a radically different approach:

Groq achieves ~500 tokens/sec on LLaMA-2 70B (across 8 LPUs), with single-digit ms latency per token. The tradeoff: limited SRAM means model weights must be distributed across many chips. Inference only — not designed for training.

AWS Trainium & Inferentia

Amazon’s custom chips, integrated directly with AWS services:

Programmed via AWS Neuron SDK, which compiles PyTorch models to run on Trainium/Inferentia. Used by Anthropic to train Claude models.

Cerebras WSE

The Wafer-Scale Engine (WSE-3) is a single massive chip occupying an entire silicon wafer:

• 4 trillion transistors (56× more than H100)
• 900,000 AI cores
• 44 GB on-chip SRAM
• 21 PB/s memory bandwidth (on-chip)

The entire model fits on-chip, eliminating the memory wall entirely. Cerebras targets training large models without the complexity of distributed parallelism.

Apple Silicon

Increasingly relevant for local inference and research:

Unified memory means the CPU and GPU share the same memory — no PCIe transfer overhead. The M4 Ultra with 256 GB can run LLaMA-3 70B at 4-bit quantization entirely in memory. Tools like llama.cpp with Metal backend and MLX are optimized for Apple Silicon.

Choosing Hardware

What’s Next

With hardware and systems covered, we move to the software ecosystem. The next chapter introduces the Hugging Face ecosystem — the libraries, hubs, and tools that connect models to practitioners.

← Previous: Chapter 26 — Distributed Training · Next: Chapter 28 — The Hugging Face Ecosystem →

Last updated: April 2026