Model Training Infrastructure and Distributed Training Questions

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

Deciding between buying a larger GPU cluster to finish training faster versus training longer on cheaper hardware is a cost-per-unit-of-training-progress comparison: model each option's total cost (hardware/rental cost times duration) against the value of finishing sooner (if there's a genuine time-value, e.g. faster iteration enabling more experiments, or a hard deadline), and choose whichever option minimizes total cost for the actual constraint that matters (fixed budget, fixed deadline, or a genuine trade-off between the two).

Structured elaboration

• Cost model inputs: for the larger/faster cluster option, hourly cost per GPU (higher-end or more numerous GPUs typically cost more per hour) times the (shorter) training duration; for the smaller/cheaper option, lower hourly cost per GPU times the (longer) training duration; both need to be normalized to the SAME unit (total dollar cost for the SAME amount of total training progress, e.g. the same number of tokens/epochs). A fourth input, time-to-convergence curves (loss or eval-metric versus wall-clock time for each candidate configuration, not just raw throughput), matters because scaling to more GPUs often forces a larger global batch size, which can shift how many optimizer steps are needed to reach the same loss target; comparing options purely on throughput/GPU-hours without checking their actual convergence curves risks concluding the larger cluster is cheaper or faster when, after accounting for it needing more total steps to converge at the larger batch size, it may not be.
• When faster wins on pure cost: if the larger cluster achieves more than proportionally faster completion for its proportionally higher cost (e.g. due to per-GPU efficiency gains from better interconnect at larger scale, or economy-of-scale pricing), total dollar cost can actually be LOWER for the faster option, not just faster; this isn't guaranteed and needs to be checked with real numbers, not assumed.
• When cheaper-and-slower wins on pure cost: if scaling up doesn't achieve proportional speedup (the scaling-efficiency degradation discussed elsewhere in this topic), the larger cluster's extra cost isn't fully offset by proportional time savings, making the smaller, cheaper, longer-running option cheaper in total dollars, even though it finishes later.
• Beyond pure dollar cost, the value of time: if there's a genuine business reason time matters independent of dollar cost (a hard product deadline, or the ability to run more sequential experiments and therefore find a better model faster), that needs to be explicitly incorporated as a value (even if roughly estimated) added to the "faster" option's side of the comparison, not left as an unstated, implicit tiebreaker. This is the model's opportunity cost term: capital or GPU-hours committed to this training run cannot be spent on other experiments in the same window, so the relevant comparison is not just this run's dollar cost but that cost plus the estimated value of whatever else those GPU-hours or that calendar time could have produced instead, made explicit as a number (even a rough one) rather than an unstated intuition.

Worked example

Option A (larger cluster): 256 GPUs at $2/GPU-hour, completing training in 48 hours: total cost = $256 \times 2 \times 48 = \$24{,}576$. Option B (smaller cluster): 64 GPUs at $2/GPU-hour (same per-GPU rate, different scale, assuming the smaller cluster achieves close to linear scaling since 64 GPUs is a less demanding scaling regime than 256), completing the same total training progress in roughly 4x the wall-clock time due to 4x fewer GPUs (assuming near-linear scaling holds at this smaller scale): $64 \times 2 \times 192 = \$24{,}576$, the same total dollar cost in this case, since near-linear scaling at 64 GPUs means the two options are dollar-equivalent, and the actual decision then hinges entirely on whether the 48-hour-versus-192-hour completion-time difference has business value worth paying for (it doesn't cost extra in this particular numeric example, so faster genuinely wins here with no downside once scaling holds).

Trade-offs & pitfalls

The most common mistake in this kind of analysis is comparing only wall-clock time (bigger is obviously faster) or only sticker hourly cost (smaller is obviously cheaper) without actually computing total dollar cost for equivalent training progress at both options, which is the only apples-to-apples comparison; the real answer depends entirely on how scaling efficiency behaves at the specific GPU counts being compared, which must be measured (per the scaling-efficiency-measurement discussion elsewhere in this topic), not assumed to be either perfectly linear or badly degraded.

Direct answer

A GPU's memory hierarchy has several tiers with very different capacity and bandwidth trade-offs, from large-but-slower host (CPU) memory, through the GPU's own high-bandwidth DRAM (HBM), down to small-but-extremely-fast on-chip caches and shared memory close to the compute units; efficient training keeps data as close to the compute units as possible for anything accessed repeatedly, since moving data down this hierarchy costs bandwidth and latency at every level.

Structured elaboration

• Host (CPU) memory: the largest tier (often hundreds of GB to TB), but connected to the GPU only via PCIe (or NVLink, if available), meaningfully slower and higher-latency to access than the GPU's own memory; used for staging data before transfer, or for offloaded optimizer state/parameters in memory-constrained training. Data batches typically originate here (loaded and preprocessed on the CPU) before being transferred to the GPU for consumption, which is exactly the host-to-device transfer whose cost is analyzed below.
• HBM (GPU DRAM): the GPU's own dedicated high-bandwidth memory (tens to over a hundred GB on modern data-center GPUs), where model weights, gradients, optimizer state, and activations live during training; much faster than host memory but still an order of magnitude or more slower than the on-chip caches below it.
• L2 cache: a chip-wide cache shared across the GPU's compute units, smaller than HBM (tens of MB) but faster, automatically caching recently-accessed HBM data to reduce redundant HBM round-trips.
• L1 cache / shared memory: per-compute-unit (per streaming multiprocessor), very small (tens to low hundreds of KB) but extremely fast, used explicitly by well-optimized kernels to hold data being actively reused within a tight computational loop (e.g. matrix-multiplication tiling keeps a working tile in shared memory rather than re-reading from HBM for every operation on it).
• Where each quantity should reside: model parameters, gradients, and optimizer state persist in HBM for the full duration of training (they're read/written every step and are far too latency-sensitive to page in from host memory on the hot path); activations persist in HBM only for as long as needed (through the forward pass and until consumed by the corresponding backward pass, which is exactly what activation checkpointing and offloading techniques manage more aggressively when HBM is the binding constraint); data batches originate in host memory (where the CPU data-loading pipeline produces them) and are transferred to HBM just-in-time, one prefetched batch ahead of when the GPU needs it, rather than staging the entire dataset in either host or GPU memory at once.
• PCIe versus NVLink transfer cost and latency: for host-to-device transfer specifically (as opposed to GPU-to-GPU transfer), the relevant link is PCIe in the overwhelming majority of systems, since host memory connects to the GPU over the PCIe bus regardless of whether the GPUs also have NVLink to each other; PCIe generation and lane count set the achievable host-to-device bandwidth (for example, PCIe 4.0 x16 tops out around 32GB/s, PCIe 5.0 x16 around 64GB/s), meaningfully lower than intra-node NVLink's GPU-to-GPU bandwidth, and every host-to-device transfer also pays a fixed per-transfer latency overhead (kernel launch and DMA setup cost) that matters disproportionately for many small transfers versus fewer, larger ones. NVLink does not change host-to-device transfer at all on typical systems (host memory is not NVLink-connected); NVLink's advantage is specific to GPU-to-GPU traffic, which is why minimizing host-device movement (below) is a PCIe-bandwidth problem specifically, distinct from the GPU-to-GPU NVLink-versus-PCIe question covered elsewhere in this topic.
• Practical rules of thumb for minimizing host-device data movement: (1) use pinned (page-locked) host memory for any tensor being transferred, since pinned memory supports faster, asynchronous DMA transfer than regular pageable memory; (2) issue the transfer asynchronously and overlap it with GPU compute on the previous batch (the same prefetching principle as the data-loading pipeline, applied specifically to the final host-to-device copy step) rather than blocking the GPU on each transfer; (3) batch small transfers into fewer, larger ones wherever possible, since each individual transfer pays fixed overhead independent of its size; (4) avoid unnecessary round-trips, such as pulling a GPU tensor back to host memory for logging/debugging on every step, which silently reintroduces PCIe-bound stalls into an otherwise GPU-resident training loop; and (5) keep anything reused across steps (model parameters, optimizer state) resident on the GPU permanently rather than re-transferring it, reserving host-device transfer specifically for the genuinely new data (the next batch) each step actually needs.

Impact on training: operations that are compute-bound (dense matrix multiplies, well-optimized to reuse data from fast on-chip memory) achieve close to the GPU's peak theoretical throughput; operations that are memory-bound (elementwise operations processing each value once, with little reuse) are limited by HBM bandwidth, not compute throughput, regardless of how fast the compute units themselves are, which is a large part of why operator fusion (combining several memory-bound operations into one kernel to reduce redundant HBM round-trips) is a meaningful optimization technique.

Worked example

A well-optimized matrix multiplication kernel tiles the computation so each tile of the input matrices is loaded into fast shared memory once and reused for many multiply-accumulate operations before being evicted, achieving high compute-unit utilization; a naive elementwise operation (like adding a bias vector) that reads and writes HBM directly for every element with no reuse is instead limited by HBM's bandwidth ceiling, explaining why such operations show comparatively low measured throughput relative to the GPU's advertised peak FLOPS, despite being computationally simple.

Trade-offs & pitfalls

A common misconception is assuming a GPU's advertised peak FLOPS figure represents achievable throughput for any workload; the memory hierarchy means actual achieved throughput depends heavily on how much data reuse a given operation allows, with memory-bound operations achieving only a fraction of peak FLOPS regardless of how fast the compute units theoretically are.

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

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.