Training Run Operations

It's 3:15 AM. Your training cluster dies. You check the storage bucket: yesterday's weights are intact. You load them into your script, hit resume, and watch yesterday's training samples replay from the start while the learning rate restarts its warmup from zero. The model weights survived, but the training run died.

We'll use a concrete run named policy-sft-v4 to walk through every recovery boundary: four data-parallel replicas, two sequences per microbatch, eight microbatches per accumulation window, and a committed checkpoint at update 1840. The same operational rules govern full-weight pretraining and parameter-efficient fine-tuning (LoRA): which state survived, which data sample comes next, and did your replacement hardware alter the optimization trajectory?

Weights are only one part of a continuation

Training loops introduced the optimizer and scheduler. A restart must restore their internal state, not merely reconstruct the model architecture. You need to distinguish three operations before calling a checkpoint loader:

For AdamW, optimizer state includes the first and second gradient moments ($m$ and $v$) alongside coordinate step counters ^{[1]Reference 1Decoupled Weight Decay Regularization.https://arxiv.org/abs/1711.05101}. For mixed precision, you must include the gradient scaler when the recipe uses one. Save per-rank random number generator (RNG) states and the data loader resumable cursor. Preserve the model, tokenizer, chat template, dataset slice, and training arguments needed to interpret them. Evaluation history helps select a checkpoint, but it isn't a substitute for optimization state.

Adapter checkpoints also require the correct frozen base weights and adapter configuration. A merged inference export isn't a resumable adapter checkpoint. Tooling frameworks distinguish intermediate training checkpoints from final exports; torchtune documents separate resume requirements for full-weight and LoRA recipes ^{[2]Reference 2Checkpointing in torchtune.https://meta-pytorch.org/torchtune/stable/deep_dives/checkpointer.html}.

Continuation isn't a guarantee of bitwise replay. A different cluster topology can change all-reduce reduction trees, random streams, and data partition assignments. Restore the intended state, document any topology changes, and compare loss trajectories against an acceptable tolerance. Exact replay demands strict control over CUDA kernels, non-deterministic operations, and sample batching.

What lost optimizer moments break

What happens when you restore weights but let AdamW start with zeroed moments? AdamW normalizes gradient steps by running second moments: $\theta_{t+1} = \theta_t - \eta_t \left( \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda \theta_t \right)$ When $m$ and $v$ reset to zero at step 1840, the bias-corrected variance $\hat{v}_1$ collapses to $(1 - \beta_2) g_1^2 / (1 - \beta_2) = g_1^2$. In coordinates where gradients are tiny, dividing by $\sqrt{g_1^2} + \epsilon$ creates massive, erratic step sizes. In steep directions where past variance previously damped updates, the optimizer stumbles. The loss curve shows an artificial jump that wastes hundreds of steps re-accumulating historical momentum.

This scalar Adam example minimizes $(w-y)^2/2$ over a repeating target sequence. It tracks first and second moments, bias-correction steps, a step-based learning rate, and a consumed data cursor. JSON round-tripping models a serialized restart:

Even keeping the step counter doesn't reconstruct the lost moments. In a real cluster restore drill, compare the next sample IDs, current learning rate, optimizer moments, and next gradient update against an uninterrupted reference run. A smooth-looking loss curve doesn't prove mathematical continuity.

Fault tolerance, cluster MTBF, and asynchronous checkpointing

Training large models requires reckoning with cluster reliability physics. An individual GPU server might average a Mean Time Between Failures (MTBF) of roughly 26,000 hours (about 3 years). On a small 32-GPU pool, an unrecoverable failure strikes once every 800 hours (about a month).

At scale, the failure arithmetic changes drastically. In the Meta Llama 3 405B training run, 16,384 H100 GPUs trained concurrently ^{[3]Reference 3The Llama 3 Herd of Models.https://arxiv.org/abs/2407.21783}. Cluster MTBF scales inversely with device count: $\text{MTBF}_{\mathrm{cluster}} = \frac{\text{MTBF}_{\mathrm{single}}}{N_{\mathrm{GPUs}}} = \frac{26000\text{ hours}}{16384} \approx 1.59\text{ hours}$ When hardware faults strike every 90 to 120 minutes, synchronous blocking checkpointing becomes an operational bottleneck. If each checkpoint takes 12 minutes to serialize and upload, saving every hour consumes 20% of your cluster budget in pure idle time, with another 10% lost replaying uncommitted steps.

Modern distributed frameworks solve this dilemma through a three-tier asynchronous checkpointing pipeline:

1. Tier 1 (GPU HBM): Live model parameters and optimizer states mutate during forward and backward passes.
2. Tier 2 (Host Pinned CPU Memory): At the checkpoint step, each rank copies its sharded state tensors across PCIe into host pinned memory using non-blocking CUDA streams. On PCIe Gen5 links, copying 100 GB of sharded weights and optimizer state takes roughly 3 to 5 seconds. Once the device-to-host (D2H) copy finishes, the GPU unblocks and immediately begins the next microbatch.
3. Tier 3 (Distributed Remote Storage): An asynchronous background thread pool streams host memory buffers over the storage fabric (Lustre, Ceph, GPFS, or object store buckets) over the next several minutes without taking GPU compute cycles.

Atomic manifest publishing and data cursor alignment

Never write sharded files directly into the active checkpoint directory. An unexpected node crash mid-write leaves truncated files that crash subsequent restore attempts.

Each rank writes its local shard into an isolated staging directory (such as checkpoints/step_1840.tmp/). Ranks flush and fsync their files. When all ranks confirm completion, Rank 0 publishes an authoritative manifest file or atomically updates a directory symlink (checkpoints/latest -> step_1840). If preemption or a host kernel panic hits during storage writes, the recovery job ignores the incomplete .tmp directory and rolls back cleanly to the last committed manifest pointer.

A clean weight restore still fails if the data cursor drifts. PyTorch data loaders with multi-worker prefetching buffer dozens of batches ahead in host memory. Recording the main thread iteration number doesn't tell you how many samples were actually consumed by the model before preemption. The checkpoint manifest must store the deterministic data loader state: the random epoch seed, the active shard file index, and the exact byte offset or sample ID processed at the completed optimizer boundary.

Preemption is a strict deadline, not a save command

Suppose update 1840 is durably committed to storage, and a preemption warning signal (such as a cloud spot termination notice or a cluster scheduler SIGTERM) arrives after microbatch 5 of the next 8-microbatch accumulation window.

Your training policy saves only at completed optimizer boundaries. There are two valid operational outcomes:

1. All nodes remain healthy and there's enough time to complete the remaining 3 microbatches, perform the optimizer step, stage the checkpoint to host RAM, and update the manifest for step 1841.
2. Time runs out or a rank dies abruptly. The orchestrator terminates, discards partial in-flight state, and reboots from durable step 1840, rewinding the data cursor to sample 117,760.

Never pair weights from step 1840 with a data cursor that advanced past microbatch 5. That mistake silently drops the gradient contributions of those five microbatches without ever applying them to the model weights. Saving mid-window requires persisting unscaled accumulated gradients, the microbatch accumulation counter, and exact intermediate RNG states. Unless your framework explicitly supports and tests mid-window recovery, stick to completed update boundaries.

Signal handlers should only set an internal shutdown flag, not initiate distributed network I/O. The main training loop coordinates the save at the next clean accumulation boundary. When an uncoordinated hardware failure kills a rank, surviving ranks can't complete collective all-reduce operations; periodic durable checkpoints remain mandatory even with preemption warning systems.

Budget the entire remaining path

A 90-second checkpoint doesn't fit into a 90-second warning if you also need to finish active microbatches and publish manifests. The preemption budget equation is: $T_{\mathrm{warning}} \ge T_{\mathrm{window}} + T_{\mathrm{save}} + T_{\mathrm{margin}}$ Here $T_{\mathrm{save}}$ includes time through durable publication, and $T_{\mathrm{margin}}$ absorbs network and disk jitter:

A blocking save every five minutes of compute spends 23.1% of cluster wall-clock time writing files. Increasing the interval to twenty minutes drops save overhead to 7.0%, but exposes up to twenty minutes of uncheckpointed work to hardware failure. Choose your checkpoint interval based on measured save latency, cluster MTBF, and acceptable replay cost.

Sharded saves need a compatible restore path

In distributed training with Fully Sharded Data Parallel (FSDP) or DeepSpeed ZeRO-3, no individual GPU holds the complete set of parameters or optimizer moments ^{[4]Reference 4ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.https://arxiv.org/abs/1910.02054}. This raises two operational questions: did the save complete reliably, and can your restore loader interpret the sharded layout?

The naive approach gathers all shards to Rank 0 to write a single monolithic state dictionary. For a 70B parameter model, the weights take 140 GB in BF16, and the AdamW moments take 280 GB in FP32. Gathering 420 GB of state onto Rank 0 causes an instant Out-Of-Memory (OOM) crash or stalls the cluster for twenty minutes over host PCIe busses.

PyTorch Distributed Checkpoint (DCP) avoids this bottleneck by saving sharded state directly from each GPU ^{[5]Reference 5PyTorch Distributed Checkpoint (DCP) Recipehttps://docs.pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html}. Each rank writes its local tensor slice alongside global coordinate metadata: the fully qualified parameter name, global tensor dimensions, data type, and the coordinate bounding box (such as [0:4096, 0:2048]).

Resharding across topologies without single-node bottlenecks

When a cluster shrinks or grows (for instance, resuming a 4-node run on 8 nodes after replacement hardware arrives, or dropping from 8 nodes to 7 after a node failure), your loader must reshard parameters across different world sizes ($N \ne M$).

Under DCP, resharding happens without central gathering:

1. The destination model initializes its partitioned DTensor structures according to the new $M$-rank topology.
2. Each destination rank consults the shared .metadata index file.
3. Each rank calculates which source shard files intersect its assigned tensor slice.
4. All $M$ ranks read their required byte ranges directly and concurrently from shared storage.

DCP doesn't guarantee backward compatibility across arbitrary PyTorch major versions ^{[6]Reference 6Distributed Checkpointhttps://docs.pytorch.org/docs/stable/distributed.checkpoint.html}. Exercise a 4-to-8 rank restore drill on your target software stack before launching multi-week training. DeepSpeed Universal Checkpointing provides similar cross-topology conversion for ZeRO runs, requiring an explicit conversion step rather than reading arbitrary rank files ^{[7]Reference 7Universal Checkpointing with DeepSpeed: A Practical Guide.https://www.deepspeed.ai/tutorials/universal-checkpointing/}.

This code demonstrates the coordinate-indexing principle by slicing a 16-element vector into four source shards and reconstructing eight destination shards:

Keep the update size separate from the GPU count

For complete accumulation windows with equal microbatch sizes: $B_{\mathrm{global}} = B_{\mathrm{micro}} \times K \times D_{\mathrm{replica}}$ Here $B_{\mathrm{micro}}$ is the sequence count per device, $K$ is the gradient accumulation steps per update, and $D_{\mathrm{replica}}$ is the count of independent data-parallel replicas.

Don't confuse total physical GPUs with data-parallel degree. In 3D parallelism, tensor-parallel ($\mathrm{TP}$) and pipeline-parallel ($\mathrm{PP}$) workers collaborate on a single model instance: $D_{\mathrm{replica}} = \frac{N_{\mathrm{GPUs}}}{\mathrm{TP} \times \mathrm{PP} \times \mathrm{CP}}$ Eight GPUs configured with $\mathrm{TP}=2$ and $\mathrm{PP}=1$ form $D=4$ replicas, not eight. Halving accumulation steps while keeping $D=4$ cuts the global batch in half.

Sequences versus supervised tokens and loss reduction

Sequence counts don't equal token throughput. In instruction tuning and conversational SFT, sequence packing, system prompt masking, and variable response lengths mean that two batches with identical sequence counts carry different numbers of loss-bearing tokens: $\text{Tokens per update} = \sum_{i=1}^{B_{\mathrm{global}}} \text{Length}_{\mathrm{unmasked}}(i)$ If replica 0 processes 4,800 supervised tokens while replica 1 processes 9,600, averaging their per-device mean loss values equally computes $\frac{1}{2}(\bar{\mathcal{L}}_0 + \bar{\mathcal{L}}_1)$. That naive average assigns twice as much weight to tokens on replica 0 as tokens on replica 1, distorting the true gradient.

Distributed loss computation must use global token-weighted reduction: $\mathcal{L}_{\mathrm{global}} = \frac{\sum_{r} \sum_{j} \mathcal{L}_{r, j}}{\sum_{r} N_{\mathrm{tokens}, r}}$ Frameworks handle this by executing an all-reduce sum over unmasked token counts before dividing the accumulated loss. Don't divide by accumulation steps twice when the framework already normalizes by global tokens.

A 4-replica job with microbatch 2 and accumulation 8 restarts on eight GPUs configured as tensor parallel 2 and data parallel 4. Should you halve accumulation to preserve the batch?

Answer

No. Data parallelism is still four, so the batch remains 2 × 8 × 4 = 64. Halving accumulation would reduce it to 32. Count independent replicas, not devices.

Preserve the scheduler before tuning a new batch

When continuing a run at the same effective batch size, restore the scheduler phase directly. Don't restart warmup just because the process restarted.

If you intentionally double the global batch size, you're changing the optimization problem. At a fixed token budget, doubling sequences per update halves the total number of updates. A learning rate schedule defined in steps will now traverse through warmup and decay twice as fast relative to data volume. Always log whether schedules are indexed by optimizer updates, examples, or total consumed tokens.

Why AdamW rejects linear learning rate scaling

The popular linear scaling rule ($\eta \propto B$) originated in large-batch stochastic gradient descent (SGD) experiments for computer vision. Goyal et al. scaled ResNet-50 minibatch size up to 8,192 on 256 GPUs by multiplying the learning rate by $k$ when multiplying batch size by $k$, pairing the rule with gradual warmup ^{[8]Reference 8Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour.https://arxiv.org/abs/1706.02677}. In SGD, a $k$-times larger batch reduces gradient noise variance by $1/k$. Taking a step $k$ times larger in the same direction approximates taking $k$ consecutive small-batch steps when curvature is gentle.

That logic doesn't hold for AdamW SFT:

1. Adaptive coordinate normalization: AdamW divides gradients by running second moment roots ($\sqrt{v_t} + \epsilon$). Increasing batch size shrinks gradient variance, but the magnitude of normalized updates remains bounded by coordinate geometry.
2. Square root scaling: Empirical and theoretical studies demonstrate that for adaptive optimizers, learning rate scaling typically follows a square root rule ($\eta \propto \sqrt{B}$) or an exponent between 0.5 and 0.7. Doubling the batch and doubling AdamW's learning rate frequently triggers divergence or sharp quality degradation.
3. Critical batch size limits: McCandlish et al. formalized the critical batch size $B^*$, governed by the gradient noise scale ^{[9]Reference 9An Empirical Model of Large-Batch Training.https://arxiv.org/abs/1812.06162}: $\mathcal{B}_{\mathrm{noise}} = \frac{\text{Tr}(\Sigma)}{|G|^2}$ When batch size is below $B^*$, increasing batch size yields roughly linear speedup in wall-clock time. As batch size approaches and exceeds $B^*$, gradient noise is already saturated; further batch increases yield diminishing returns where more compute no longer reduces the total optimization steps needed.

Hardware telemetry and run diagnostics: silent killers at scale

A loss spike doesn't reveal its cause on its own. On large GPU clusters, hardware degradation and silent data anomalies cause subtle failures that look like modeling divergence.

In bulk-synchronous parallel execution (All-Reduce or Reduce-Scatter), every GPU in the replica group waits at the barrier until the slowest GPU finishes. A single GPU throttling its clocks from 1980 MHz to 1050 MHz due to fan failure or VRM overheating drags down the entire 16,384-GPU cluster. Continuous telemetry must track per-rank backward pass latency, NCCL barrier wait times, and GPU core clock distributions.

InfiniBand link flapping presents an even trickier failure mode. A dirty optical connector or loose fiber transceiver produces intermittent physical bit errors. The link renegotiates or drops from 400 Gbps (NDR) to 50 Gbps without throwing an unhandled socket exception. The job stays alive, but all-reduce latency jumps tenfold. Monitor physical fabric counters to evict degraded nodes automatically.

Silent Data Corruption (SDC) occurs when hardware defects or cosmic rays cause bit flips in compute ALUs without triggering ECC memory alerts ^{[3]Reference 3The Llama 3 Herd of Models.https://arxiv.org/abs/2407.21783}. An undetected bit flip during matrix multiplication can produce an erroneous activation or poisoned gradient. Canary batches (running an identical small matrix multiplication periodically across all ranks and verifying checksums) catch corrupt silicon before poisoned weights propagate to checkpoints.

Automated loss-spike rollback tripwires

When a loss spike hits at 3:00 AM, waiting for a human on-call engineer burns thousands of idle GPU hours. Modern training harnesses implement automated tripwire circuits:

1. Anomaly tripwire: If step loss exceeds the rolling median by three standard deviations ($\mathcal{L}_t > \text{median}(\mathcal{L}_{t-50:t}) + 3\sigma$), or if unclipped gradient norm exceeds a threshold ($||g||_2 > g_{\mathrm{max}}$), pause the run immediately.
2. Automated rewind: The orchestrator rewinds model and optimizer state to two checkpoints prior ($t - 2$), well before the destabilizing batch entered the accumulation window.
3. Data quarantine: The offending data partition is logged, quarantined, and skipped in the data loader.
4. Seed perturbation: The data loader perturbs its shuffle seed, and training resumes automatically. If a second spike hits immediately, the orchestrator triggers an alert and halts the job safely.

Choose the objective, trainable weights, and precision separately

Training recipes aren't a single dropdown choice. SFT and continued pretraining define what the model learns. LoRA defines which parameters update. QLoRA adds 4-bit base-model quantization to adapter tuning. Distillation defines learning from teacher logits ^{[10]Reference 10Distilling the Knowledge in a Neural Network.https://arxiv.org/abs/1503.02531}.

Domain-adaptive pretraining can improve specialized comprehension, but unfamiliar vocabulary alone doesn't prove it's necessary ^{[11]Reference 11Don't Stop Pretraining: Adapt Language Models to Domains and Tasks.https://aclanthology.org/2020.acl-main.740/}. Small labeled datasets make LoRA SFT an attractive starting point, though adapters aren't guaranteed to beat full fine-tuning ^{[12]Reference 12LoRA: Low-Rank Adaptation of Large Language Models.https://arxiv.org/abs/2106.09685}.

Reject impossible memory plans before launching

The original QLoRA paper demonstrated fine-tuning a 65B model on a single 48GB GPU ^{[13]Reference 13QLoRA: Efficient Finetuning of Quantized LLMshttps://arxiv.org/abs/2305.14314}. That empirical result doesn't mean a 70B model fits into a consumer 24 GiB GPU.

At an ideal four bits per parameter, a 70B model requires: $\text{Memory}_{\mathrm{base}} = \frac{70 \times 10^9 \times 4\text{ bits}}{8\text{ bits/byte}} = 35 \times 10^9\text{ bytes} \approx 32.6\text{ GiB}$ Base weights alone exceed 24 GiB by 8.6 GiB before allocating a single byte for LoRA adapters, activations, gradients, optimizer moments, or KV caches.

For 70B models on 24 GiB hardware, you need multi-GPU sharding or CPU offloading. Passing this arithmetic check doesn't guarantee your job fits; failing it guarantees an OOM crash.

Rehearse the restart before the long run

Don't wait for your first midnight preemption to discover that your checkpoint loader crashes. Run a restart drill on small test jobs:

• Local state continuity: Compare five uninterrupted steps against five steps resumed from an intermediate checkpoint. Verify that weights, Adam moments, learning rates, and data sample IDs match within floating-point tolerance.
• Distributed topology restoration: Save a checkpoint on four GPUs, then load it onto eight GPUs (or six GPUs). Ensure your DCP or ZeRO loader resshards state cleanly without Rank-0 memory exhaustion.
• Preemption deadline drill: Trigger a simulated termination signal (kill -15). Verify that in-flight accumulation windows finish within the safety margin, the atomic manifest pointer updates, and partial .tmp files are ignored if killed prematurely.
• Hardware telemetry and tripwire check: Verify that stragglers and simulated loss anomalies trigger automated rollback and quarantine routines rather than hanging the cluster.

Distinguish between your latest resumable checkpoint and your best evaluation checkpoint. Step 1840 might be your operational recovery point, while step 1600 holds your highest validation score. Store both with explicit metadata tags.