System Design Lab

LLM pretraining is bounded by GPU memory and inter-GPU communication; both force you to split the model across many accelerators.

Scale model parameters, GPU count, global batch, sequence length, interconnect bandwidth, and checkpoint interval. The design moves from a single GPU to data parallelism with all-reduce, to tensor and pipeline parallelism for a model too big for one device, to full 3D parallelism with activation checkpointing, and finally to thousands of GPUs that demand fast collectives plus frequent checkpoint-and-recover.

Read the full tutorial

Guided walkthrough

Reason about it one step at a time

Normal evolution scenarios

Click left to right for the intended demo path. Each card changes the workload inputs.

Workload

These are inputs, not preset architecture stages.

Model parameters Total trainable parameters. Weights + gradients + optimizer state set the per-device memory floor. GPU count Accelerators in the training mesh, factored into data x tensor x pipeline parallel groups. Global batch Tokens processed per optimizer step across the whole mesh. Larger batches need more activation memory and data throughput. Sequence length Context window in tokens. Activation memory scales with sequence length. Interconnect bandwidth Per-GPU collective bandwidth (NVLink within a node, InfiniBand across nodes). Gradient all-reduce and tensor-parallel traffic ride on it. Data pipeline throughput Tokens per second the data loaders can stream (tokenized shards from object storage). Must keep the mesh fed. Checkpoint interval Minutes of compute between full checkpoints. Shorter loses less work on a failure but costs more store bandwidth. Tensor + pipeline parallel (3D) Split each layer across GPUs (tensor) and split layers into stages (pipeline) so a model larger than one device can train. Activation checkpointing Recompute activations in the backward pass instead of storing them, trading extra compute for much lower memory.

Recommended shape

Current architecture path

Dataset

Token corpus tokenized training shards held in object storage and streamed during the run

Data loaders

Data loaders prefetch, shuffle, and pack token batches so GPUs never wait on input

Training mesh

Data-parallel replicas each train on a slice of the batch and all-reduce gradients each step

Tensor-parallel splits each layer matrix across GPUs that exchange activations within a node

Pipeline-parallel assigns contiguous layer stages to GPUs and streams micro-batches between them

Comms + checkpoint

Collective fabric NVLink and InfiniBand carry gradient all-reduce and tensor-parallel exchanges

Checkpoint store persists sharded weights and optimizer state so a run can resume after a failure

Ops

Metrics logs loss, throughput, and hardware health off the training hot path

Fault recovery detects dead nodes and restarts the mesh from the latest checkpoint

Bottlenecks

Per-GPU memory pressure

Collective comms load

Data pipeline load

Activation memory

Checkpoint / failure exposure

Why this changes

Decision tradeoffs

Parallelism strategy

Gradient sync / collectives

Activation memory

Checkpoint + fault tolerance

Data pipeline throughput

Interconnect topology

Source-backed rules

These are the durable system-design claims behind the model. The exact slider thresholds are deliberately labeled as teaching assumptions.

Verified rule

Tensor parallelism splits each layer across GPUs to train models too big for one device

Megatron-LM partitions the attention and MLP matrices across GPUs so a single transformer layer fits in aggregate device memory, exchanging activations over the fast intra-node interconnect.

Megatron-LM

Verified rule

ZeRO shards optimizer state, gradients, and parameters to cut the per-GPU memory floor

Standard data parallelism replicates the full ~16 bytes/param of optimizer state on every GPU; ZeRO partitions that state across data-parallel ranks so memory falls roughly linearly with the number of devices.

DeepSpeed ZeRO (arXiv:1910.02054)

Verified rule

Pipeline parallelism splits layers into stages and overlaps micro-batches to keep GPUs busy

GPipe assigns contiguous layer stages to different accelerators and pipelines micro-batches through them, with re-materialization (activation recompute) to bound the activation memory per stage.

GPipe (arXiv:1811.06965)

Verified rule

Combining tensor, pipeline, and data parallelism scales transformer training to thousands of GPUs

Efficient large-model training composes the three parallelism axes (3D parallelism), placing high-bandwidth tensor-parallel traffic within a node and lower-bandwidth data-parallel all-reduce across nodes.

Megatron-LM scaling (arXiv:1909.08053)

Teaching assumptions

Memory is modeled as ~16 bytes/param for weights, gradients, and Adam optimizer state, plus an activation term that grows with batch and sequence length; real frameworks vary with precision and ZeRO sharding.
Per-GPU memory is a fixed teaching figure (80 GB) and collective/data thresholds are round numbers, not vendor benchmarks.
Failure rate is approximated as proportional to GPU count, so expected lost work scales with both fleet size and checkpoint interval.