Where AI Data Lives

The idea in one minute#

An AI system moves six kinds of data, and each has a natural home in the memory hierarchy:

Each tier down is bigger, slower and cheaper per byte. Most performance work in AI infrastructure is keeping the hot data one tier higher than it would otherwise be. Checked on 16 October 2026.

A picture#

How it really works#

The tiers, with numbers#

Training data: the full pipeline#

A training run reads its tokenised dataset from object storage, caches shards on local NVMe, loads batches into CPU memory, and copies them to HBM, with prefetching so the GPU never waits. Text is small (a 15T-token dataset is tens of terabytes tokenised); images and video are not, which is why multimodal training is often storage-bound.

Optimizer state: the hidden multiplier#

Mixed-precision training with Adam keeps, per parameter: BF16 weights (2 bytes), BF16 gradients (2), an FP32 master copy (4) and Adam’s two moments (4 + 4): 16 bytes per parameter before activations. A 70B model needs 1.1 TB of state: no single GPU holds it. ZeRO / FSDP shards it:

• Stage 1 shards the optimizer state (12 bytes) across data-parallel GPUs.
• Stage 2 also shards gradients.
• Stage 3 (FSDP) shards the weights too, gathering each layer just before it is used.

When even that does not fit, offload moves optimizer state to CPU memory or NVMe; on Grace superchips the 900 GB/s C2C link makes CPU offload much less painful than over PCIe.

Activations: recompute instead of store#

The backward pass needs the forward pass’s activations. Storing them all for long sequences can exceed the weights. Activation checkpointing stores only some and recomputes the rest during the backward pass: spend FLOPs to save HBM. Inside a kernel, FlashAttention applies the same idea at the SRAM level: it never writes the full attention matrix to HBM.

KV cache: the inference memory problem#

During serving the KV cache grows with context length times concurrent users and often exceeds the weights. It now has its own hierarchy: hot blocks in HBM (PagedAttention), warm blocks offloaded to CPU memory, cold or shared prefixes on NVMe or a shared store reachable by every replica (LMCache, NVIDIA Dynamo’s KV block manager, Mooncake-style designs). Reusing a cached prefix saves recomputing it, which is why fast tiers below HBM pay for themselves.

Checkpoints: the burst#

A training job writes its full state every N steps: for a 1T-parameter model, about 14 TB of weights and optimizer state. Writing it must not stall thousands of GPUs, so it is written first to local NVMe or CPU memory asynchronously, then drained to the parallel file system and finally to object storage. How fast you can checkpoint sets how often you can afford to, and so how much work a failure throws away.

Code#

Per-GPU training memory for a 70B model under each ZeRO stage, and how long a checkpoint takes to write at different storage speeds.

// state.go — where a training run's bytes live, per GPU.
package main

import "fmt"

func main() {
	const params = 70e9
	const gb = 1e9
	weights, grads, optim := 2*params, 2*params, 12*params // bytes: BF16, BF16, FP32 master+Adam

	fmt.Println("data-parallel GPUs   ZeRO-0   ZeRO-1   ZeRO-2   ZeRO-3/FSDP   (GB per GPU, before activations)")
	for _, n := range []float64{8, 64, 512} {
		z0 := weights + grads + optim
		z1 := weights + grads + optim/n
		z2 := weights + (grads+optim)/n
		z3 := (weights + grads + optim) / n
		fmt.Printf("%18.0f %8.0f %8.0f %8.0f %13.0f\n", n, z0/gb, z1/gb, z2/gb, z3/gb)
	}

	ckpt := weights + optim // a resumable checkpoint: weights + optimizer state
	fmt.Printf("\ncheckpoint size: %.2f TB\n", ckpt/1e12)
	for _, s := range []struct {
		where string
		gbs   float64
	}{
		{"object storage, 20 GB/s", 20},
		{"parallel file system, 500 GB/s", 500},
		{"local NVMe on 64 nodes, 64 × 30 GB/s", 64 * 30},
	} {
		fmt.Printf("  %-40s %7.1f s\n", s.where, ckpt/gb/s.gbs)
	}
}

Without sharding, every GPU needs 1.1 TB, impossible on any GPU sold today. FSDP across 64 GPUs brings it to about 18 GB each, leaving the rest of HBM for activations. The checkpoint numbers explain the layered write path: almost a minute straight to object storage, two seconds to a parallel file system, half a second to NVMe spread across the nodes.

Remember this#

• Six kinds of data: activations, weights, KV cache, optimizer state, checkpoints, training data.
• Weights and KV cache need HBM to run fast; everything else can be pushed down a tier.
• Adam training needs ~16 bytes per parameter; ZeRO/FSDP shards it, offload pushes it to the CPU.
• Activation checkpointing trades FLOPs for HBM; FlashAttention does the same at the SRAM level.
• KV cache now has its own hierarchy (HBM → CPU → NVMe → shared store).
• Checkpoints go to the fastest tier first and drain downwards.

Try it#

1. Run state.go. Change the model to 405B. How many GPUs do you need under FSDP to stay under 100 GB per GPU?
2. Using lesson 02’s KV numbers, how much CPU memory would hold a day’s worth of popular prefixes for a 70B model if 10,000 distinct 4k-token prefixes recur?
3. On a Grace-Hopper or Grace-Blackwell machine, compare optimizer offload time over C2C (900 GB/s) with PCIe Gen 5 (64 GB/s) for a 70B model’s 840 GB of optimizer state.

Check yourself#

1. What are the 16 bytes per parameter of mixed-precision Adam?
2. What does each ZeRO stage shard?
3. Why do checkpoints go to local NVMe before the parallel file system?
4. Name three places a KV cache block can live, from fastest to slowest.

Sources#

• ZeRO: memory optimizations toward training trillion-parameter models (Rajbhandari et al., 2020)
• PyTorch FSDP (Zhao et al., 2023)
• Efficient memory management for LLM serving with PagedAttention (Kwon et al., 2023)
• Mooncake: a KVCache-centric architecture for LLM serving (2024)
• NVIDIA GPUDirect Storage