AI Knowledge Tree

Training infrastructure overview

Survey the cluster pieces needed to train modern models: storage, schedulers, NCCL / collective comms, and observability.

⏱ 22 min

Training a frontier model is half algorithm and half infrastructure. The optimizer, the model code, and the data pipeline only work if the underlying cluster of thousands of GPUs cooperates: high-bandwidth interconnect, fast storage, robust schedulers, instant failure detection. This lesson surveys the parts of training infrastructure that algorithms papers usually take for granted.

The components

A modern training cluster has roughly seven parts:

  1. Compute: GPUs / TPUs / specialized accelerators.
  2. Interconnect: NVLink within nodes, InfiniBand / Ethernet between nodes.
  3. Storage: shared filesystems for datasets and checkpoints.
  4. Collective communications: libraries (NCCL, RCCL, XLA) that coordinate across devices.
  5. Schedulers: SLURM, Kubernetes, custom systems that allocate jobs to hardware.
  6. Observability: monitoring, logging, alerting on training health.
  7. Fault tolerance: checkpointing, retry / resume.

Each one can break a training run. The job of an ML platform team is to make all of them invisible to the model developer.

Interconnect

The performance of distributed training depends on how fast gradients can move.

  • NVLink (within a node, 8 GPUs sharing a high-bandwidth fabric): ~600 GB/s per GPU. Fast enough that tensor parallelism within a node is cheap.
  • InfiniBand (between nodes, 200-400 Gbit/s typical, up to 800 Gbit/s on recent setups): much slower than NVLink. Communication here is often the bottleneck.
  • Ethernet RoCE (RDMA over Converged Ethernet): cheaper alternative to InfiniBand. Used by some clouds.

The all-reduce bandwidth across the cluster sets the maximum useful batch size: above that, communication time exceeds compute time.

Storage

Datasets are huge (LAION-5B is 240 TB; The Pile is 800 GB; Common Crawl is petabytes). Training reads through them many times.

Patterns:

  • Sharded files (WebDataset format): each shard is a tar with a few thousand examples. Streamable, parallel-readable.
  • Object stores (S3, GCS): for the largest datasets. With careful prefetching, throughput can match local SSDs.
  • Local NVMe scratch: each node has fast local storage for the current chunk of data and recent checkpoints.

Checkpoints are also gigantic: a 7B model's checkpoint is 28 GB at fp32, but with optimizer state and shards from FSDP, the per-step checkpoint can be 100 GB+. Saving every KK steps to remote storage takes time; staggered saves and async I/O are standard.

NCCL and collective communications

NCCL (NVIDIA Collective Communications Library) provides implementations of:

  • All-reduce: combine values from every GPU and distribute the result back. The workhorse of data parallelism.
  • All-gather: collect values from every GPU into a single concatenated tensor. Used by FSDP.
  • Reduce-scatter: sum and shard. Used by FSDP backward.
  • All-to-all: every GPU sends a portion of its data to every other GPU. Used by expert parallelism and sequence parallelism.

NCCL chooses algorithms based on topology: ring all-reduce on NVLink, tree reductions across nodes, hierarchical algorithms for cluster-wide operations.

PyTorch, JAX, TensorFlow all wrap NCCL transparently[1].

Schedulers

Cluster management:

  • SLURM: the standard for academic supercomputers. Mature, batch- oriented.
  • Kubernetes (with Kubeflow, Volcano): general-purpose container orchestration. Used at many companies.
  • Custom systems: every hyperscaler has their own (Borg at Google, Tupperware at Meta, etc.). Closer integration with quotas and priority.

The scheduler decides which job runs where, when to preempt, how to restart. For long training runs (weeks), preemptions are inevitable and the system must resume from the last checkpoint.

Observability

A training job at scale produces:

  • Per-step loss and metrics (logged to TensorBoard, WandB, or a custom dashboard).
  • System metrics: GPU utilization, memory, network throughput, storage I/O.
  • Anomaly detection: NaN losses, gradient norm spikes, loss plateaus.
  • Profiles: occasional snapshots of where time is spent.

Frontier teams have dedicated dashboards and on-call rotations for training jobs. A training run that wastes 10% of compute by going unnoticed is millions of dollars in wasted compute.

Fault tolerance

Hardware fails:

  • GPU memory errors: silent or noisy.
  • Node failures: power, networking, kernel panics.
  • Storage hiccups: file system stalls, S3 errors.

The training pipeline must detect failures and resume from the most recent checkpoint without losing too much work. Standard checkpoint cadence: every 100-1000 steps (a few minutes of compute), depending on the failure rate of the cluster.

Elastic training (PyTorch, JAX): jobs can scale up or down at runtime as nodes come and go[1].

Cost considerations

Modern training is expensive:

  • A 7B model: ~$10K-100K on cloud.
  • A 70B model: ~$1M-10M.
  • Frontier models (GPT-4-class): tens to hundreds of millions of dollars per training run.

Most of the cost is compute (GPU-hours). Storage, networking, and ops are smaller line items.

Spot / preemptible instances (50-70% cheaper than on-demand) are heavily used with checkpoint-resume to handle preemptions.

What this means for ML practitioners

You almost never interact with this layer directly as a model developer:

  • Hyperscalers and labs have entire teams owning each piece.
  • Open-source training stacks (Megatron-DeepSpeed, NeMo, MaxText) wrap most of the complexity.
  • Cloud services (AWS Sagemaker Hyperpod, GCP Vertex Training, etc.) provide managed clusters.

But knowing the pieces helps you:

  • Choose the right scaling strategy for your job size.
  • Diagnose performance problems ("why is my GPU at 50% utilization?").
  • Avoid pathological costs (write rate to checkpoint storage, fall-back to TCP from RDMA, etc.).

Exercises

Exercise 1. A cluster has 512 GPUs. Each data-parallel step must all-reduce a 28 GB gradient buffer (per GPU, after the (G−1)/G(G-1)/G ring factor this is the bytes each GPU moves). The compute per step takes 900 ms. Compare the step time when the all-reduce runs over NVLink-class bandwidth (600 GB/s) versus inter-node InfiniBand (200 Gbit/s), and state the implication for batch size.

Solution

Step 1. NVLink path: 28 GB/600 GB/s≈0.04728 \text{ GB} / 600 \text{ GB/s} \approx 0.047 s =47= 47 ms. Step time is max⁡\max/overlap of compute and comm; even fully serialized it is 900+47=947900 + 47 = 947 ms, so communication is about 5% overhead.

Step 2. InfiniBand at 200 Gbit/s is 200/8=25200/8 = 25 GB/s. The same 28 GB takes 28/25=1.1228 / 25 = 1.12 s =1120= 1120 ms, longer than the 900 ms of compute.

Step 3. When all-reduce time exceeds compute time, adding more data-parallel work no longer scales: communication dominates. The cluster's all-reduce bandwidth sets a ceiling on useful global batch size, which is why gradient buffers are kept off the slow inter-node fabric where possible (large TP groups inside nodes, gradient compression, overlap of comm with backward).

Exercise 2. A 70B model is checkpointed with FSDP. The checkpoint holds fp32 master weights (4 bytes/param) plus two fp32 Adam moments (8 bytes/param). Compute the checkpoint size. If the run saves every 500 steps, each step takes 1.2 s, and the remote write sustains 5 GB/s, find the write time and the fraction of a save interval it consumes. Comment on async I/O.

Solution

Step 1. Bytes per parameter in the checkpoint: 4+8=124 + 8 = 12. Total: 12×70×109=840×10912 \times 70 \times 10^9 = 840 \times 10^9 bytes =840= 840 GB.

Step 2. Write time at 5 GB/s: 840/5=168840 / 5 = 168 s.

Step 3. A save interval is 500×1.2=600500 \times 1.2 = 600 s of compute. A synchronous 168 s write would stall the run for 168/600=28%168/600 = 28\% of the interval, a large waste. Async/staggered I/O overlaps the write with continued training (snapshot to local NVMe fast, then stream to remote in the background), so the 168 s is hidden behind compute instead of blocking it.

Exercise 3. A 512-GPU run has a mean time between node failures of 18 hours. Restarting from a checkpoint loses all steps since the last save plus a fixed 10-minute reload. The team weighs checkpointing every 30 minutes versus every 5 minutes. Estimate the expected wasted time per failure for each cadence and state the tradeoff that picks the cadence.

Solution

Step 1. On a failure, expected lost compute is on average half the checkpoint interval (the failure lands uniformly in the interval), plus the 10-minute reload.

Step 2. 30-minute cadence: expected loss ≈15+10=25\approx 15 + 10 = 25 minutes per failure. 5-minute cadence: expected loss ≈2.5+10=12.5\approx 2.5 + 10 = 12.5 minutes per failure.

Step 3. Frequent checkpointing cuts recompute loss but adds steady write overhead (Exercise 2: each save costs I/O time and bandwidth). The cadence balances expected failure loss against per-save cost: with an 18-hour MTBF and cheap async saves, the 5-minute cadence's lower expected loss usually wins; if each save were expensive or failures rare, the longer interval would be cheaper overall.

This closes pillar D

D ran from the basic MLP, through CNNs and Transformers, into the modern landscape of generative models, self-supervision, and graph networks, finishing with the training-systems infrastructure that makes large-scale deep learning practical.

The thread: each subsection deepens the toolkit for a different data modality or problem class, but they all share a small set of building blocks (residual connections, attention, normalization, gradient-based optimization). The next pillar (E) covers NLP and language models specifically: a vertical application of D's foundations.


References

  1. [1]Paszke, A., et al. (2019). "PyTorch: an imperative style, high-performance deep learning library". NeurIPS. arXiv:1912.01703. The framework whose distributed and elastic-training APIs wrap NCCL and the cluster collectives surveyed here.↩