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 を使用するケース:

• 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

checkpoints ディレクトリを確認します:

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

高度なトピック

対応アーキテクチャ: 20 種類以上のモデルファミリーのサイズと機能の完全な一覧は、references/supported-models.md を参照してください。

学習レシピ: 事前学習とファインチューニングのための実績のあるハイパーパラメータ設定は、references/training-recipes.md を参照してください。

FSDP の設定: Fully Sharded Data Parallel を使用したマルチ GPU 学習は、references/distributed-training.md を参照してください。

カスタムアーキテクチャ: LitGPT のスタイルで新しいモデルアーキテクチャを実装する方法は、references/custom-models.md を参照してください。

ハードウェア要件

• 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など）