Model Training Infrastructure and Distributed Training Questions

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.

Direct answer

Data parallelism replicates the entire model on every device and splits the training data across devices; model parallelism splits the model itself (different parameters on different devices) while typically keeping the same data on each; pipeline parallelism is a specific form of model parallelism that splits the model into sequential stages placed on different devices, with different devices processing different micro-batches at different pipeline stages simultaneously to keep all devices busy.

Structured elaboration

• Data parallelism: every device holds an identical, full copy of the model; each device processes a different slice of the mini-batch, computing its own local gradients; gradients are then synchronized (typically via all-reduce) across devices before the (identical, since it started from identical weights and applies an identical aggregated gradient) parameter update. Chosen when the model fits comfortably in a single device's memory but training would benefit from more parallel compute across more data.
• Model parallelism: the model's parameters themselves are split across devices (e.g. different layers, or different slices of very large layers via tensor parallelism), with each device holding only a portion of the full model; necessary specifically when the model is too large to fit on a single device, regardless of batch size or data volume. Its communication pattern is denser than data parallelism's: because a split weight matrix's output activations need to be reassembled (or its input activations need to be replicated) across the devices holding the split, tensor parallelism requires an all-gather or reduce-scatter/all-reduce at essentially every split-and-rejoin point in both the forward and backward pass, on every layer that's split, which is why tensor parallelism is generally kept within a single node over fast NVLink rather than spread across a slower network.
• Pipeline parallelism: a specific model-parallelism strategy that splits the model into sequential stages (contiguous groups of layers), each stage on a different device; to avoid devices sitting idle waiting for the previous stage to finish an entire batch, the batch is further divided into smaller micro-batches, which flow through the pipeline stages in an overlapped fashion (while stage 2 processes micro-batch 1's forward pass, stage 1 can already be processing micro-batch 2's forward pass), analogous to an assembly line. Its communication pattern is much sparser than tensor parallelism's: since each device holds a distinct, non-overlapping group of layers, communication between adjacent stages is limited to point-to-point send/recv of activations (forward direction) and gradients (backward direction) at each stage boundary, once per micro-batch, rather than the dense collective communication tensor parallelism requires on every split layer; this is why pipeline parallelism tolerates being spread across nodes with slower cross-node networking far better than tensor parallelism does.
• Combining all three (3D parallelism): for very large models on large clusters, these three axes are commonly combined: tensor/model parallelism within a node (exploiting fast NVLink), pipeline parallelism across a few node groups (tolerating the necessarily-less-frequent cross-node communication), and data parallelism to scale out replicas of that combined unit across many such groups, as discussed in the earlier large-model-training-design survivors.

Worked example

Training a model that fits comfortably on one GPU but would benefit from more compute: use pure data parallelism, replicating the model across, say, 8 GPUs, each processing 1/8 of the mini-batch. Training a model too large for even one GPU's memory: use model parallelism (splitting the model's parameters across devices) as a prerequisite before data parallelism can even be applied at all, since data parallelism assumes the whole model fits on each device to begin with. Training a very large model across many GPUs with a deep layer stack: pipeline parallelism specifically addresses the idle-time problem that would otherwise occur if you naively split layers across devices without micro-batching, since without micro-batches, each device would sit completely idle while waiting for the previous device to finish the ENTIRE batch's forward pass before it could start.

Trade-offs & pitfalls

A common conceptual confusion is treating "model parallelism" and "pipeline parallelism" as synonyms; pipeline parallelism is specifically ONE WAY to do model parallelism (splitting by sequential layer groups, with micro-batching to keep devices busy), while tensor parallelism (splitting individual weight matrices) is a different way to do model parallelism, and the choice between them (or combining both) depends on the specific memory/communication trade-offs of the model and cluster in question.

Direct answer

Gradient clipping caps the magnitude of gradients before the optimizer update to prevent occasional very large gradients (from an unstable loss landscape, a bad mini-batch, or numerical issues) from causing a destructively large parameter update; the two common strategies are global-norm clipping (rescale the entire gradient vector if its overall norm exceeds a threshold) and value clipping (clamp each individual gradient component to a fixed range independently).

Structured elaboration

• Global norm clipping: compute the L2 norm of the concatenated gradient across all parameters; if that norm exceeds a threshold max_norm, rescale every gradient component by the same factor max_norm / norm so the overall gradient direction is preserved but its magnitude is capped. This is the more common choice for deep networks (especially RNNs and transformers) because it preserves the gradient's relative direction across parameters, only shrinking its overall size.
• Value clipping: clamp each individual gradient component independently to a fixed range (e.g. [-c, c]), regardless of the overall norm. This is simpler and cheaper to compute (no need to first compute a global norm across all parameters), but can distort the gradient's direction, since components near the clip boundary get truncated disproportionately relative to smaller components, unlike norm clipping's uniform rescaling.
• Impact on training dynamics: clipping prevents a single unstable step from derailing training (a common cause of NaN/divergence in RNNs and transformers, especially early in training or with an aggressive learning rate), at the cost of somewhat biasing the effective update whenever clipping actually triggers (the update is smaller and, for value clipping, differently-directed than the "true" unclipped gradient would have produced).

Worked example

A gradient vector with an unusually large spike (norm 50, versus a typical norm around 2) and a clip threshold max_norm=1.0: global norm clipping rescales every component by 1.0/50 = 0.02, preserving the gradient's direction while capping magnitude to exactly 1.0; value clipping with a threshold c=0.1 instead clamps any individual component exceeding 0.1 to exactly 0.1 regardless of the others, which for a gradient with most components small but a few very large ones changes the update's direction, not just its magnitude, since the large components get disproportionately truncated relative to the small ones.

Trade-offs & pitfalls

Norm clipping's threshold needs to be chosen relative to the model's typical (unclipped) gradient norm during stable training, not an arbitrary small number; a threshold set far below typical gradient norms clips essentially every step, meaningfully slowing learning by systematically shrinking updates that weren't actually problematic.

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

GPUs are the default choice for most training workloads today because of their broad software ecosystem and flexibility across model architectures, while TPUs offer an edge specifically for very large, highly-regular workloads (especially transformers at scale) where their systolic-array matrix-multiply design and purpose-built high-bandwidth interconnect (ICI) pay off most.

Structured elaboration

• Workload characteristics favoring GPUs: irregular or research-stage architectures (custom ops, dynamic control flow, frequent architecture changes), smaller-scale training, and any workload needing the widest possible library/framework support (most cutting-edge research code targets GPU first).
• Workload characteristics favoring TPUs: very large, well-established architectures (transformers) trained at massive scale, where TPU pods' purpose-built interconnect and systolic-array matrix units are specifically tuned for the dense matrix multiplies dominating transformer compute, often at a favorable cost-per-FLOP for sustained large training runs.
• Programming model differences: GPU programming (CUDA, and the PyTorch/TensorFlow ecosystems built on it) is eager-execution-friendly and broadly flexible; TPUs traditionally require (or strongly favor) XLA-compiled, graph-based execution, which handles static, regular computation graphs excellently but is less forgiving of dynamic shapes or data-dependent control flow, and has historically meant a narrower (though growing) software ecosystem outside Google's own frameworks (JAX, TensorFlow).
• Availability and lock-in: GPUs are available across essentially every cloud provider and on-prem; TPUs are Google Cloud-specific, which is a real practical constraint for teams not already committed to that ecosystem.

Worked example

A research team iterating rapidly on novel architectures with custom CUDA kernels and dynamic control flow would find GPUs' flexibility essential; the same team, once they've settled on a stable, well-established transformer architecture and are scaling a single well-defined training run to thousands of accelerators, might specifically evaluate TPU pods for that production-scale run given TPUs' interconnect and systolic-array advantages at that specific workload shape, while continuing GPU-based experimentation for anything still in flux.

Precision support and instance-type selection

• Precision support: both platforms support bf16, but their history differs. GPUs (from the Volta generation onward for fp16, Ampere onward for native bf16 Tensor Core support) support fp16 well but fp16's narrow dynamic range typically requires loss scaling; TPUs were designed around bf16 from early generations specifically to avoid the loss-scaling machinery fp16 needs, so bf16-native training is the more idiomatic default on TPU, while GPU training commonly still supports both fp16 (with loss scaling) and bf16 depending on the GPU generation and framework defaults.
• Practical considerations for instance-type selection: cloud GPU instances (e.g. A100/H100-based VM families) are available in flexible single-GPU-to-multi-GPU configurations across essentially every major cloud provider, so a team can start small and scale incrementally; TPU pods are typically procured in fixed pod-slice sizes (e.g. v4/v5 slices of a given chip count) on Google Cloud specifically, which means committing to a minimum scale upfront and to a single cloud provider, a real practical constraint when a team wants provider flexibility or needs to start at very small scale before committing to a large training run.

Trade-offs & pitfalls

The GPU-versus-TPU choice is as much about ecosystem and organizational fit (existing tooling, cloud provider commitments, team CUDA expertise) as it is about a workload's raw performance characteristics; a team defaulting to GPUs purely out of ecosystem familiarity, even for a workload that would benefit from TPU's advantages, is a common and often reasonable trade-off, not a mistake, given the switching cost.