Bhaskar
← All writing
first principlesAug 19, 2026· 41 min read

From One GPU to Distributed Language-Model Training

A first-principles guide to data, memory, parallelism, networking, recovery, evaluation, and release

Current-source note: the implementation references in this article use the PyTorch stable documentation available on August 19, 2026. Distributed APIs continue to evolve; verify signatures against the linked documentation before adopting the code in a production run.

Training a language model is easy to describe: show it tokens, ask it to predict the next token, measure how wrong it was, and adjust its parameters. Training infrastructure exists because this simple procedure eventually collides with four limits:

  1. The data no longer fits on one machine or cannot be read fast enough.
  2. The model, optimizer state, and activations no longer fit on one GPU.
  3. One GPU cannot perform the required computation before the deadline.
  4. A long, distributed run is likely to fail before it finishes.

Every major training technique is a response to one of those failures. Mixed precision changes the representation of numbers. Activation checkpointing exchanges computation for memory. Data parallelism exchanges communication for throughput. Fully sharded data parallelism exchanges more frequent communication for lower per-GPU state. Tensor, pipeline, and context parallelism divide different axes of the computation. Distributed checkpointing converts inevitable worker failures from run-ending events into bounded delays.

The aim of this article is to derive those mechanisms from the smallest possible training loop, not to hide them behind a trainer.

1. The smallest system: one sequence, one model, one machine

Assume that our corpus is the string:

the cat sat

A character tokenizer might assign:

" " -> 0, "a" -> 1, "c" -> 2, "e" -> 3,
"h" -> 4, "s" -> 5, "t" -> 6

The text becomes a sequence of integers. If the context length is four, one training example can be constructed by shifting the same token window by one position:

input:  [t, h, e,  ]
target: [h, e,  , c]

The model is not asked to reproduce the entire target at once. At each position it produces a vector of unnormalized scores—logits—over the vocabulary. The target at that position tells us which logit should become relatively larger.

1.1 The language-model training objective

For tokens (x_1,\ldots,x_T), an autoregressive language model represents

[ p(x_1,\ldots,x_T)=\prod_{t=1}^{T}p(x_t\mid x_{<t}). ]

Training minimizes average negative log-likelihood, usually implemented as token-level cross-entropy:

[ \mathcal{L}=-\frac{1}{N}\sum_{i=1}^{N}\log p_\theta(y_i\mid x_{i,<t}). ]

Here, (N) must mean the number of valid target tokens, not necessarily batch size times sequence length. Padding, masked prefixes, or ignored document-boundary positions must not enter the denominator. PyTorch's CrossEntropyLoss combines log_softmax and negative log-likelihood and accepts logits directly; applying softmax first is both unnecessary and less numerically stable.

Perplexity is (\exp(\mathcal{L})) when the loss is natural-log cross-entropy over the same tokenizer. It is interpretable as an effective branching factor, but perplexities from different tokenizers are not directly comparable because the units differ.

1.2 Tokenized datasets, input/target sequences, and batches

A tokenized dataset is not merely a large integer array. It also needs:

  • tokenizer identity and version;
  • document boundaries and source provenance;
  • train, validation, and test assignments;
  • a policy for end-of-document tokens;
  • packing and padding rules;
  • dtype and binary format;
  • enough metadata to reproduce the exact sample order.

For a decoder-only model, inputs and targets are usually the same packed stream offset by one. A batch has conceptual shape [batch, sequence]. The model produces logits of shape [batch, sequence, vocabulary].

Packing multiple short documents into a fixed-length sequence improves utilization. But the attention mask determines whether tokens in one document may attend into the previous document. Allowing that attention is computationally efficient and common in pretraining, but it changes the task: the model observes artificial adjacency. Segment-aware masks avoid that leakage at some implementation cost.

1.3 Forward pass, loss, backpropagation, and optimizer

The forward pass transforms token IDs into logits:

  1. embedding lookup maps token IDs to vectors;
  2. positional information is added or applied;
  3. transformer blocks mix information across positions and channels;
  4. an output projection maps hidden vectors to vocabulary logits.

The loss compresses all prediction errors into a scalar. Backpropagation applies the chain rule in reverse through the operations recorded during the forward pass. It computes a gradient for each trainable parameter:

[ g_t=\nabla_\theta \mathcal{L}(\theta_t). ]

An optimizer converts that gradient into an update. Plain stochastic gradient descent uses

[ \theta_{t+1}=\theta_t-\eta_t g_t. ]

Adam additionally maintains moving estimates of the first and second moments of gradients. That often improves optimization, but those moment tensors become a major memory cost later.

The learning rate (\eta_t) is rarely constant. A typical schedule has:

  • warm-up, preventing unstable early updates while optimizer statistics and activations are poorly calibrated;
  • a peak learning rate;
  • cosine or linear decay toward a smaller terminal value.

The scheduler advances on optimizer updates, not on microbatches. Confusing these clocks silently changes the schedule when gradient accumulation is introduced.

1.4 The irreducible PyTorch training loop

model.train()

for step, (input_ids, targets) in enumerate(loader):
    input_ids = input_ids.to(device)
    targets = targets.to(device)

    optimizer.zero_grad(set_to_none=True)
    logits = model(input_ids)                         # forward
    loss = F.cross_entropy(
        logits.reshape(-1, logits.size(-1)),
        targets.reshape(-1),
    )
    loss.backward()                                   # gradients
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()                                  # parameters change
    scheduler.step()                                  # update clock advances

This loop contains the semantics that every larger system must preserve. Distributed training is correct only if its collective operations, accumulation, sharding, and recovery produce the intended effective update.

Experiment 1: train a character-level model

Use a small text file, a character vocabulary, contexts of 64–256 tokens, and an embedding-plus-transformer model small enough for CPU or one GPU. Record:

  • initial and final train loss;
  • validation loss on a document-disjoint split;
  • generated samples from a fixed prompt and fixed sampling seed;
  • tokens processed and tokens per second;
  • parameter count.

The success criterion is not fluent prose. It is that the training loss falls, held-out loss improves before eventually overfitting, and checkpointed runs reproduce the same next update within the expected numerical tolerance.

2. The first scaling failure: the desired batch no longer fits

Larger batches can improve hardware utilization and reduce gradient noise, but a batch stores activations for every included token. Eventually the microbatch that fits in GPU memory is smaller than the batch needed for the optimization experiment.

2.1 Gradient accumulation

Split one logical batch into (A) microbatches. Backpropagate each microbatch without stepping the optimizer, then update once:

optimizer.zero_grad(set_to_none=True)

for micro_step in range(accumulation_steps):
    x, y = next_batch()
    logits = model(x)
    loss = F.cross_entropy(
        logits.flatten(0, 1),
        y.flatten(),
        reduction="mean",
    )
    (loss / accumulation_steps).backward()

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
scheduler.step()

For equal-size microbatches with the same number of valid tokens, dividing each mean loss by (A) produces the mean gradient of the logical batch. For variable-length or differently masked microbatches, it does not. The correct method is to accumulate loss sums and normalize gradients by the total number of valid tokens.

The effective global batch in tokens is approximately

[ B_{tokens}=B_{micro/GPU}\times S\times A\times D, ]

where (S) is valid tokens per sequence and (D) is the data-parallel world size.

Accumulation is not perfectly equivalent to one large physical batch when the model contains batch-dependent operations, randomness differs, gradient clipping occurs per microbatch, or floating-point summation order changes. Standard transformers generally avoid BatchNorm, which removes one important discrepancy.

2.2 Gradient clipping

Exploding gradients can turn a recoverable bad batch into NaNs. Global norm clipping computes

[ \tilde g = g\cdot\min\left(1, \frac{c}{\lVert g\rVert_2+\epsilon}\right). ]

Clipping is a safety rail, not a cure for a consistently unstable learning rate or corrupt inputs. Track both the pre-clip norm and the fraction of steps clipped. If nearly every step clips, the optimizer is operating under a different effective update rule than intended.

With gradient accumulation, clip once after the logical gradient is complete. With FP16 loss scaling, unscale first and then clip.

Experiment 2: prove accumulation correctness

From identical model and RNG states, compare:

  • one batch of eight sequences;
  • four microbatches of two sequences with loss divided correctly.

Disable dropout for the comparison. Compare each parameter after one optimizer update. Then repeat with unequal padding and demonstrate why normalizing each microbatch independently gives the wrong token-weighted gradient.

3. The second failure: arithmetic is slow and memory is scarce

Modern accelerators execute lower-precision matrix multiplications much faster than FP32. The challenge is to use low precision for operations that tolerate it while retaining enough range and accuracy for stable learning.

3.1 FP32, FP16, and BF16

Format Bits Exponent bits Fraction bits Practical consequence
FP32 32 8 23 Broad range and high precision; expensive in memory and often lower tensor-core throughput
FP16 16 5 10 More precision than BF16 near 1, but narrow range; small gradients can underflow and large values can overflow
BF16 16 8 7 FP32-like exponent range with lower precision; commonly easier for large-model training

Mixed precision does not mean every tensor is forced to one dtype. PyTorch automatic mixed precision uses torch.autocast to select lower precision for eligible operations while keeping numerically sensitive operations in a safer dtype. FP16 training commonly uses GradScaler to multiply the loss before backpropagation, preventing small gradients from becoming zero, and then unscales before the optimizer update. BF16 usually does not require loss scaling because it has FP32's exponent range.

scaler = torch.amp.GradScaler("cuda", enabled=use_fp16)

optimizer.zero_grad(set_to_none=True)
with torch.autocast("cuda", dtype=amp_dtype):
    logits = model(x)
    loss = F.cross_entropy(logits.flatten(0, 1), y.flatten())

scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(optimizer)
scaler.update()

Monitor scaler reductions and skipped steps in FP16. Repeated scale collapse signals overflow. For either FP16 or BF16, monitor finite loss, gradients, parameters, and optimizer state; a finite scalar loss does not prove every tensor is finite.

3.2 Mixed precision changes speed, not just capacity

Its benefits can include:

  • smaller activations and communication payloads;
  • faster tensor-core matrix multiplication;
  • more arithmetic per memory transaction;
  • capacity for a larger microbatch, which may improve GPU utilization.

But the optimizer may still keep FP32 master parameters or FP32 moment estimates. Therefore the total memory reduction is smaller than “32 bits to 16 bits” suggests.

Experiment 3: measure precision rather than assume it

Run the same short training job in FP32, BF16, and FP16 where supported. Compare:

  • loss curve and validation loss;
  • tokens per second;
  • maximum allocated and reserved GPU memory;
  • number of non-finite or skipped optimizer steps;
  • gradient norm distribution.

Use a meaningful warm-up before timing because compilation, allocator growth, and kernel initialization distort the first iterations.

4. The memory ledger: what actually occupies a GPU

An out-of-memory error is not explained by parameter count alone. Peak training memory includes:

[ M_{peak}\approx M_{params}+M_{grads}+M_{optimizer}+M_{activations}+M_{temp}+M_{allocator}. ]

4.1 Model parameters, gradients, and optimizer state

For (P) trainable parameters, a rough, implementation-dependent ledger for FP32 Adam is:

  • parameters: (4P) bytes;
  • gradients: (4P) bytes;
  • first moment: (4P) bytes;
  • second moment: (4P) bytes.

That is roughly 16 bytes per parameter, before activations, temporary buffers, fragmentation, and framework bookkeeping.

Mixed FP16/BF16 Adam may contain a low-precision parameter copy, low-precision gradients, FP32 master weights, and FP32 moments. Depending on optimizer and framework implementation, model-state memory can remain near 16 bytes per parameter or differ materially. Do not budget from folklore; inspect actual tensor dtypes and measure peak memory.

4.2 Activations

Backpropagation requires intermediate values from the forward pass. Their memory grows with microbatch size, context length, hidden size, and layer count. Naive attention also materializes score or probability tensors proportional to (S^2), where (S) is sequence length. Memory-efficient attention kernels can avoid storing the full matrix, but the arithmetic cost of dense attention remains quadratic.

This distinction matters:

  • parameter-state pressure grows mainly with model size;
  • activation pressure grows mainly with batch, depth, width, and context;
  • attention pressure becomes especially severe as context grows.

4.3 Temporary and allocator memory

Kernels need workspaces. Collective communication can require buckets or gathered parameters. The CUDA caching allocator reserves blocks that may exceed live tensor allocations. Consequently:

  • memory_allocated describes live tensor memory;
  • memory_reserved describes memory held by the allocator;
  • the driver view may include additional context and library allocations.

PyTorch provides peak counters and CUDA memory snapshots for allocator-level diagnosis. A snapshot can distinguish a true live-tensor peak from fragmentation or an accidental retained graph.

4.4 Activation checkpointing

If activations dominate, store only selected boundary activations during the forward pass. During backward, rerun the omitted forward subgraph to regenerate them. This is also called activation recomputation or rematerialization.

The trade is explicit:

[ \text{less activation memory} \quad\Longleftrightarrow\quad \text{more forward computation during backward}. ]

Checkpoint transformer blocks rather than tiny operations; otherwise launch and framework overhead can dominate. Checkpointing functions with randomness requires correct RNG-state handling. PyTorch's torch.utils.checkpoint documentation warns that moving tensors to unanticipated devices inside the checkpointed region can make deterministic equivalence impossible.

Checkpointing does not reduce parameter or optimizer state. If those dominate, sharding is the relevant mechanism.

Experiment 4: construct the memory equation

For a small transformer, independently sweep:

  • model depth while holding batch and context fixed;
  • context length while holding model and batch fixed;
  • microbatch size while holding model and context fixed;
  • checkpointing off versus every block versus every few blocks.

Record peak allocated memory and step time. Fit a simple empirical model to the measurements. The aim is not a universal formula; it is to learn which term dominates your architecture and kernel stack.

5. The data pipeline becomes a distributed system

Once the corpus is larger than local storage or preprocessing cannot keep up with the GPU, data engineering becomes part of model training correctness.

5.1 Dataset streaming

A map-style dataset assumes random access by index. A streaming or iterable dataset reads sequential shards from local or remote storage and may decode or tokenize on the fly. PyTorch distinguishes map-style and iterable-style datasets in torch.utils.data.

Streaming solves storage pressure, but it creates new requirements:

  • partition shards among ranks and worker processes without duplication;
  • prefetch far enough to hide storage latency;
  • cache hot shards on node-local NVMe;
  • bound the shuffle buffer;
  • store resumable cursor state;
  • tolerate corrupt records without silently changing global sample counts;
  • measure data wait time separately from GPU computation.

If GPU utilization falls while kernels are fast, the bottleneck may be object storage, decompression, tokenization, Python workers, or host-to-device copies.

5.2 Shuffling at scale

An exact random permutation of trillions of tokens is expensive to materialize. Practical systems randomize at several levels:

  1. shuffle shard order per epoch;
  2. assign disjoint shards to ranks;
  3. shuffle records or packed sequences within a bounded buffer;
  4. vary seeds by epoch while preserving a resumable global order.

The quality question is whether correlated examples arrive close enough together to bias optimization. The correctness question is whether samples are duplicated or omitted across ranks. Log stable sample IDs for a small audit window and test partition logic independently.

5.3 Deduplication

Duplicate text changes the intended sampling distribution, wastes compute, can amplify memorization, and can leak validation or benchmark data into training. Three useful levels are:

  • exact document or content-hash deduplication;
  • near-duplicate document detection, often using shingles and MinHash/LSH;
  • repeated-substring detection for boilerplate copied across otherwise different pages.

Deduplicate across the complete mixture, not only inside each source. Also deduplicate or decontaminate across train, validation, and evaluation boundaries. Public pipelines such as FineWeb demonstrate that extraction, filtering, and deduplication choices materially alter the resulting corpus.

Deduplication is not free: aggressive thresholds can erase legitimate repeated structures, minority-language data, templates that teach useful syntax, or independently corroborated facts. Treat the threshold as an experimental variable and retain provenance explaining why each document was removed.

5.4 Data filtering

Filtering may remove malformed pages, navigation boilerplate, spam, unsafe material, personally sensitive content, disallowed licenses, language mismatches, or low-information text. Every filter changes the model's learned distribution. A “quality classifier” can also encode cultural, domain, and stylistic preferences that reduce coverage.

The honest test is downstream: train matched small models on unfiltered and differently filtered mixtures at equal token and compute budgets. Recent work continues to challenge universal filtering rules; for example, the 2026 preprint A Bitter Lesson for Data Filtering reports a high-compute, data-scarce regime in which aggressive filtering can lose to retaining more data. That result is not permission to ingest indiscriminately—it is evidence that quality, diversity, repetition, legality, safety, and compute interact, so filtering policies require controlled ablations.

5.5 Tokenizer training

Character tokenization has a tiny vocabulary but creates long sequences. Word tokenization creates an unbounded vocabulary and brittle unknown words. Subword tokenizers balance these extremes.

BPE repeatedly merges frequent adjacent units; Unigram begins with candidates and removes units under a probabilistic objective. SentencePiece can train directly on raw text and supports BPE and Unigram-style models.

A tokenizer determines:

  • tokens per byte or word—often called fertility;
  • effective text that fits into a fixed context window;
  • embedding/output-matrix size;
  • representation quality across languages, code, numbers, and domain terminology;
  • behavior around whitespace, normalization, invalid bytes, and special tokens.

Train the tokenizer on a representative, deliberately sampled subset rather than the easiest source to access. Reserve and document BOS, EOS, padding, unknown, and any control tokens. Freeze the tokenizer before tokenizing the full corpus; changing it changes every model input and invalidates checkpoints.

Experiment 5: make data quality measurable

Build three small corpora: raw, deduplicated, and deduplicated-plus-filtered. Train identical tokenizers and identical small transformers at the same token budget. Compare:

  • tokenizer fertility by language/domain;
  • train and validation loss;
  • memorization probes on repeated strings;
  • evaluation performance by source group;
  • data-pipeline throughput and rejection counts.

The data pipeline should emit a versioned manifest containing source identifiers, transformations, counts, hashes, licenses, and split rules.

6. More GPUs for throughput: data parallelism and DDP

Suppose the model fits on one GPU, but one GPU would take a year to process the token budget. The simplest solution is to replicate the model and divide the batch.

6.1 Data parallelism

With (D) workers:

  1. every GPU holds the same model and optimizer structure;
  2. each GPU processes a different microbatch;
  3. every GPU computes local gradients;
  4. gradients are reduced across workers;
  5. every optimizer applies the same update.

Synchronous data parallelism preserves one logical model. If gradients are averaged, the result corresponds approximately to the gradient of the union of the rank-local batches.

6.2 Distributed Data Parallel

PyTorch DistributedDataParallel runs one process per GPU and synchronizes gradient buckets. As backward computes gradients, ready buckets can begin asynchronous all-reduce, overlapping communication with the remaining backward computation.

dist.init_process_group(backend="nccl")
torch.cuda.set_device(local_rank)

model = TinyGPT(config).to(local_rank)
model = DDP(model, device_ids=[local_rank])

sampler = DistributedSampler(dataset, shuffle=True)
loader = DataLoader(dataset, sampler=sampler, ...)

for epoch in range(num_epochs):
    sampler.set_epoch(epoch)
    for x, y in loader:
        optimizer.zero_grad(set_to_none=True)
        loss = compute_loss(model, x, y)
        loss.backward()       # DDP hooks reduce gradient buckets
        optimizer.step()

During gradient accumulation, avoid unnecessary synchronization on intermediate microbatches by using DDP's no_sync() context, then synchronize on the final microbatch.

For ring all-reduce, each rank transfers roughly

[ 2\frac{D-1}{D}G ]

bytes for a gradient payload of size (G), ignoring protocol overhead. This approaches (2G) per rank as the world grows. More GPUs therefore do not reduce replicated parameter, gradient, or optimizer memory; they add communication and expand the global batch.

6.3 Correctness when global batch changes

Adding GPUs without reducing per-GPU batch multiplies the global batch. That changes gradient noise, number of optimizer updates per token, warm-up length, and often the suitable learning rate. “Same epochs” is not a stable comparison when token counts or packing differ. Use tokens and optimizer updates as explicit clocks.

Experiment 6: measure scaling efficiency

Run the same model on one, two, four, and eight GPUs, keeping the global batch constant first and then keeping the per-GPU batch constant. Measure:

[ \text{speedup}(D)=\frac{T_1}{T_D}, \qquad \text{efficiency}(D)=\frac{\text{speedup}(D)}{D}. ]

Profile compute, exposed communication, and input wait. Compare a parameter tensor after one controlled update against the single-GPU reference. A fast job that silently samples duplicate data or averages gradients incorrectly is not a successful scale-up.

7. The model state no longer fits: fully sharded data parallelism

DDP replicates everything. If parameters, gradients, or optimizer states exceed per-GPU memory, those states must be partitioned.

7.1 Parameter, gradient, and optimizer-state sharding

The ZeRO formulation distinguishes three cumulative stages:

Strategy Parameters Gradients Optimizer state Main benefit Main new cost
DDP replicated replicated replicated simple, strong throughput no model-state memory reduction
Optimizer sharding replicated replicated sharded removes redundant Adam moments optimizer coordination
+ gradient sharding replicated sharded sharded removes redundant gradient storage reduce-scatter-style gradient flow
+ parameter sharding / full shard sharded at rest sharded sharded model-state memory approaches (1/D) parameter all-gathers and tighter scheduling

The original ZeRO paper describes this elimination of data-parallel memory redundancy. PyTorch's newer composable fully_shard API is commonly called FSDP2; it represents sharded parameters using DTensor and all-gathers parameters before computation, then frees or reshards them afterward.

7.2 The FSDP execution pattern

For each sharded module:

  1. all-gather parameters so every rank in the shard group can run the layer;
  2. compute forward;
  3. release or reshard full parameters;
  4. all-gather again when backward needs the parameters, unless they were retained;
  5. compute local gradients;
  6. reduce-scatter gradients, leaving each rank with its owned shard;
  7. update only local optimizer shards.

This lowers persistent state but creates transient full-parameter buffers and more latency-sensitive communication. Wrapping granularity matters. A single giant unit creates a large peak gather; extremely small units create many collectives. Prefetching can overlap the next gather with current compute but increases peak memory.

Sharding degree need not equal total GPU count. A hybrid design may shard within a node and replicate across nodes, trading more memory for avoiding fine-grained parameter gathers over slower inter-node links.

7.3 FSDP is not tensor parallelism

FSDP reconstructs a complete layer on each participating rank for computation, then shards it at rest. Tensor parallelism keeps a single layer's matrix computation partitioned while it executes. The former attacks redundant state; the latter attacks the inability to compute a layer on one device and can reduce per-device matmul dimensions.

Experiment 7: identify the real sharding win

Train the same model with DDP, optimizer-state sharding, and full sharding. Keep global tokens per update fixed. Record:

  • persistent and peak memory;
  • parameter all-gather and gradient reduce-scatter time;
  • step time and tokens per second;
  • checkpoint size written per rank;
  • final loss and parameter equivalence after a controlled step.

Test several wrapping granularities. The configuration with the lowest memory is not necessarily the configuration with the best tokens per dollar.

8. One layer or one sequence no longer fits: model parallelism

Sharding state at rest is insufficient when the live computation itself is too large. We then divide the model's work.

8.1 Tensor parallelism

Consider a linear layer (Y=XW). Split (W) by columns across (T) GPUs. Each GPU computes part of (Y); a later operation gathers or reduces partial results. Transformer MLP and attention projections can be arranged in complementary column- and row-parallel forms so only a small number of collectives occurs per block. Megatron-LM established this practical intra-layer pattern.

Tensor parallelism reduces per-GPU parameter and activation work inside each layer, but invokes collectives frequently—often every transformer block. It therefore prefers the fastest, lowest-latency fabric, usually GPUs connected by NVLink/NVSwitch inside a node. Extending a wide tensor-parallel group across ordinary inter-node networking can make communication dominate.

8.2 Pipeline parallelism

Split the network by depth: early layers run on stage 0, later layers on subsequent stages. Stage boundaries send activations forward and activation gradients backward. If a whole batch passes through one stage at a time, most stages sit idle. Split the batch into microbatches and pipeline them.

With (p) stages and (m) microbatches, insufficient (m) leaves a pipeline bubble. Larger (m) improves utilization but changes activation residency, communication frequency, and scheduling overhead. Stage imbalance is equally important: throughput is constrained by the slowest stage. GPipe derives synchronous microbatch pipelining; PyTorch's current pipeline package separates model partitioning, stage runtime, and schedules.

Pipeline parallelism communicates boundary activations rather than every layer's partial matmul. It can work across nodes better than fine-grained tensor parallelism, but partitioning, bubbles, and failure recovery are more complex.

8.3 Context parallelism

Long contexts can make per-sequence activations and attention computation too large even after parameter sharding. Context or sequence parallelism partitions tokens along the sequence dimension.

For attention, each rank initially owns only part of the queries, keys, and values. Exact attention requires every query block to interact with all causally valid key/value blocks. Ring Attention circulates key/value blocks among ranks while computing blockwise attention, reducing per-rank memory without approximating dense attention; see Ring Attention with Blockwise Transformers.

The exchange is again explicit: longer sequences become possible, but key/value communication and load balancing enter the critical path. Causal masking creates unequal useful work for naïve partitions, so block assignment matters.

Do not confuse:

  • tensor parallelism: shard hidden dimensions or operators;
  • pipeline parallelism: shard layers;
  • context parallelism: shard sequence positions;
  • data parallelism/FSDP: shard examples and possibly model state across replicas.

These dimensions can be composed. If

[ N_{GPU}=D\times T\times P\times C, ]

then (D,T,P,C) are the data, tensor, pipeline, and context degrees. The product is easy; choosing topology-aware groups that actually perform well is the engineering problem.

Experiment 8: compare parallelism by the failure it solves

Use a model that fits under FSDP but increase one axis at a time:

  • widen hidden layers until a live layer becomes the pressure point;
  • deepen the model until stage partitioning becomes attractive;
  • lengthen context until activations dominate.

For each case, test only the corresponding parallelism. Measure memory relief, collective volume, communication overlap, idle time, and implementation complexity. This prevents treating all forms of “model parallelism” as interchangeable.

9. Communication collectives, NCCL, and network bandwidth

Distributed training is a graph of local tensor operations connected by communication collectives. NCCL implements GPU-focused collectives and point-to-point operations over transports such as NVLink and high-speed networks.

9.1 The essential collectives

Collective Result Common training use
Broadcast one rank's tensor copied to all initialization or metadata distribution
All-reduce values reduced and result copied to all DDP gradient synchronization
Reduce-scatter values reduced, each rank retains one shard sharded gradients
All-gather rank-local shards concatenated on all ranks reconstruct sharded parameters
All-to-all every rank sends a distinct shard to every rank token routing or some sequence/expert layouts
Send/receive point-to-point transfer pipeline stage boundaries or rings

An all-reduce can be understood as reduce-scatter followed by all-gather. This identity explains the relationship between replicated gradients and sharded gradient ownership.

9.2 Latency, bandwidth, and arithmetic intensity

A rough communication time model is

[ T_{comm}\approx \alpha\cdot n_{messages}+\frac{bytes}{effective\ bandwidth}, ]

where (\alpha) represents latency and software/protocol overhead. Large collectives are mainly bandwidth-bound; many tiny collectives are latency-bound.

Peak link bandwidth is not application bandwidth. Topology, contention, protocol, message size, rank placement, PCIe paths, NIC affinity, and collective algorithm all matter. Measure collectives using the exact node layout used by training.

9.3 Overlap and stragglers

Communication that occurs concurrently with useful compute may be mostly hidden. Exposed communication extends the step. Bucketing, prefetching, and schedule design aim to create overlap, but excessive concurrency can cause resource contention.

Synchronous training advances at the speed of the slowest rank. A single worker with slow input, thermal throttling, network retransmission, or a noisy neighbor can stall every collective. Monitor per-rank step components and collective latency distributions, not just cluster averages.

Experiment 9: benchmark the fabric before the model

Run size sweeps for all-reduce, all-gather, reduce-scatter, and point-to-point traffic:

  • inside one node;
  • across two nodes;
  • across the intended full topology.

Then compare microbenchmark bandwidth with application traces. A healthy fabric test plus slow training usually points toward bad overlap, small collectives, load imbalance, or data stalls rather than raw network capacity.

10. Long runs make failure a normal state

Let each worker fail independently at rate (1/M), where (M) is its mean time between failures. For (N) workers, the approximate job-level failure probability during duration (T) is

[ P(\text{at least one failure})\approx 1-e^{-NT/M}. ]

Scaling workers and duration therefore makes “restart from the beginning” economically unacceptable.

10.1 What a training checkpoint must contain

A weight file alone resumes a model, not a training run. A correct checkpoint normally includes:

  • model parameters and buffers;
  • optimizer states;
  • learning-rate scheduler state;
  • AMP scaler state for FP16;
  • optimizer-update number and tokens processed;
  • CPU and accelerator RNG states;
  • data sampler, shard order, streaming cursor, and shuffle-buffer state or a deterministic reconstruction rule;
  • gradient-accumulation position if mid-update checkpoints are allowed;
  • tokenizer and dataset-manifest hashes;
  • model, optimizer, and precision configuration;
  • code revision, container/runtime versions, world topology, and checkpoint schema version.

Checkpoint only at optimizer-update boundaries unless there is a compelling reason not to. Mid-accumulation recovery requires saving partial gradients and exact microbatch position.

10.2 Distributed checkpoints

A large sharded run should not funnel all state through rank 0. Each rank writes its local shards in parallel, while a manifest maps logical tensors to physical files. PyTorch Distributed Checkpoint supports parallel save/load and load-time resharding, which can allow a checkpoint to load under a different compatible topology.

A robust checkpoint protocol is transactional:

  1. write shards under a temporary checkpoint ID;
  2. calculate sizes and checksums;
  3. make shard durability explicit;
  4. publish the manifest or completion marker last;
  5. treat only committed checkpoints as recoverable;
  6. retain multiple known-good checkpoints;
  7. regularly restore one in an isolated test job.

Checkpoint time matters because a synchronous save pauses expensive GPUs. Async staging can reduce the pause but needs enough host memory and I/O bandwidth and must not report success before durable storage commits.

10.3 Choosing checkpoint frequency

If a checkpoint costs (C) time units, checkpoints occur every (I), and job failures occur with mean interval (M), a simple expected-overhead model is

[ \frac{C}{I}+\frac{I}{2M}. ]

The first term is checkpoint overhead; the second is expected lost work. Minimizing the approximation gives (I\approx\sqrt{2CM}). Real systems add restart time, correlated failures, storage contention, and checkpoint validation, but the derivation shows why the interval must depend on both save cost and observed failure rate.

10.4 Fault recovery correctness

After resume, verify:

  • the next sample IDs match the uninterrupted run or the documented replay policy;
  • token and optimizer clocks continue, not restart;
  • the next learning rate is correct;
  • optimizer moments are present and correctly sharded;
  • loss does not jump beyond expected numerical variation;
  • no rank loads stale or partial state.

At-least-once sample replay after a failure may be acceptable if bounded and measured. Silent skips or unbounded repeats are not.

Experiment 10: kill the job intentionally

Checkpoint a short run, terminate a worker at several points, and resume. Compare its next few batches, loss, learning rate, and parameters against an uninterrupted control. Measure:

  • checkpoint pause and total checkpoint duration;
  • detection time;
  • scheduler/relaunch time;
  • load and reshard time;
  • replayed work;
  • total recovery time.

Fault tolerance that has never been tested is only an assumption.

11. Reproducibility is a spectrum, not a seed

PyTorch's reproducibility guidance explicitly notes that complete reproducibility is not guaranteed across releases, commits, platforms, or even CPU and GPU execution with identical seeds.

Useful reproducibility levels are:

  1. run reconstruction: same code, data manifest, tokenizer, configuration, and environment can be assembled;
  2. sample-order reproducibility: ranks consume the same examples in the same order;
  3. numerical reproducibility: metrics stay within an agreed tolerance;
  4. bitwise reproducibility: every value is identical—expensive and often not portable.

Record all random seeds and RNG states, but also deterministic-algorithm settings, library versions, compiler flags, GPU type, world size, rank mapping, data-worker seeds, and collective configuration. Non-associative floating-point addition means a different reduction tree can change low-order bits and eventually training trajectory.

The production goal is usually explainable and bounded variance, not universal bitwise identity. Keep repeat small runs to estimate ordinary seed variance before attributing every metric change to a code change.

12. Training observability: know whether compute becomes learning

A run can remain alive while wasting millions of tokens. Monitoring must cover optimization, performance, data, distributed health, checkpoints, and cost.

12.1 Loss curves

Track token-weighted training loss, held-out validation loss, and optionally perplexity. Interpret shapes causally:

  • smooth decline: optimization is learning the measured distribution;
  • sudden spike on all ranks: bad batch, LR issue, numerical overflow, or data-distribution change;
  • one-rank anomaly before a collective stall: data or hardware problem on that rank;
  • train loss falls while validation rises: overfitting, train/validation mismatch, or contamination/duplication effects;
  • flat loss: too-small LR, frozen parameters, broken labels, insufficient capacity, or an optimizer bug;
  • discontinuity after resume: incomplete state, wrong data cursor, or scheduler reset.

Validation must run on a stable, document-disjoint set with a fixed tokenizer and masking policy. Otherwise changes in evaluation mechanics masquerade as learning.

12.2 Gradient and parameter statistics

Track at least:

  • global gradient norm before clipping;
  • clipped-step fraction;
  • non-finite gradients and parameters;
  • update norm and parameter norm;
  • update-to-weight ratio by layer;
  • activation or logit outliers for selected layers;
  • FP16 scaler and skipped updates when applicable.

Logging full histograms every step is expensive. Sample layers and use a slower cadence, with a higher-resolution capture triggered by anomalies.

12.3 Systems metrics

The core counters are:

  • valid training and validation tokens processed;
  • sequences, microbatches, and optimizer updates;
  • tokens per second per GPU and for the whole job;
  • GPU compute utilization and achieved FLOP rate;
  • allocated/reserved/peak memory;
  • input wait, forward, backward, optimizer, and exposed communication time;
  • collective latency by type and message size;
  • checkpoint pause, write bandwidth, duration, and failures;
  • restart count, lost tokens, and recovery time;
  • energy or cloud cost and cost per billion tokens.

Use torch.profiler for bounded trace windows, not continuously at maximum detail. Aggregate dashboards find when a run is unhealthy; traces explain why.

12.4 Throughput and utilization

Always use valid tokens rather than padded tokens when reporting useful throughput:

[ \text{tokens/s}=\frac{\sum \text{non-masked target tokens}}{\text{wall time}}. ]

Model FLOP utilization (MFU) compares estimated achieved model FLOPs to hardware peak. A common dense-transformer approximation starts near six parameter-FLOPs per training token—forward plus backward—but attention, embeddings, recomputation, sparsity, and fused kernels make the true count architecture-dependent. Publish the exact FLOP accounting used; otherwise MFU comparisons are ambiguous.

13. Scaling laws and compute budgets

Training is an allocation problem: for a fixed budget, how much should go into model parameters, training tokens, context, data quality experiments, and failed-run insurance?

Empirical scaling laws show that language-model loss often follows approximate power-law relationships with model size, data, and compute over measured regimes. The Chinchilla study, based on hundreds of training runs, showed that many earlier large models were undertrained and that compute-optimal parameter count and training tokens should grow together under its studied conditions.

These are empirical planning tools, not physical laws. Architecture, tokenizer, data distribution and quality, optimization, target tasks, inference cost, and reuse of a released model can move the economic optimum.

13.1 Compute budget accounting

Before a large run, budget:

  • pilot and ablation runs;
  • the main run's planned token count;
  • validation and checkpoint overhead;
  • expected lost work from failures;
  • reruns due to software or data errors;
  • evaluation and safety analysis;
  • storage, networking, and artifact retention;
  • post-training and inference costs if choosing model size.

Measure actual accelerator-hours and cost rather than multiplying list-price GPUs by scheduled duration. Queue delay, idle allocations, low utilization, retries, and storage/network charges all alter effective cost.

13.2 Use small runs to choose a large run

Train a grid of smaller models and token budgets. Fit loss-versus-compute trends, but also evaluate target capabilities. Use those experiments to select model/data allocation and to estimate variance. Extrapolate only within a defensible range and carry uncertainty into the budget.

A scaling curve cannot tell you whether a source is legally usable, whether a safety filter erased a language, whether a benchmark is contaminated, or whether the cluster can sustain the predicted throughput. It complements, rather than replaces, data and systems validation.

14. Training evaluation and pretraining contamination

Training loss measures next-token prediction on the training mixture. It does not alone establish useful capability, generalization, safety, or absence of memorization.

14.1 Evaluation layers

Use several levels:

  1. held-out language-model loss by source, language, and domain;
  2. capability evaluations aligned to intended uses;
  3. behavior and safety evaluations;
  4. memorization and privacy probes;
  5. efficiency evaluation: model size, inference memory, latency, and throughput.

Evaluate at planned token checkpoints to understand emergence, regression, and overtraining, but do not repeatedly tune against a supposedly untouched final test set.

14.2 Pretraining contamination

If benchmark questions, answers, paraphrases, source documents, or generated derivatives occur in pretraining data, benchmark scores can overestimate generalization. Exact-string decontamination detects only part of the problem; near-duplicates, translations, templated variants, and semantic equivalents may remain.

A defensible workflow:

  • freeze an evaluation registry before final corpus construction;
  • search for normalized exact matches, long n-gram overlaps, and near-duplicates;
  • remove matched documents or define conservative exclusion windows;
  • record which benchmark items were affected;
  • report results with contaminated items excluded where possible;
  • use post-training or private time-split evaluations when available;
  • retain the data lineage needed for later audits.

Decontamination itself has false positives and false negatives. Report the method and thresholds rather than stating only “decontaminated.”

15. Model release and documentation

A trained weight tensor is not a complete release. A responsible model package should include:

  • architecture and parameter count;
  • tokenizer files, normalization, vocabulary, and special-token semantics;
  • context length and attention behavior;
  • training objective, token count, batch schedule, precision, optimizer, LR schedule, clipping, and checkpoint policy;
  • data-source categories, time range, filtering, deduplication, decontamination, languages, licenses, and known omissions;
  • evaluation methods, results, uncertainty, and contamination caveats;
  • intended uses, prohibited or unsupported uses, risks, and limitations;
  • inference requirements and numerical formats;
  • model license and attribution requirements;
  • code revision and dependency/container information;
  • security and integrity hashes for released artifacts;
  • enough run metrics to audit compute and training stability.

Document unsuccessful or ambiguous results too: instability interventions, changed data mixtures, resumed checkpoints, known data gaps, and evaluation regressions. This is part of scientific provenance and operational trust, not merely release marketing.

16. The complete practical project

The project below follows the exact progression from a local character model to a recoverable multi-GPU transformer.

Stage 1 — Train a character-level language model

Implement character vocabulary construction, shifted examples, a tiny causal model, cross-entropy, backward, optimizer, and validation. No trainer. Confirm overfitting on one batch, then generalization on a document-disjoint validation split.

Deliverable: train_char.py, configuration, metrics JSONL, and generated sample.

Stage 2 — Train a small tokenizer

Train BPE or Unigram SentencePiece on a representative subset. Measure fertility across ordinary English, code, numbers, whitespace, Telugu or other target languages, and domain text. Freeze special tokens and hash the model.

Deliverable: tokenizer model, vocabulary, training-corpus manifest, and tokenizer evaluation report.

Stage 3 — Train a small transformer

Implement embeddings, positional mechanism, masked multi-head self-attention, MLP, residual paths, normalization, and output projection. Weight tying is optional but must be recorded. Unit-test causal masking and input/target alignment.

Deliverable: readable model code and an explicit parameter-count breakdown.

Stage 4 — Add mixed precision

Add autocast and FP16 scaling where needed. Keep a FP32 baseline and compare stability, throughput, and memory.

Gate: no unexplained non-finite values or systematic validation regression.

Stage 5 — Add gradient accumulation

Separate microstep and optimizer-step counters. Normalize by valid tokens and advance the scheduler once per update.

Gate: one controlled update matches a physical large-batch reference within tolerance.

Stage 6 — Measure memory usage

Log parameter bytes by dtype, gradient bytes, optimizer-state bytes, peak allocated/reserved memory, and a sampled memory snapshot. Sweep context and microbatch size with and without activation checkpointing.

Gate: predict a safe microbatch for a new context length from the measured ledger.

Stage 7 — Add multi-GPU data parallelism

Use one process per GPU, a distributed sampler, rank-aware seeds, DDP, and no_sync() during intermediate accumulation microsteps. Aggregate global token-weighted metrics.

Gate: no duplicate rank-local samples in the audit window; controlled update matches the single-GPU reference; scaling efficiency is reported.

Stage 8 — Add checkpoint and resume

Save all training and data-order state at optimizer boundaries. Write transactionally, retain at least two generations, and add a kill-and-resume integration test.

Gate: resumed and uninterrupted controls consume the intended next batches and produce matching or bounded-equivalent updates.

Stage 9 — Evaluate the model

Report validation loss by source and sequence length, perplexity under the fixed tokenizer, small capability tests appropriate to model size, memorization checks, and generation samples using fixed prompts.

Gate: results include uncertainty and negative results; no claim exceeds the measurement.

Stage 10 — Document the complete run

Create a run card containing data manifest, tokenizer hash, code revision, environment, model config, optimizer/schedule, global batch in tokens, precision, hardware, topology, checkpoints, interruptions, final metrics, cost, and known limitations.

Gate: another engineer can reconstruct the pipeline and explain every number on the dashboard.

Minimal run record

run_id: tiny-gpt-001
code_revision: <git-sha>
dataset_manifest_sha256: <hash>
tokenizer_sha256: <hash>
model:
  layers: 8
  hidden_size: 512
  heads: 8
  context_length: 1024
optimization:
  optimizer: adamw
  precision: bf16
  microbatch_sequences_per_gpu: 4
  gradient_accumulation_steps: 8
  data_parallel_world_size: 4
  max_grad_norm: 1.0
clocks:
  optimizer_updates: 0
  valid_tokens_processed: 0
artifacts:
  checkpoint_schema_version: 1

The YAML is not the source of truth for measured counters; it identifies configuration and artifact versions. Runtime metrics remain append-only time-series data.

17. A larger distributed architecture—designed, not necessarily executed

Consider a dense model large enough that replicated Adam state does not fit on a GPU, a multi-trillion-token corpus, and 64 GPUs arranged as eight nodes of eight GPUs each.

17.1 Data plane

  • Immutable, versioned raw and processed datasets live in durable object storage.
  • Pre-tokenized, checksummed shards are large enough for efficient sequential I/O but small enough for balancing and retry.
  • Each node stages upcoming shards on local NVMe.
  • A deterministic global planner assigns shards and seeds; workers maintain resumable positions.
  • The loader prefetches into pinned host memory and performs asynchronous host-to-device copies.
  • The pipeline emits source, language, rejection, duplication, and wait-time metrics.

17.2 Compute topology

Start with the least complicated strategy that fits:

  1. BF16 mixed precision.
  2. Memory-efficient attention kernels.
  3. Activation checkpointing at transformer-block granularity.
  4. FSDP/full sharding to remove replicated model state.
  5. Tensor parallelism only if live layers or per-rank computation still require it.
  6. Pipeline parallelism only if depth partitioning is required or beneficial after measuring bubbles.
  7. Context parallelism only for sequence lengths that remain infeasible after activation/kernel optimizations.

Keep frequent tensor-parallel collectives inside the eight-GPU high-bandwidth node when possible. Use slower inter-node links for coarser data/FSDP groups or pipeline boundaries. One possible layout is tensor parallel degree 8 inside each node and data-parallel degree 8 across nodes. Another is FSDP across eight GPUs per node with eight replicated node groups. The correct choice depends on whether parameter memory, live-layer memory, or inter-node bandwidth is the dominant limit.

17.3 Control plane

  • A scheduler launches one process per GPU with explicit rank and topology metadata.
  • Preflight tests validate CUDA/NCCL versions, peer connectivity, collective bandwidth, storage access, clock synchronization, and sufficient checkpoint space.
  • A coordinator tracks heartbeats and progress but does not sit in the tensor data path.
  • Failure policy terminates the inconsistent gang, replaces bad capacity, and resumes the complete job from the last committed checkpoint.
  • Canary runs validate new images, kernels, data manifests, and checkpoint compatibility before allocating the full cluster.

17.4 Checkpoint plane

  • Every rank writes local logical shards in parallel.
  • A completion manifest commits only after checksum verification.
  • Async staging is bounded so it cannot exhaust host memory.
  • Checkpoint cadence is chosen from observed failure and write costs.
  • Daily restore drills load the newest checkpoint under an isolated job; periodic tests reshard to an alternate compatible topology.
  • Retention includes recent rolling checkpoints, milestone checkpoints, and the final release state.

17.5 Observability and cost

Dashboards connect four clocks: wall time, optimizer updates, valid tokens, and consumed accelerator-hours. Alerts cover:

  • non-finite loss or gradients;
  • unexpected loss slope or validation regression;
  • repeated clipping or FP16 skipped updates;
  • rank-level step-time skew;
  • exposed collective time and bandwidth degradation;
  • GPU underutilization and input starvation;
  • memory approaching the tested safety boundary;
  • checkpoint duration, age, or verification failure;
  • tokens or cost deviating from the plan.

The primary efficiency objective is not GPU utilization alone. It is validated learning progress per unit of time and cost. A system can show high utilization while recomputing excessively, training on duplicates, or using the wrong sample order.

18. A decision table for scaling techniques

Observable failure First technique to test What it trades Do not expect it to solve
Desired batch does not fit Gradient accumulation more microsteps per update parameter-state memory
Low-precision hardware is underused BF16/FP16 mixed precision numerical margin and dtype complexity all optimizer-state memory
Activations dominate memory Activation checkpointing recomputation for memory parameters/Adam state
Input pipeline starves GPU Streaming, sharding, prefetch, local cache complexity and buffering slow model kernels
Model fits, run is too slow DDP/data parallelism gradient communication, larger global batch per-GPU model memory
Replicated Adam/gradients/parameters do not fit FSDP/ZeRO-style sharding gathers/scatters and transient buffers a single live operator that cannot fit
A live layer is too large Tensor parallelism frequent low-latency collectives pipeline depth imbalance
Model depth cannot fit or needs coarser partition Pipeline parallelism bubbles, scheduling, boundary traffic long-context attention alone
One long sequence cannot fit Context parallelism KV/attention communication parameter-state pressure
Failure risk makes restart unacceptable Distributed transactional checkpoints I/O, pause, storage corrupt or unversioned data

The ordering is deliberate. Each added dimension increases the number of ways a run can be slow, numerically different, or unrecoverable.

19. Mastery gate: reconstruct the system without memorizing labels

You have mastered the material when you can produce and defend these artifacts:

  1. A handwritten PyTorch loop whose input/target shift, loss denominator, backward pass, optimizer step, and scheduler clock are correct.
  2. A byte ledger separating parameters, gradients, optimizer states, activations, temporary buffers, and allocator reserve.
  3. A gradient-accumulation equivalence test, including the variable-valid-token case.
  4. A one-page comparison showing which tensor dimension or state each parallelism strategy shards and which collective it introduces.
  5. A topology-aware choice of data, full-shard, tensor, pipeline, and context degrees for a specified model and cluster.
  6. A transactional distributed-checkpoint design that captures optimizer, RNG, data cursor, clocks, and artifact identities.
  7. A kill-and-resume experiment with measured checkpoint and recovery time.
  8. A dashboard explanation that connects a loss anomaly to gradients, data, communication, and per-rank systems evidence.
  9. A compute budget that includes pilots, failures, checkpointing, evaluation, and idle inefficiency.
  10. A complete lineage from raw source and tokenizer through token batches, optimizer updates, checkpoints, evaluation, and release documentation.

If you can explain not only what technique is used but which previous system failed, which bytes or operations moved, which collective appears, how correctness is tested, and what new failure mode was introduced, then you understand language-model training infrastructure from first principles.

Primary references and current implementation documentation

PyTorch and NCCL documentation

Primary research