File size: 5,863 Bytes
236083b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
"""Pre-flight validation for LitGPT checkpoints.
Usage:
litgpt validate --checkpoint_dir checkpoints/meta-llama/...
"""
import sys
from pathlib import Path
import torch
import yaml
from litgpt.config import Config
from litgpt.utils import (
check_valid_checkpoint_dir,
estimate_model_memory,
validate_checkpoint,
)
def validate_setup(
checkpoint_dir: Path,
model_filename: str = "lit_model.pth",
dtype: str = "float32",
training: bool = False,
) -> None:
"""Run pre-flight validation on a checkpoint directory.
This checks everything without actually running training or generation:
1. Checkpoint directory structure (required files exist)
2. Model config loading
3. Tokenizer loading
4. Checkpoint key/shape validation against the model
5. Memory estimation
Args:
checkpoint_dir: Path to the checkpoint directory.
model_filename: Name of the checkpoint file (default: ``lit_model.pth``).
dtype: Data type for memory estimation (``float32``, ``float16``, ``bfloat16``).
training: If ``True``, estimate memory for training (includes optimizer states).
"""
checkpoint_dir = Path(checkpoint_dir)
print(f"{'=' * 60}")
print("LitGPT Pre-flight Validation")
print(f"Checkpoint: {checkpoint_dir}")
print(f"{'=' * 60}\n")
all_passed = True
# --- Step 1: Checkpoint directory structure ---
print("[1/5] Checking checkpoint directory structure...")
try:
check_valid_checkpoint_dir(
checkpoint_dir,
model_filename=model_filename,
)
print(" ✓ All required files found.\n")
except (FileNotFoundError, SystemExit) as e:
print(f" ✗ Directory validation failed: {e}\n", file=sys.stderr)
all_passed = False
# --- Step 2: Load model config ---
print("[2/5] Loading model config...")
config = None
config_path = checkpoint_dir / "model_config.yaml"
try:
if config_path.is_file():
with open(config_path, encoding="utf-8") as f:
config_dict = yaml.safe_load(f)
config = Config(**config_dict)
print(f" ✓ Config loaded: {config.name or 'unnamed'}")
print(
f" n_layer={config.n_layer}, n_embd={config.n_embd}, "
f"n_head={config.n_head}, vocab_size={config.vocab_size}\n"
)
else:
print(f" ✗ Config file not found: {config_path}\n", file=sys.stderr)
all_passed = False
except Exception as e:
print(f" ✗ Failed to load config: {e}\n", file=sys.stderr)
all_passed = False
# --- Step 3: Tokenizer ---
print("[3/5] Checking tokenizer...")
try:
from litgpt.tokenizer import Tokenizer
tokenizer = Tokenizer(checkpoint_dir)
# Do a simple encode/decode round-trip
test_text = "Hello"
tokens = tokenizer.encode(test_text)
decoded = tokenizer.decode(tokens)
print(f" ✓ Tokenizer loaded (backend={tokenizer.backend})")
print(f' Round-trip test: "{test_text}" → {tokens.tolist()} → "{decoded}"\n')
except Exception as e:
print(f" ✗ Tokenizer failed: {e}\n", file=sys.stderr)
all_passed = False
# --- Step 4: Checkpoint validation ---
print("[4/5] Validating checkpoint against model...")
checkpoint_path = checkpoint_dir / model_filename
if config is not None and checkpoint_path.is_file():
try:
from litgpt import GPT
with torch.device("meta"):
model = GPT(config)
result = validate_checkpoint(checkpoint_path, model, verbose=False)
if result.is_valid:
print(" ✓ Checkpoint keys and shapes match the model.\n")
else:
all_passed = False
print(f" ✗ {result.summary()}\n", file=sys.stderr)
except Exception as e:
print(f" ✗ Checkpoint validation error: {e}\n", file=sys.stderr)
all_passed = False
elif not checkpoint_path.is_file():
print(f" ⊘ Skipped (checkpoint file not found: {checkpoint_path})\n")
else:
print(" ⊘ Skipped (config not loaded)\n")
# --- Step 5: Memory estimation ---
print("[5/5] Estimating memory requirements...")
if config is not None:
mem = estimate_model_memory(config, dtype=dtype, training=training)
print(f" Estimated parameters: {mem['param_count']:,}")
print(f" Parameter memory: {mem['param_memory_gb']:.2f} GB ({dtype})")
mode_str = "training (params + grads + optimizer)" if training else "inference (params only)"
print(f" Estimated total ({mode_str}): {mem['estimated_total_gb']:.2f} GB")
if mem["available_gpu_memory_gb"] is not None:
print(f" Available GPU memory: {mem['available_gpu_memory_gb']:.2f} GB")
if mem["fits_in_memory"]:
print(" ✓ Model should fit in GPU memory.\n")
else:
print(" ⚠ WARNING: Model may NOT fit in GPU memory!\n", file=sys.stderr)
all_passed = False
else:
print(" ⊘ No GPU detected, skipping memory fit check.\n")
else:
print(" ⊘ Skipped (config not loaded)\n")
# --- Summary ---
print(f"{'=' * 60}")
if all_passed:
print("✓ All validation checks passed!")
else:
print("✗ Some validation checks failed. See details above.", file=sys.stderr)
print(f"{'=' * 60}")
if not all_passed:
raise SystemExit(1)
|