LLMTailor: A Layer-wise Tailoring Tool for Efficient Checkpointing of Large Language Models

Abstract: Checkpointing is essential for fault tolerance in training LLMs. However, existing methods, regardless of their I/O strategies, periodically store the entire model and optimizer states, incurring substantial storage overhead and resource contention. Recent studies reveal that updates across LLM layers are highly non-uniform. Across training steps, some layers may undergo more significant changes, while others remain relatively stable or even unchanged. This suggests that selectively checkpointing only layers with significant updates could reduce overhead without harming training. Implementing such selective strategies requires fine-grained control over both weights and optimizer states, which no current tool provides. To address this gap, we propose \texttt{LLMTailor}, a checkpoint-merging framework that filters and assembles layers from different checkpoints to form a composite checkpoint. Our evaluation indicates that LLMTailor can work with different selective checkpointing strategies and effectively reduce checkpoint size (e.g., 4.3 times smaller for Llama3.1-8B) and checkpoint time (e.g., 2.8 times faster for Qwen2.5-7B) while maintaining model quality.

• The paper introduces LLMTailor, a YAML-driven framework that merges individual transformer layers, optimizer states, auxiliary modules, and metadata into fully resumable LLM checkpoints.
• Parity merging cuts checkpoint size by about 50% and reduces checkpoint-time ratios by roughly 40%, while resumed losses and benchmark scores closely match uninterrupted training.
• Filtered checkpointing achieves up to 4.3× lower storage and 2.8× lower checkpoint-time ratios, but inconsistent quality results show that adaptive layer-selection policies remain necessary.

Checkpointing is the dominant fault-tolerance mechanism in LLM training, yet existing systems uniformly save all model and optimizer state at every checkpoint interval, incurring substantial I/O overhead—reported to consume between 12% and 43% of total training time. The paper introduces LLMTailor, a checkpoint-merging framework that assembles fully resumable checkpoints by selecting individual layers (with their associated optimizer states) from multiple source checkpoints (2602.22158). Its central premise is that LLM layers update non-uniformly during training, so selectively checkpointing only significantly changing layers can reduce overhead without harming model quality.

Motivation: non-uniform layer updates

The authors ground their approach in prior observations of layer heterogeneity: distinct transformer layers encode different linguistic information, fine-tuned models show block-structured representation similarity across layers, and metrics such as $PL\_Alpha\_Hill$ demonstrate that layers train at different speeds under SGD and Adam. These findings motivate partial checkpointing—for example, saving half the layers per interval and merging two partial checkpoints into one complete state. Realizing this requires fine-grained manipulation of both weights and optimizer states, which no existing tool provides. MergeKit, the closest system, merges weights only via its passthrough method and omits optimizer files, auxiliary layers (embeddings, lm_head), and configuration metadata, making resumed training impossible.

Design

LLMTailor extends MergeKit's YAML-driven interface to full checkpoint merging through three mechanisms:

Separable optimizer construction. Optimizer files store flattened tensors organized into coarse parameter groups (typically two groups split by weight-decay eligibility), which cannot be split along layer boundaries. LLMTailor regroups parameters before training begins so that each transformer layer maps to two parameter groups—one weight-decay-exempt, one not—expanding the group count from 2 to $2L + x$ for $L$ transformer layers and $x$ auxiliary layers. Because neither parameter values nor hyperparameters change, training semantics are preserved; only a small computational cost is added.

Optimizer merging. Given a YAML recipe specifying base model, source layers, and target positions, LLMTailor indexes each layer's parameter groups across ZeRO-3 shards and copies them into the assembled checkpoint. The dominant cost is I/O over up to $N \times (L+3)$ optimizer shard files for $N$ GPUs; the tool parallelizes shard loading with Python's ProcessPoolExecutor to reduce latency.

Auxiliary layers and configuration. Unlike MergeKit, LLMTailor explicitly splits and merges embed_tokens, lm_head, and normalization modules, and copies configuration/metadata files (training step, learning rate) from the most recent checkpoint to preserve training continuity.

Evaluation

Experiments run on an 8×A100-80GB cluster with DeepSpeed ZeRO-3 and AdamW, evaluating Llama-3.2-1B, Llama-3.1-8B, and Qwen2.5-7B on continual pre-training (PubMed-Summarization) and supervised fine-tuning (MedQA). Two use cases are studied against default transformers-library checkpointing.

Parity merging alternates odd/even layers between consecutive checkpoints, halving storage per interval. This reduces checkpoint size by ~50% for both models and cuts the checkpoint-time-to-training-time ratio from 4.99% to 3.03% (Llama3.1-8B) and from 20.63% to 12.76% (Qwen2.5-7B)—roughly a 40% time reduction. Resumed training matches the original final train/eval losses exactly (1.58/1.60 for Qwen SFT; 1.58/1.58 for Llama CPT), and zero-shot benchmark scores (MMLU, MedMCQA, MedQA, PubMedQA) remain within noise of the uninterrupted baseline, confirming correctness of merged optimizer states.

Filtered checkpointing saves only the first and last two layers at every interval and half the remaining layers at $5\times$ the original interval, exploiting evidence that deeper layers contribute less to reasoning. This yields the paper's strongest results: 4.3× smaller total checkpoint storage for Llama3.1-8B (1799.52 GB → 420 GB) and a 2.8× reduction in checkpoint-time ratio for Qwen2.5-7B (20.63% → 7.26%). Quality outcomes are mixed and the authors are candid about this: in the SFT task, the filtered Qwen checkpoint degrades noticeably (e.g., MMLU_med drops from 89.00 to 84.00; MMLU from 73.14 to 71.64), while in the CPT task the filtered Llama3.1 model slightly outperforms the baseline (MMLU 62.06 vs. 60.00). Final losses also shift marginally (e.g., 1.60 vs. 1.58 train loss for Qwen SFT). The implication is that rule-based filtering trades some quality risk for large overhead savings, and that more dynamic layer-selection strategies could improve robustness—an open question rather than a settled result.

Merge overhead. Assembling a composite checkpoint costs seconds to minutes depending on model size and merge pattern: e.g., 332 s for a two-checkpoint merge of Llama3.1-8B versus 16.8 s for loading a single baseline checkpoint, rising to 1027 s for interleaved parity-style access because optimizer shards must be fully loaded before any layer is accessible (no lazy loading is possible). The authors argue this overhead is small relative to multi-hour or multi-day training runs, and note it would shrink substantially under a natively layer-wise checkpointing format.

Limitations and open questions

The paper concedes several constraints. First, the prototype operates only on local checkpoints; integration with other checkpointing frameworks remains future work. Second, the evaluation covers post-training workloads (CPT/SFT) on medical datasets at 7–8B scale; whether selective checkpointing holds during full pre-training, where update dynamics differ, is untested. Third, the filtered strategy's inconsistent quality results (degradation in SFT, improvement in CPT) indicate that static, hand-crafted selection rules are insufficient, leaving open how to design dynamic, quality-aware policies for choosing which layers to checkpoint and when. Finally, the interleaved-loading penalty exposes a structural limitation: without lazy loading of optimizer shards, merge cost scales poorly with the number of source checkpoints.

Conclusion

LLMTailor demonstrates that layer-granular checkpoint composition—including optimizer states, auxiliary layers, and configuration—is feasible and preserves resumability. Its two use cases show up to 4.3× storage and 2.8× checkpoint-time reductions with largely maintained model quality, establishing partial checkpointing as a viable complement to existing I/O optimizations such as asynchronous, compressed, and in-memory checkpointing. The main unresolved issue is replacing fixed selection heuristics with adaptive strategies whose quality guarantees hold across training regimes.