implementing-llms-litgpt

LitGPT - 简洁的 LLM 实现

快速入门

LitGPT 提供 20+ 种预训练 LLM 实现，代码简洁易读，并配备可用于生产环境的训练工作流。

安装：

pip install 'litgpt[extra]'

加载并使用任意模型：

from litgpt import LLM

# Load pretrained model
llm = LLM.load("microsoft/phi-2")

# Generate text
result = llm.generate(
    "What is the capital of France?",
    max_new_tokens=50,
    temperature=0.7
)
print(result)

列出可用模型：

litgpt download list

常见工作流

工作流 1：在自定义数据集上微调

复制此清单：

Fine-Tuning Setup:
- [ ] Step 1: Download pretrained model
- [ ] Step 2: Prepare dataset
- [ ] Step 3: Configure training
- [ ] Step 4: Run fine-tuning

步骤 1：下载预训练模型

# Download Llama 3 8B
litgpt download meta-llama/Meta-Llama-3-8B

# Download Phi-2 (smaller, faster)
litgpt download microsoft/phi-2

# Download Gemma 2B
litgpt download google/gemma-2b

模型会保存到 checkpoints/ 目录。

步骤 2：准备数据集

LitGPT 支持多种格式：

Alpaca 格式（指令-响应）：

[
  {
    "instruction": "What is the capital of France?",
    "input": "",
    "output": "The capital of France is Paris."
  },
  {
    "instruction": "Translate to Spanish: Hello, how are you?",
    "input": "",
    "output": "Hola, ¿cómo estás?"
  }
]

保存为 data/my_dataset.json。

步骤 3：配置训练

# Full fine-tuning (requires 40GB+ GPU for 7B models)
litgpt finetune \
  meta-llama/Meta-Llama-3-8B \
  --data JSON \
  --data.json_path data/my_dataset.json \
  --train.max_steps 1000 \
  --train.learning_rate 2e-5 \
  --train.micro_batch_size 1 \
  --train.global_batch_size 16

# LoRA fine-tuning (efficient, 16GB GPU)
litgpt finetune_lora \
  microsoft/phi-2 \
  --data JSON \
  --data.json_path data/my_dataset.json \
  --lora_r 16 \
  --lora_alpha 32 \
  --lora_dropout 0.05 \
  --train.max_steps 1000 \
  --train.learning_rate 1e-4

步骤 4：运行微调

训练完成后，检查点会自动保存到 out/finetune/。

监控训练：

# View logs
tail -f out/finetune/logs.txt

# TensorBoard (if using --train.logger_name tensorboard)
tensorboard --logdir out/finetune/lightning_logs

工作流 2：在单块 GPU 上进行 LoRA 微调

这是最节省内存的方案。

LoRA Training:
- [ ] Step 1: Choose base model
- [ ] Step 2: Configure LoRA parameters
- [ ] Step 3: Train with LoRA
- [ ] Step 4: Merge LoRA weights (optional)

步骤 1：选择基础模型

GPU 内存有限（12-16GB）时：

• Phi-2（2.7B）- 质量与模型大小的最佳平衡
• Llama 3 1B - 体积最小、速度最快
• Gemma 2B - 推理能力良好

步骤 2：配置 LoRA 参数

litgpt finetune_lora \
  microsoft/phi-2 \
  --data JSON \
  --data.json_path data/my_dataset.json \
  --lora_r 16 \          # LoRA rank (8-64, higher=more capacity)
  --lora_alpha 32 \      # LoRA scaling (typically 2×r)
  --lora_dropout 0.05 \  # Prevent overfitting
  --lora_query true \    # Apply LoRA to query projection
  --lora_key false \     # Usually not needed
  --lora_value true \    # Apply LoRA to value projection
  --lora_projection true \  # Apply LoRA to output projection
  --lora_mlp false \     # Usually not needed
  --lora_head false      # Usually not needed

LoRA 秩指南：

• r=8：轻量级，适配器大小为 2-4MB
• r=16：标准配置，质量良好
• r=32：容量较高，适用于复杂任务
• r=64：最高质量，适配器大小增加 4 倍

步骤 3：使用 LoRA 进行训练

litgpt finetune_lora \
  microsoft/phi-2 \
  --data JSON \
  --data.json_path data/my_dataset.json \
  --lora_r 16 \
  --train.epochs 3 \
  --train.learning_rate 1e-4 \
  --train.micro_batch_size 4 \
  --train.global_batch_size 32 \
  --out_dir out/phi2-lora

# Memory usage: ~8-12GB for Phi-2 with LoRA

步骤 4：合并 LoRA 权重（可选）

将 LoRA 适配器合并到基础模型中，以便部署：

litgpt merge_lora \
  out/phi2-lora/final \
  --out_dir out/phi2-merged

现在使用合并后的模型：

from litgpt import LLM
llm = LLM.load("out/phi2-merged")

工作流 3：从头开始预训练

使用你的领域数据训练新模型。

Pretraining:
- [ ] Step 1: Prepare pretraining dataset
- [ ] Step 2: Configure model architecture
- [ ] Step 3: Set up multi-GPU training
- [ ] Step 4: Launch pretraining

步骤 1：准备预训练数据集

LitGPT 需要经过分词的数据。使用 prepare_dataset.py：

python scripts/prepare_dataset.py \
  --source_path data/my_corpus.txt \
  --checkpoint_dir checkpoints/tokenizer \
  --destination_path data/pretrain \
  --split train,val

步骤 2：配置模型架构

编辑配置文件或使用现有配置：

# config/pythia-160m.yaml
model_name: pythia-160m
block_size: 2048
vocab_size: 50304
n_layer: 12
n_head: 12
n_embd: 768
rotary_percentage: 0.25
parallel_residual: true
bias: true

步骤 3：设置多 GPU 训练

# Single GPU
litgpt pretrain \
  --config config/pythia-160m.yaml \
  --data.data_dir data/pretrain \
  --train.max_tokens 10_000_000_000

# Multi-GPU with FSDP
litgpt pretrain \
  --config config/pythia-1b.yaml \
  --data.data_dir data/pretrain \
  --devices 8 \
  --train.max_tokens 100_000_000_000

步骤 4：启动预训练

在集群上进行大规模预训练：

# Using SLURM
sbatch --nodes=8 --gpus-per-node=8 \
  pretrain_script.sh

# pretrain_script.sh content:
litgpt pretrain \
  --config config/pythia-1b.yaml \
  --data.data_dir /shared/data/pretrain \
  --devices 8 \
  --num_nodes 8 \
  --train.global_batch_size 512 \
  --train.max_tokens 300_000_000_000

工作流 4：转换并部署模型

导出 LitGPT 模型以用于生产环境。

Model Deployment:
- [ ] Step 1: Test inference locally
- [ ] Step 2: Quantize model (optional)
- [ ] Step 3: Convert to GGUF (for llama.cpp)
- [ ] Step 4: Deploy with API

步骤 1：在本地测试推理

from litgpt import LLM

llm = LLM.load("out/phi2-lora/final")

# Single generation
print(llm.generate("What is machine learning?"))

# Streaming
for token in llm.generate("Explain quantum computing", stream=True):
    print(token, end="", flush=True)

# Batch inference
prompts = ["Hello", "Goodbye", "Thank you"]
results = [llm.generate(p) for p in prompts]

步骤 2：量化模型（可选）

在尽可能少损失质量的情况下缩小模型体积：

# 8-bit quantization (50% size reduction)
litgpt convert_lit_checkpoint \
  out/phi2-lora/final \
  --dtype bfloat16 \
  --quantize bnb.nf4

# 4-bit quantization (75% size reduction)
litgpt convert_lit_checkpoint \
  out/phi2-lora/final \
  --quantize bnb.nf4-dq  # Double quantization

步骤 3：转换为 GGUF（用于 llama.cpp）

python scripts/convert_lit_checkpoint.py \
  --checkpoint_path out/phi2-lora/final \
  --output_path models/phi2.gguf \
  --model_name microsoft/phi-2

步骤 4：通过 API 部署

from fastapi import FastAPI
from litgpt import LLM

app = FastAPI()
llm = LLM.load("out/phi2-lora/final")

@app.post("/generate")
def generate(prompt: str, max_tokens: int = 100):
    result = llm.generate(
        prompt,
        max_new_tokens=max_tokens,
        temperature=0.7
    )
    return {"response": result}

# Run: uvicorn api:app --host 0.0.0.0 --port 8000

何时使用 LitGPT，而非其他方案

适合使用 LitGPT 的情况：

• 想了解 LLM 架构（代码简洁易读）
• 需要可用于生产环境的训练方案
• 用于教育或研究
• 为新的模型创意制作原型
• Lightning 生态系统用户

适合改用其他方案的情况：

• Axolotl/TRL：微调功能更多，支持 YAML 配置
• Megatron-Core：针对 >70B 模型实现最高性能
• HuggingFace Transformers：模型支持范围最广
• vLLM：仅用于推理（不支持训练）

常见问题

问题：微调期间内存不足

改用 LoRA，而非全量微调：

# Instead of litgpt finetune (requires 40GB+)
litgpt finetune_lora  # Only needs 12-16GB

或者启用梯度检查点：

litgpt finetune_lora \
  ... \
  --train.gradient_accumulation_iters 4  # Accumulate gradients

问题：训练速度太慢

启用 Flash Attention（内置，在兼容硬件上自动启用）：

# Already enabled by default on Ampere+ GPUs (A100, RTX 30/40 series)
# No configuration needed

使用更小的微批次并进行梯度累积：

--train.micro_batch_size 1 \
--train.global_batch_size 32 \
--train.gradient_accumulation_iters 32  # Effective batch=32

问题：模型无法加载

检查模型名称：

# List all available models
litgpt download list

# Download if not exists
litgpt download meta-llama/Meta-Llama-3-8B

确认检查点目录：

ls checkpoints/
# Should see: meta-llama/Meta-Llama-3-8B/

问题：LoRA 适配器太大

降低 LoRA 秩：

--lora_r 8  # Instead of 16 or 32

减少应用 LoRA 的层数：

--lora_query true \
--lora_value true \
--lora_projection false \  # Disable this
--lora_mlp false  # And this

进阶主题

支持的架构：参见 references/supported-models.md，查看 20+ 种模型系列的完整列表，以及各自的大小和能力。

训练方案：参见 references/training-recipes.md，查看经过验证的预训练和微调超参数配置。

FSDP 配置：参见 references/distributed-training.md，了解如何使用 Fully Sharded Data Parallel 进行多 GPU 训练。

自定义架构：参见 references/custom-models.md，了解如何按照 LitGPT 风格实现新的模型架构。

硬件要求

• GPU：NVIDIA（CUDA 11.8+）、AMD（ROCm）、Apple Silicon（MPS）
• 内存：
  • 推理（Phi-2）：6GB
  • LoRA 微调（7B）：16GB
  • 全量微调（7B）：40GB+
  • 预训练（1B）：24GB
• 存储：每个模型 5-50GB（取决于大小）

资源

• GitHub：https://github.com/Lightning-AI/litgpt
• 文档：https://lightning.ai/docs/litgpt
• 教程：https://lightning.ai/docs/litgpt/tutorials
• 模型库：20+ 种预训练架构（Llama、Gemma、Phi、Qwen、Mistral、Mixtral、Falcon 等）

Agent 操作流程

1. 检查环境。 检查 GPU 类型、内存和驱动/CUDA 版本（nvidia-smi）、框架版本，以及数据集的位置和大小。
2. 明确输入。 根据数据或用户确认格式、标识符和参数。任何会改变结果的值都应先询问，不要猜测。
3. 先运行小规模版本。 进行冒烟测试：使用小型模型或子集，运行少量步骤，并确认损失下降且检查点已保存。
4. 按照上述说明和参考资料执行完整任务。
5. 验证结果。 在留出数据上跟踪指标，与基线比较，并记录随机种子、配置和硬件。
6. 报告。 说明运行了什么（版本、命令、参数）、检查了什么，以及还有哪些不确定之处。

如果发生这种情况	请执行以下操作
CUDA 显存不足	减小批次大小，启用梯度累积/检查点或混合精度，或对模型进行分片。
损失为 NaN 或发散	降低学习率，检查数据中是否有无效值，并启用梯度裁剪。
这些说明中提到的函数、标志或端点在已安装版本中不存在	查阅已安装版本自己的文档（help()、--help、官方文档），进行调整，并告知用户。切勿编造 API。
必需的输入、标识符或参数不明确	询问用户，或在运行前明确说明所作的假设。

诚信规则

• 切勿捏造结果、参数、标识符、引文或统计数据。如果某项内容无法运行或验证，请明确说明。
• 没有日志时，切勿声称训练取得了某种结果；启动大型任务前先估算计算成本。
• 此处与版本相关的细节可能已经过时：请对照已安装版本的官方文档进行确认。
• 对于需要花费资金、占用共享 GPU 或云配额、涉及个人或患者数据，或无法撤销的操作，请先征求同意。

相关技能

• unsloth：提供使用 Unsloth 微调大模型的指导。Unsloth 是一个利用 LoRA 和 QLoRA 实现更快、低内存训练的库，基础……
• peft-fine-tuning：使用 Hugging Face PEFT 微调 LLM，采用 LoRA、QLoRA、IA3、AdaLoRA、前缀调优和提示词调优，使得不到 1% 的参数参与……