LayerCheck: Adaptive Layer-wise Checkpointing for Large Language Model Post-training

• Fixed-interval checkpointing trades I/O and synchronization overhead against rollback recomputation, while prior pipeline, compression, and selective-persistence approaches retain scaling or recovery costs.Checkpointing can consume up to 12% of total training time and as much as 43% in extreme cases; existing approaches also face widening GPU–storage gaps, unchanged I/O volume, or tracking and recovery overheads.Source: §1 · p. 2
• Because transformer layers evolve at different paces, LayerCheck persists layers only when their accumulated drift is significant and reconstructs recovery state by merging each layer’s latest persisted version.This avoids writing the entire model at every checkpoint and avoids retraining or replay during reconstruction.Source: §1 · pp. 2–3
• LayerCheck accumulates per-layer drift, tracks lightweight layer-level update statistics, and uses an update-scaled threshold alongside an independent staleness safeguard.If a layer reaches the user-specified staleness budget Smax, it is persisted even when its drift is below the threshold.Source: §1 · p. 2
• LayerCheck reduces total checkpoint size by up to 22.6×, improves total training time by up to 1.31×, and reduces recovery cost by up to 3.4× versus state-of-the-art checkpointing systems.The evaluation covers multiple open-source LLMs and reports no measurable model-quality degradation in the evaluated settings.Source: §1 · p. 2
• Post-restart loss deviates from the failure-free trajectory by at most 0.54%, while downstream benchmark scores match failure-free fine-tuning.A per-layer staleness guard bounds the age of unsaved layers; under this invariant, recovery does not change the asymptotic convergence rate.Source: §1 · pp. 2–3

2.1 Problem Setup and Goals

• LayerCheck targets fail-stop LLM post-training by restoring each layer from its latest checkpoint, including associated optimizer state and configuration, to reduce I/O while preserving recovery behavior and final quality.Source: §2.1 · p. 3

2.2 Layer Structure of LLMs

• Trainable parameters are grouped into N=B+3 module-level layers: embeddings, each transformer block, normalization, and lm_head; Qwen2.5-7B has 28 transformer blocks.In smaller models, lm_head may be weight-tied to embed_tokens to reduce parameter count.Source: §2.2 · p. 3, Figure 1 · p. 4

2.3 Layer-wise Unbalanced Updating

• Across Qwen2.5-7B’s 31 layers, accumulated updates vary across layers and steps, motivating asynchronous layer-wise saves over uniform all-layer checkpointing.Source: §2.3 · p. 3, Figure 2 · p. 4

• LayerCheck targets redundant full-model snapshots by exploiting the observation that Transformer layers drift at different paces.Source: §3 · p. 5
• LayerCheck monitors per-layer drift, persists selected layers’ full weights and optimizer tensors, and merges their latest saved versions after failure.The merged layers may come from different iterations, avoiding the need for delta replay during recovery.Source: Figure 3 · p. 4, §3 · p. 5
• A bounded-staleness invariant (Smax) limits the mixed-timestamp differences among restored layers and their Adam moments.The invariant also underpins the convergence argument in §3.6.Source: §3 · p. 5

3.1 Overall Workflow

• In Figure 3, each iteration block represents the full training state, and each layer box contains that layer’s parameters and associated optimizer tensors.The workflow is shown along the training timeline, with iterations on the x-axis.Source: §3.1 · p. 5
• LayerCheck persists a layer when its accumulated normalized drift since its last save exceeds a threshold, refreshing fast-changing layers more often and skipping slow-changing layers.This selective persistence reduces checkpoint I/O and storage compared with saving the entire model and optimizer state at each checkpoint.Source: §3.1 · p. 5
• At recovery, LayerCheck assembles each layer’s latest successfully persisted state at or before failure, including its corresponding optimizer tensors, and resumes at iteration tfail+1 without rollback.The resulting checkpoint can contain mixed timestamps across layers, while each layer’s parameters and optimizer tensors remain internally consistent.Source: §3.1 · p. 5

3.2 Layer-wise Monitoring and Checkpointing

• LayerCheck tracks drift at layer granularity rather than per parameter, avoiding auxiliary tracking state comparable in size to the model.The stated goal is for recovery from a composite checkpoint to yield statistically indistinguishable model quality from a failure-free run.Source: §3.2 · p. 5
• For each layer, LayerCheck accumulates a normalized update-magnitude estimate over training iterations.The accumulator starts at zero and is updated after the optimizer computes each parameter update.Source: §3.2 · p. 5
• The drift estimate uses absolute values to avoid coordinate cancellation and a mean to reduce outlier sensitivity and stabilize the layer-level signal.This conservative estimate favors persisting slightly more layers over missing meaningful drift.Source: §3.2 · p. 5
• A layer is persisted when its accumulated drift exceeds a threshold scaled from a base value using the current learning rate; persistence writes the full layer state and resets its accumulator.The scaling function is derived from the learning-rate schedule to account for warmup and decay under AdamW.Source: §3.2 · pp. 5–6

3.3 Bounding Mixed-Timestamp Staleness

• Layer staleness is the iterations since a layer’s last persistence, while overall checkpoint staleness is weighted by layer size.For a failure at iteration τ, layer ℓ last persisted at tℓ has staleness τ − tℓ.Source: §3.3 · p. 6
• At Smax, LayerCheck forces the layer’s persistence at the next checkpoint opportunity and resets its drift accumulator, bounding maximum per-layer staleness.The guard addresses slowly drifting layers that might otherwise remain unpersisted for long periods.Source: §3.3 · p. 6
• With K0 = 10^-3, most layers refresh roughly every 200 iterations, while Smax = 1000 is a rarely triggered backstop and parameter-weighted average staleness remains low.Figure 4 shows the maximum driven mainly by a normalization layer from steps 200–1000, while the weighted average better reflects mean staleness across layers.Source: §3.3 · p. 6, Figure 4 · p. 6

3.4.1 Layer-aligned optimizer state layout.

• Optimizer parameter groups mirror the model’s layer hierarchy, enabling per-layer state serialization and reconstruction while preserving weight-decay policy.Each Transformer layer uses separate with- and without-weight-decay groups; auxiliary layers use one group when they contain only one type.Source: §3.4.1 · p. 7

3.4.2 Reusing the optimizer update to compute drift.

• Drift computation piggybacks on the optimizer update path, avoiding extra tensor copies and keeping monitoring overhead low.Source: §3.4.2 · p. 7

3.4.3 Threshold defaults and special-casing embeddings.

• LR-adaptive thresholding uses K0 = 10^-3 for most layers and a larger 10^-2 example threshold for embeddings.Selective persistence is disabled for the first 200 steps, with a full base checkpoint retained as fallback for rarely updated layers.Source: §3.4.3 · p. 7

3.5 Recovery from LayerCheck Checkpoints

• At failure, LayerCheck finds each layer’s latest checkpoint at or before the failure step and assembles the selected states into a full checkpoint.If a layer has no prior layer checkpoint, recovery uses the base full checkpoint.Source: §3.5 · p. 7
• Each selected layer state includes its weights and corresponding optimizer tensors, such as Adam moments.Recovery places each loaded layer state back into its original layout.Source: §3.5 · p. 7
• Per-layer persistence-history metadata drives a reconstruction plan that assembles the latest valid layer states into a DeepSpeed-compatible checkpoint.Training resumes from the reconstructed composite checkpoint.Source: §3.5 · p. 7

3.6.2 Composite recovery model

• Bounded staleness adds a bounded second-moment term to Adam’s standard regret bound without changing its asymptotic average-regret rate.The proof-of-concept bound assumes one recovery and the stated boundedness conditions.Source: §3.6 · p. 7, §3.6.2 · p. 8

3.6.1 Setup and assumptions

• The analysis follows Adam’s standard regret setup with convex per-iteration loss, bounded parameters, a recovered state in the same domain, and zero-initialized moments.It adopts the standard boundedness assumptions and requires 1/√β2 < 1.Source: §3.6.1 · p. 7

3.6.2 Composite recovery model

• Recovery combines each layer’s latest persisted checkpoint with its layer-local optimizer tensors, under a guard limiting every coordinate’s staleness to Smax.The analysis also bounds the recovered parameter and first-moment deviations by D∞ and 2G∞, respectively.Source: §3.6.2 · p. 8

3.6.3 Multiple failures.

• With a fixed number of failures independent of T, the staleness terms sum to O(T), leaving the asymptotic average-regret rate unchanged.Source: §3.6.3 · p. 8

4.1.1 Testbed / 4.1.2 Evaluation Methodology

• Evaluation used a single 8-GPU A100 server, AdamW, three open-source LLMs, and MedQA and OpenThoughts fine-tuning workloads.ZeRO-3 sharded model and optimizer states across GPUs; recovery fidelity was evaluated using post-restart loss trajectories and final downstream benchmark scores.Source: §4.1.1 · p. 9, §4.1.2 · p. 9

4.2 Recovery Fidelity: Training After Mixed-Timestamp Recovery

• After repeated failures at K=10−3, maximum absolute loss deviation was 0.0024 (0.54% relative); runs across three K thresholds stayed close to baseline under the Smax guard.Smaller K refreshed layers more aggressively and reduced post-recovery deviation, while larger K skipped more writes but kept deviation bounded rather than amplifying.Source: Figure 5 · p. 9, §4.2 · p. 10

4.3.1 Time overhead

• Across all three models, LayerCheck had the shortest one-epoch training time on OpenThoughts, with 11.7–20.0% overhead versus checkpoint-free training; it was 1.31× faster than LowDiff on Llama3.2-1B and 3.76× faster than GEMINI on Qwen2.5-7B.Checkpoint I/O dominated the overhead; subtracting checkpoint time, remaining training time essentially matched tnet across all three models.Source: §4.3.1 · pp. 10–11

4.3.2 I/O volume

• LayerCheck reduced total checkpoint storage versus LowDiff by up to 22.6× (average ∼17.1×) and mean per-checkpoint size by up to 6.6× (average ∼5.0×).At K = 10−3, LayerCheck generated many checkpoint events containing only a subset of layers, rather than rewriting every layer each time.Source: §4.3.2 · p. 11, Table 1 · p. 12

4.4 Recovery Overhead

• For Qwen2.5-7B under matched freshness, LayerCheck's t⋆ recovery was 10s versus LowDiff's 34s (3.4× faster).The comparison matched rollback-based methods to LayerCheck's average staleness of 106 training steps and excluded method-specific reconstruction time to preserve cross-method comparability.Source: §4.4 · p. 12, Figure 7 · p. 13

4.5 Scalability under Repeated Failures

• At MTBF=0.5 hours, LayerCheck achieved 92.3% ETTR, versus 88.8% for LowDiff and 84.5% for GEMINI, and had the highest ETTR across all tested MTBF values.ETTR measures the fraction of wall-clock time spent on useful training work rather than checkpointing and recovery.Source: §4.5 · p. 13, Figure 8 · p. 14

4.6 Recovered Model Quality

• Across datasets and model sizes, recovered runs closely matched failure-free final benchmark scores and showed no systematic degradation.Table 3 compares failure-free fine-tuning with fine-tuning followed by an injected failure, recovery, and continued training.Source: §4.6 · p. 13, Table 3 · p. 14

4.7 Summary and Limitations

• Improvement magnitude depends on storage backend, and cross-method recovery comparisons exclude implementation-dependent reconstruction time.The evaluation focuses on AdamW; applying the persistence principle to SGD would require calibrating K0 to the optimizer's update pattern, trading more writes for lower staleness or fewer writes for potentially greater recovery deviation.Source: §4.7 · pp. 14–15

Checkpointing for Deep Learning and LLM Training

• Prior checkpointing systems overlap snapshotting with computation, improve resilience through elasticity or redundancy, or reduce writes and reconstruction through partial, incremental, and differential methods.Examples include CheckFreq, GEMINI, PCcheck, Datastate-LLM [17], FlowCheck, CPR, SCAR [24], and LowDiff.Source: §5 · p. 15
• LowDiff is LayerCheck’s primary online-I/O comparison, while LLMTailor uses fixed layer-merging recipes without online selection or convergence analysis, and filtered variants can degrade quality.LLMTailor is closest in spirit because it assembles resumable checkpoints from existing ones.Source: §5 · p. 15
• LayerCheck selects layers online at each step using accumulated drift and LR-aware thresholds, guarantees bounded staleness for mixed-timestamp recovery, and avoids replay while reducing checkpoint I/O.This distinguishes it from fixed, rule-based layer-merging recipes.Source: §5 · p. 15

Checkpoint Compression and Model Reduction

• Checkpoint compression methods shrink stored states, but are primarily designed for local checkpoints or smaller federated settings rather than online, large-scale distributed LLM training.Examples include ExCP [19], Inshrinkerator [50], QD-Compressor, AutoCheck [7], ADVICE, and replay-based ZOCheck.Source: §5 · p. 15
• Quantization, pruning, and structural compression suggest layer redundancy and sensitivity differ, but LayerCheck selectively persists full layer states online instead of compressing the deployed model.Related structural methods include LLM-Streamline, LLM-Pruner [57], LaCo [58], and ShortGPT.Source: §5 · p. 15

6 Conclusion

• Up to 22.6× reduction in total checkpoint size and up to 1.31× reduction in checkpointing time versus existing SOTA preserve final model quality.Experiments also show shorter end-to-end training time.Source: §6 · p. 15
• LayerCheck tracks layer drift using already-available training signals, adding negligible per-step overhead while selectively persisting layers with significant accumulated drift.Selective persistence reduces checkpoint I/O and shortens end-to-end training time while preserving final model quality.Source: §6 · p. 15
• Integrating LayerCheck with overlap-oriented checkpoint pipelines is identified as a next step to overlap computation, communication, and I/O.This would let reduced checkpoint work benefit from end-to-end overlap.Source: §6 · p. 15