Model Training Infrastructure and Distributed Training Questions

Direct answer

GPU memory capacity limits how large a batch and model you can fit at all, while memory bandwidth (how fast data moves between GPU memory and compute units) limits how fast you can actually process that data once it fits, so a training configuration can be capacity-bound, bandwidth-bound, or compute-bound depending on which resource is the tightest constraint for a given model and batch size.

Structured elaboration

• Capacity constraint: parameters, gradients, optimizer state, and activations must all fit within the GPU's total memory (e.g. 80GB on an H100); exceeding this causes an out-of-memory error regardless of how fast the GPU could otherwise compute, forcing a smaller batch size, model, or the memory-saving techniques (mixed precision, checkpointing, sharding) discussed elsewhere in this topic.
• Bandwidth constraint: even when everything fits, operations that move a lot of data relative to the compute they perform (elementwise operations, some normalization layers, small matrix multiplies) can be bottlenecked by how fast data streams between HBM and the GPU's compute units rather than by the compute units' raw throughput; this shows up as low GPU utilization percentage even though the workload is running continuously.
• Batch size trade-off: larger batch sizes generally improve compute efficiency (better utilization of the GPU's parallel compute, amortizing fixed per-kernel-launch overhead) up to the point where capacity is exhausted; very small batch sizes tend to be more bandwidth-bound (relatively more data movement per unit of useful compute) since fixed overheads and memory-bound elementwise operations don't scale down favorably.
• Model size trade-off: a larger model directly increases capacity pressure (more parameters, more activation memory per sample) which forces either a smaller batch size or additional memory-saving techniques, in turn affecting how compute-efficient a given step actually is.

Worked example

Training a model where activation memory dominates: at batch size 8, activations plus weights plus optimizer state fit within an 80GB GPU with headroom, and the GPU runs near its compute-bound peak throughput; increasing batch size to 32 might exceed the 80GB budget entirely (capacity-bound failure), while decreasing to batch size 1 for debugging purposes leaves so much of each kernel launch as fixed overhead relative to the tiny amount of actual compute that the GPU spends much of its time bandwidth- or launch-overhead-bound rather than compute-bound.

Trade-offs & pitfalls

A common mistake is assuming more GPU memory alone fixes a slow training run; if the workload is bandwidth-bound rather than capacity-bound, adding memory (a bigger GPU) doesn't address the actual bottleneck, only profiling (checking whether the GPU is near its compute-throughput ceiling or spending time waiting on memory transfers) reveals which constraint is actually binding.

Direct answer

An effective checkpoint for iterative training persists everything needed to resume training exactly where it left off: model weights, optimizer state, the RNG state, and the current training position (epoch/step and, ideally, data-loader offset), saved at a cadence balancing recovery-time risk against checkpoint overhead, and written to durable object storage (S3/GCS) through an atomic write pattern so a crash mid-upload never leaves a corrupted checkpoint that a resume could silently load.

Structured elaboration

• Model weights (state_dict): the parameters themselves, obviously necessary to resume with the same trained values rather than starting from scratch.
• Optimizer state: for Adam-family optimizers, the first and second moment estimates per parameter; omitting this and just restoring weights effectively resets momentum, causing a noticeable disruption to the loss curve right after resume as the optimizer "relearns" its momentum estimates from scratch.
• RNG (random number generator) state: seeds and generator states for Python's random, NumPy, and the deep learning framework's own RNG (CPU and, separately, each GPU's RNG), needed for bit-for-bit reproducibility of anything stochastic that happens after resume (dropout masks, data shuffling order, data augmentation).
• Training position: current epoch, step count, and ideally enough information to resume the data loader at the correct position rather than restarting the epoch from the beginning (which would either skip or duplicate some samples relative to a run that never failed).
• Scheduler state: if using a learning-rate scheduler, its internal state (e.g. current step count for a warmup or decay schedule) so the LR trajectory continues correctly rather than restarting the schedule.
• Atomic uploads to object storage (S3/GCS): object stores do not offer atomic "in-place" file updates, so a checkpoint is first written to a temporary key (or local disk) and only made visible to resume logic via a single atomic operation, either an object-store copy/rename to the final key, or a separate small "pointer" object (e.g. latest.json naming the last fully-uploaded checkpoint key) that is overwritten only after the full checkpoint upload succeeds. This guarantees a resume never reads a partially-uploaded, truncated checkpoint: a crash mid-upload leaves the old pointer intact and the incomplete object simply unreferenced.
• Fast resume: resume time is bounded by parallelizing the download of the checkpoint's shards/files (rather than one large sequential download) and by keeping the checkpoint format such that a resuming process can start pulling from object storage immediately once the pointer resolves, without waiting on any additional coordination step.

Worked example

A training run failing at step 45,231 with checkpoints saved every 1,000 steps has its last checkpoint at step 45,000; resuming from it and correctly restoring optimizer state, RNG state, and the scheduler's step counter means the run continues as if uninterrupted from step 45,000, needing to recompute only the 231 steps' worth of work between the checkpoint and the failure. If the step-45,000 checkpoint upload to S3 was itself interrupted mid-write (e.g. a node failure during the upload), the atomic pointer-swap design means the resume logic still sees the step-44,000 checkpoint as "latest" (the pointer was never advanced), correctly falling back to it instead of loading a truncated, corrupted step-45,000 object.

Trade-offs & pitfalls

Checkpointing more frequently bounds the amount of lost work on failure but adds I/O overhead (writing potentially many gigabytes of state) and can itself become a meaningful fraction of total training time if the interval is too short relative to the checkpoint's write time; the right cadence is set by balancing the expected failure rate against the checkpoint write cost. Skipping the atomic-pointer step and instead overwriting a single well-known key directly (e.g. always uploading to checkpoint-latest.pt) is a common shortcut that reintroduces exactly the corruption risk atomicity is meant to prevent: a crash mid-upload leaves that key in a truncated, unloadable state with no older fallback.

Direct answer

Gradient accumulation simulates training with a larger batch size than fits in memory by running several forward/backward passes on smaller micro-batches, summing (or averaging) their gradients without updating the weights, and only applying the optimizer update once after the desired number of micro-batches has accumulated.

Structured elaboration

• Why it's needed: GPU memory limits the largest batch that fits for a given model; gradient accumulation decouples the "effective batch size" the optimizer sees from the "micro-batch size" that actually fits in memory, letting you train with a large effective batch on hardware that can't hold that batch's activations all at once.
• Algorithm: for each of K micro-batches, run forward and backward (accumulating .grad on each parameter, since PyTorch's default is to add to existing gradients rather than overwrite), without calling optimizer.step() or zero_grad() in between; after the Kth micro-batch, call optimizer.step() once, then zero_grad() to reset for the next cycle.
• Loss scaling convention: divide each micro-batch's loss by K before calling .backward(), so the summed gradient across K micro-batches equals the average gradient a true large-batch step would have produced, keeping the effective learning rate consistent with what was tuned for that effective batch size.

Worked example

To simulate an effective batch size of 128 on hardware that only fits 32 samples at a time: run 4 micro-batches of 32 samples each, dividing each micro-batch's loss by 4 before backward, accumulating gradients across all 4, then calling optimizer.step() once; the resulting parameter update is (approximately, modulo any batch-size-sensitive layers like BatchNorm) equivalent to what a single true batch-128 step would have produced.

Interaction with learning-rate schedules

A learning-rate scheduler's .step() must be called once per completed accumulation cycle (once per true optimizer step), not once per micro-batch; calling it every micro-batch effectively runs the schedule K times faster than intended relative to the actual number of samples/effective-batches seen, so a warmup or decay schedule tuned for, say, 1000 optimizer steps would finish in 1000 micro-batches instead of 1000*K micro-batches, a large, easy-to-miss bug. The scheduler's total step count (and any warmup-step count) should be computed as total_micro_batches / accumulation_steps, matching the optimizer_step counter, not the raw micro-batch or sample index.

Trade-offs & pitfalls

Gradient accumulation increases wall-clock time per effective step (K sequential forward/backward passes instead of one) since it doesn't add compute parallelism, only memory savings; it also doesn't fix batch-size-sensitive layers like BatchNorm, whose statistics are still computed per micro-batch rather than over the full effective batch, unless specifically handled (e.g. via sync-BN-style aggregation or accepting the approximation).

Direct answer

Synchronous SGD has every worker compute a gradient against the same, current parameter values and waits for all workers before applying a single combined update, giving convergence behavior equivalent to (or very close to) single-machine SGD at a larger effective batch size; asynchronous SGD lets each worker push its gradient and pull fresh parameters independently, without waiting for others, trading some workers computing gradients against slightly outdated ("stale") parameters for higher hardware utilization.

Structured elaboration

• Synchronous: every worker's gradient this step is computed against identical parameter values (the state after the previous step's update); once all gradients arrive, they're averaged and applied as one update, after which every worker again has identical, up-to-date parameters. Convergence guarantees closely mirror standard SGD's, since the process is mathematically equivalent to computing a gradient over a larger effective batch (the concatenation of every worker's mini-batch).
• Asynchronous: a worker pulls current parameters, computes a gradient, and pushes it back independently of other workers' progress; by the time its push arrives, the server's parameters may have already been updated by other workers' pushes in the meantime, meaning the pushed gradient was computed against parameters that are now "stale" (out of date) relative to the current server state.
• Staleness and its effect: the degree of staleness (how many other updates happened between a worker's pull and its push) tends to grow with more workers and with heterogeneous worker speeds (a slow worker's gradient becomes more stale the longer it takes to compute); staleness biases the effective update direction, since it's technically a gradient of an earlier point on the loss surface being applied to a later point, which can slow or, in extreme cases, destabilize convergence if unbounded.
• When each is chosen: synchronous is the default for most modern large-scale training (predictable convergence behavior, well-supported by AllReduce-based collectives) provided stragglers are managed; asynchronous is chosen specifically when worker heterogeneity or unreliability is severe enough that waiting for the slowest worker every step would be prohibitively wasteful, accepting some convergence-quality cost in exchange for higher aggregate hardware utilization.

Worked example

With 8 workers, one of which is consistently 3x slower than the others (a straggler): synchronous training's every-step wall-clock time is bounded by that slowest worker, wasting the other 7 workers' idle time waiting each step; asynchronous training lets the 7 faster workers keep contributing updates continuously without waiting, at the cost of the slow worker's occasional contributions being noticeably stale (computed against parameters several updates out of date) by the time they arrive.

Trade-offs & pitfalls

Bounded-staleness schemes (allowing async updates but capping how stale any single contribution is allowed to be before it's rejected or down-weighted) are a common middle ground, retaining most of asynchronous training's utilization benefit while limiting the worst-case convergence-bias risk that fully unbounded asynchrony carries.

Direct answer

Mixed-precision training runs the forward and backward passes in a lower-precision format (fp16 or bfloat16) for speed and memory savings, while keeping a master copy of the weights in fp32 for the optimizer update, since accumulating many small fp16 gradient updates directly into fp16 weights loses precision that matters over the course of training.

Structured elaboration

• fp32: the standard single-precision format (8 exponent bits, 23 mantissa bits), wide dynamic range and good precision but twice the memory and compute cost of fp16/bf16.
• fp16: half precision (5 exponent bits, 10 mantissa bits); much narrower dynamic range than fp32, which is why small gradient values can underflow to zero and large activation/gradient values can overflow to infinity during training without countermeasures.
• bfloat16: also half precision but keeps fp32's 8 exponent bits (same dynamic range as fp32) while sacrificing more mantissa precision (7 bits); this trades some precision for avoiding fp16's overflow/underflow problems, which is why bf16 has become the more common default on hardware that supports it (avoids needing loss scaling in most cases).
• Throughput and memory benefits: half-precision tensors are half the memory footprint and modern GPU tensor cores execute half-precision matrix multiplies at roughly 2x (or more) the throughput of fp32, so both compute and memory bandwidth benefit.
• Common numerical pitfalls and mitigations: fp16's narrow dynamic range means small gradients can underflow to exactly zero, silently stalling learning for those parameters; loss scaling (multiplying the loss by a scale factor before backward, then dividing gradients by that same factor before the optimizer step) shifts small gradient values into fp16's representable range, preventing underflow. Dynamic loss scaling adjusts the scale factor automatically, increasing it when no overflow is observed for a while and decreasing it sharply if an overflow (inf/nan gradient) occurs.

Worked example

A gradient value of $1\times10^{-8}$ underflows to zero in fp16 (whose smallest positive normal value is around $6\times10^{-5}$). Scaling the loss by, say, 1024 before backward makes that same underlying gradient appear as $1.024\times10^{-5}$ after the scale is applied through the chain rule, still within fp16's representable range, and after backward the optimizer divides the accumulated gradient by 1024 again before applying the update, recovering the correct (unscaled) magnitude for the actual parameter update.

Selective casts and frameworks

Selective casting is the mechanism that makes 'mixed' precision actually mixed rather than all-or-nothing: rather than manually annotating every operation, frameworks provide an automatic-cast context (PyTorch's torch.autocast/torch.cuda.amp.autocast, or NVIDIA's older Apex amp) that runs numerically robust, compute-heavy ops (matrix multiplies, convolutions) in fp16/bf16 for the throughput win, while automatically keeping numerically sensitive ops (softmax, loss computation, some reductions and normalization statistics) in fp32, since those operations are more prone to precision-related instability at low bit-width. This selective-cast policy is what lets a model get fp16/bf16's throughput on the operations that benefit most without manually rewriting the model in mixed types. In terms of frameworks and hardware: PyTorch's torch.cuda.amp/torch.autocast and TensorFlow's tf.keras.mixed_precision API are the standard software layers, and both rely on NVIDIA Tensor Cores (available from the Volta generation onward, with native bf16 support from Ampere onward) to execute the low-precision matrix multiplies at higher throughput than fp32 CUDA cores, which is the hardware feature that actually makes the speedup real rather than just theoretical.

Trade-offs & pitfalls

bf16 avoids the underflow/overflow problem that motivates loss scaling in the first place, at the cost of coarser precision per value (fewer mantissa bits), which can matter for numerically sensitive operations (some normalization statistics, certain loss computations) that are often deliberately kept in fp32 even in an otherwise-bf16 training run.