LLMTailor

LLMTailor is an enhanced fork of mergekit, designed for layer-wise merging of large language models (LLMs) with extended support for:

✅ Compatible with our new checkpoint system StreamCheck ✅ Layer-wise model merging & selection
✅ Optimizer state reconstruction (supports ZeRO-3 shards)
✅ Tokenizer & embedding adaptation: these auxiliary layers in LLMs could also be selected and merged now ✅ Backward compatibility with most mergekit plans

Note: LLMTailor retains most of mergekit’s original merging capabilities while adding extensions (llmtailor.* fields in YAML) for training-oriented scenarios.

---

Citing LLMTailor

The relevant research paper will be published at PDSW25. If you reference or use LLMTailor in your research, please cite:

@inproceedings{10.1145/3731599.3767515,
author = {Sun, Minqiu and Huang, Xin and Guo, Luanzheng and Tallent, Nathan R. and Sato, Kento and Dai, Dong},
title = {LLMTailor: A Layer-wise Tailoring Tool for Efficient Checkpointing of Large Language Models},
year = {2025},
isbn = {9798400718717},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
url = {https://doi.org/10.1145/3731599.3767515},
doi = {10.1145/3731599.3767515},
abstract = {Checkpointing is essential for fault tolerance in training large language models (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 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.},
booktitle = {Proceedings of the SC '25 Workshops of the International Conference for High Performance Computing, Networking, Storage and Analysis},
pages = {1366–1374},
numpages = {9},
keywords = {Checkpoint, Large Language Model, I/O optimization},
location = {
},
series = {SC Workshops '25}
}

Installation

Required Software

• Python 3.11

conda create -n myenv python=3.11

conda activate myenv

• Clone From GitHub

git clone https://github.com/SunMinqiu/LLMTailor.git
cd LLMTailor
pip install -r requirements.txt
pip install -e .

Required Hardware

• GPU:

• CPU: At least 32 cores

• Memory: At least 200 GB

• Storage: Depending on the model and training epochs, recommend at least at least 350 GB for a 7B model and 700 GB for 14B model.

• Benchmark Running For the benchmark, we use the open source project called lm-evaluation-harness. Please follow the instructions of this project to install.

Required Hardware

• GPU: Recommend at least one node of 8 * L40s, or 4 * H100.
• CPU: At least 64 cores.
• Memory: At least 200 GB.
• Storage: Depending on the model and training epochs, recommend at least at least 350 GB for a 7B model and 700 GB for 14B model.

Quick Start

LLMTailor supports three usage modes:

Mode 1: Use Your Own YAML Config (Manual Merge)

If you already have a YAML configuration file, you can run the merge directly:

python examples/start_merge.py \
    --config_yml "path/to/your_config.yaml" \
    --output_path "path/to/output_model" \
    --num_gpu 8

Example YAML config (examples/Qwen_cut.yaml):

slices:
  - sources:
      - model: /path/to/checkpoint-100
        layer_range: [0, 14]
  - sources:
      - model: /path/to/checkpoint-200
        layer_range: [14, 28]

merge_method: passthrough
dtype: bfloat16

Mode 2: Generate YAML Only (Dry Run)

Use --dry_run to generate the YAML config from StreamCheck logs without running the merge. This is useful for reviewing the configuration before merging.

python examples/start_merge.py \
    --streamcheck_json "path/to/checkpoint_flags.jsonl" \
    --streamcheck_path "path/to/checkpoints" \
    --config_yml "output_config.yaml" \
    --failure_step 500 \
    --total_layers 28 \
    --dry_run

This will:

1. Parse the checkpoint_flags.jsonl to find the best checkpoint for each layer
2. Generate a YAML config file at the specified path
3. Exit without running the merge

You can then review the generated YAML and run Mode 1 to execute the merge.

Mode 3: One-Stop Service (Generate + Merge)

Run the complete pipeline: generate YAML from StreamCheck logs and execute the merge in one command.

python examples/start_merge.py \
    --streamcheck_json "path/to/checkpoint_flags.jsonl" \
    --streamcheck_path "path/to/checkpoints" \
    --config_yml "output_config.yaml" \
    --output_path "path/to/merged_model" \
    --failure_step 500 \
    --total_layers 28 \
    --num_gpu 8

Using the Shell Script

The shell script examples/start_merge.sh supports all parameters via command line:

# Show all available options
bash examples/start_merge.sh --help

Mode 1: Use existing YAML config

bash examples/start_merge.sh \
    --config_yml /path/to/your_config.yaml \
    --output_path /path/to/merged_model \
    --num_gpu 8

Mode 2: Generate YAML only (dry run)

bash examples/start_merge.sh \
    --streamcheck_json /path/to/checkpoint_flags.jsonl \
    --streamcheck_path /path/to/checkpoints \
    --config_yml /path/to/output_config.yaml \
    --failure_step 500 \
    --total_layers 28 \
    --dry_run

Mode 3: One-stop service (all parameters)

bash examples/start_merge.sh \
    --streamcheck_json /path/to/checkpoint_flags.jsonl \
    --streamcheck_path /path/to/checkpoints \
    --config_yml /path/to/output_config.yaml \
    --output_path /path/to/merged_model \
    --failure_step 500 \
    --total_layers 28 \
    --num_gpu 8 \
    --lora_merge_cache /tmp/cache \
    --copy_tokenizer true \
    --lazy_unpickle false \
    --low_cpu_memory false

Command Line Arguments

Argument	Description	Default
--config_yml	Path to YAML merge config	examples/Qwen_cut.yaml
--output_path	Output path for merged model	Required
--streamcheck_json	Path to checkpoint_flags.jsonl	None
--streamcheck_path	Path to checkpoints directory	None
--failure_step	Step number for base checkpoint	160
--total_layers	Number of transformer layers	28
--num_gpu	Number of GPUs to use	8
--dry_run	Only generate YAML, skip merge	False
--copy_tokenizer	Copy tokenizer to output	True
--lazy_unpickle	Low-memory model loader	False
--low_cpu_memory	Use when VRAM > RAM	False