# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file. """Tests for the validate CLI script and related error-handling utilities.""" import json import warnings from dataclasses import asdict from pathlib import Path from unittest import mock import pytest import torch import yaml from litgpt import GPT from litgpt.config import Config from litgpt.utils import ( CheckpointValidationResult, estimate_model_memory, validate_checkpoint, ) # --------------------------------------------------------------------------- # validate_checkpoint tests # --------------------------------------------------------------------------- class TestValidateCheckpoint: """Tests for the validate_checkpoint utility.""" @staticmethod def _save_model_checkpoint(model: torch.nn.Module, path: Path) -> None: torch.save(model.state_dict(), str(path)) def test_valid_checkpoint(self, tmp_path): """A checkpoint saved from the same model should pass validation.""" config = Config.from_name("pythia-14m") with torch.device("meta"): model = GPT(config) # Create a real state_dict with matching shapes real_sd = {k: torch.randn(v.shape) for k, v in model.state_dict().items()} ckpt_path = tmp_path / "lit_model.pth" torch.save(real_sd, str(ckpt_path)) result = validate_checkpoint(ckpt_path, model, verbose=False) assert result.is_valid assert result.missing_keys == [] assert result.unexpected_keys == [] assert result.shape_mismatches == [] assert result.errors == [] assert "passed" in result.summary().lower() def test_missing_keys(self, tmp_path): """Checkpoint missing some keys should report them.""" config = Config.from_name("pythia-14m") with torch.device("meta"): model = GPT(config) real_sd = {k: torch.randn(v.shape) for k, v in model.state_dict().items()} # Remove a key removed_key = list(real_sd.keys())[0] del real_sd[removed_key] ckpt_path = tmp_path / "lit_model.pth" torch.save(real_sd, str(ckpt_path)) result = validate_checkpoint(ckpt_path, model, verbose=False) assert not result.is_valid assert removed_key in result.missing_keys assert result.unexpected_keys == [] def test_unexpected_keys(self, tmp_path): """Checkpoint with extra keys should report them.""" config = Config.from_name("pythia-14m") with torch.device("meta"): model = GPT(config) real_sd = {k: torch.randn(v.shape) for k, v in model.state_dict().items()} real_sd["extra.unexpected.key"] = torch.randn(3) ckpt_path = tmp_path / "lit_model.pth" torch.save(real_sd, str(ckpt_path)) result = validate_checkpoint(ckpt_path, model, verbose=False) assert not result.is_valid assert "extra.unexpected.key" in result.unexpected_keys def test_shape_mismatch(self, tmp_path): """Checkpoint with wrong shapes should report mismatches.""" config = Config.from_name("pythia-14m") with torch.device("meta"): model = GPT(config) real_sd = {k: torch.randn(v.shape) for k, v in model.state_dict().items()} # Corrupt a shape key = "lm_head.weight" real_sd[key] = torch.randn(10, 10) # wrong shape ckpt_path = tmp_path / "lit_model.pth" torch.save(real_sd, str(ckpt_path)) result = validate_checkpoint(ckpt_path, model, verbose=False) assert not result.is_valid assert any(key in m for m in result.shape_mismatches) def test_file_not_found(self, tmp_path): """Non-existent checkpoint should report an error.""" config = Config.from_name("pythia-14m") with torch.device("meta"): model = GPT(config) result = validate_checkpoint(tmp_path / "nonexistent.pth", model, verbose=False) assert not result.is_valid assert any("not found" in e for e in result.errors) def test_corrupted_file(self, tmp_path): """A file that is not a valid PyTorch checkpoint should report an error.""" config = Config.from_name("pythia-14m") with torch.device("meta"): model = GPT(config) ckpt_path = tmp_path / "corrupted.pth" ckpt_path.write_text("this is not a checkpoint") result = validate_checkpoint(ckpt_path, model, verbose=False) assert not result.is_valid assert any("Failed to load" in e for e in result.errors) def test_model_key_wrapper(self, tmp_path): """Checkpoint wrapped under a 'model' key should be unwrapped.""" config = Config.from_name("pythia-14m") with torch.device("meta"): model = GPT(config) real_sd = {k: torch.randn(v.shape) for k, v in model.state_dict().items()} wrapped = {"model": real_sd} ckpt_path = tmp_path / "lit_model.pth" torch.save(wrapped, str(ckpt_path)) result = validate_checkpoint(ckpt_path, model, verbose=False) assert result.is_valid def test_summary_format(self): """Summary strings should be well-formed.""" result = CheckpointValidationResult( is_valid=True, missing_keys=[], unexpected_keys=[], shape_mismatches=[], errors=[] ) assert result.summary() == "Checkpoint validation passed." result = CheckpointValidationResult( is_valid=False, missing_keys=["a", "b"], unexpected_keys=["c"], shape_mismatches=["d: model=(2,3), checkpoint=(4,5)"], errors=[], ) summary = result.summary() assert "Missing keys" in summary assert "Unexpected keys" in summary assert "Shape mismatches" in summary # --------------------------------------------------------------------------- # estimate_model_memory tests # --------------------------------------------------------------------------- class TestEstimateModelMemory: """Tests for the estimate_model_memory utility.""" def test_basic_estimation(self): """Should return reasonable estimates for a known config.""" config = Config.from_name("pythia-14m") result = estimate_model_memory(config, dtype=torch.float32, training=False) assert result["param_count"] > 0 assert result["param_memory_gb"] > 0 assert result["estimated_total_gb"] > 0 # pythia-14m is ~14M params → ~0.05 GB in fp32 assert result["param_memory_gb"] < 1.0 def test_training_multiplier(self): """Training should use ~4× multiplier (params + gradients + Adam optimizer states).""" config = Config.from_name("pythia-14m") inference = estimate_model_memory(config, dtype=torch.float32, training=False) training = estimate_model_memory(config, dtype=torch.float32, training=True) assert training["estimated_total_gb"] > inference["estimated_total_gb"] # Should be approximately 4× (params + gradients + Adam optimizer states). # Bounds are loose (3.5–4.5) to absorb rounding from the two round() calls in the function. ratio = training["estimated_total_gb"] / inference["estimated_total_gb"] assert 3.5 < ratio < 4.5 def test_dtype_affects_memory(self): """Half precision should use ~half the parameter memory.""" config = Config.from_name("pythia-14m") fp32 = estimate_model_memory(config, dtype=torch.float32, training=False) fp16 = estimate_model_memory(config, dtype=torch.float16, training=False) assert fp16["param_memory_gb"] < fp32["param_memory_gb"] # Should be approximately double (exact ratio depends on estimate granularity) ratio = fp32["param_memory_gb"] / fp16["param_memory_gb"] assert 1.5 < ratio < 2.5 def test_gpu_fields(self): """GPU-related fields should be None when no GPU is available.""" config = Config.from_name("pythia-14m") with mock.patch("litgpt.utils.torch.cuda.is_available", return_value=False): result = estimate_model_memory(config, dtype=torch.float32) assert result["available_gpu_memory_gb"] is None assert result["fits_in_memory"] is None # --------------------------------------------------------------------------- # Tokenizer JSON warning test # --------------------------------------------------------------------------- def test_tokenizer_json_warning(tmp_path): """Tokenizer should emit a warning when generation_config.json has invalid JSON.""" # Set up a minimal tokenizer directory with an HF tokenizer checkpoint_dir = tmp_path / "test_model" checkpoint_dir.mkdir() # We need tokenizer.json for the HF path and a valid tokenizer_config.json # Use a minimal invalid generation_config.json invalid_json = '{\n "bos_token_id": 1,\n "eos_token_id": 2,\n}' # trailing comma (checkpoint_dir / "generation_config.json").write_text(invalid_json) (checkpoint_dir / "tokenizer_config.json").write_text(json.dumps({"tokenizer_class": "GPT2Tokenizer"})) # Create a minimal tokenizer.json that the HF tokenizer can load minimal_tokenizer_json = { "version": "1.0", "truncation": None, "padding": None, "added_tokens": [], "normalizer": None, "pre_tokenizer": None, "post_processor": None, "decoder": None, "model": {"type": "BPE", "vocab": {"": 0, "": 1}, "merges": []}, } (checkpoint_dir / "tokenizer.json").write_text(json.dumps(minimal_tokenizer_json)) from litgpt.tokenizer import Tokenizer with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") tokenizer = Tokenizer(checkpoint_dir) # Check that the warning was raised assert len(w) == 1 assert "invalid JSON" in str(w[0].message) # Check that the fix worked assert tokenizer.bos_id == 1 assert tokenizer.eos_id == 2 # --------------------------------------------------------------------------- # Validate script integration test # --------------------------------------------------------------------------- class TestValidateScript: """Integration tests for the validate CLI script.""" @staticmethod def _make_checkpoint_dir(tmp_path: Path, config_name: str = "pythia-14m") -> Path: """Create a fake but structurally valid checkpoint directory.""" checkpoint_dir = tmp_path / "checkpoints" / "test" checkpoint_dir.mkdir(parents=True) config = Config.from_name(config_name) config_dict = asdict(config) with open(checkpoint_dir / "model_config.yaml", "w") as f: yaml.dump(config_dict, f) # Create a real checkpoint with torch.device("meta"): model = GPT(config) real_sd = {k: torch.randn(v.shape) for k, v in model.state_dict().items()} torch.save(real_sd, str(checkpoint_dir / "lit_model.pth")) # Create minimal tokenizer files (checkpoint_dir / "tokenizer_config.json").write_text(json.dumps({"tokenizer_class": "GPT2Tokenizer"})) (checkpoint_dir / "tokenizer.json").write_text("{}") return checkpoint_dir def test_validate_missing_dir(self, tmp_path, capsys): """validate_setup should fail for a non-existent directory.""" from litgpt.scripts.validate import validate_setup with pytest.raises(SystemExit): validate_setup(checkpoint_dir=tmp_path / "nonexistent") def test_validate_missing_model_file(self, tmp_path, capsys): """validate_setup should fail when checkpoint file is missing.""" checkpoint_dir = tmp_path / "test" checkpoint_dir.mkdir() (checkpoint_dir / "model_config.yaml").write_text(yaml.dump(asdict(Config.from_name("pythia-14m")))) (checkpoint_dir / "tokenizer_config.json").write_text(json.dumps({"tokenizer_class": "GPT2Tokenizer"})) (checkpoint_dir / "tokenizer.json").write_text("{}") # No lit_model.pth from litgpt.scripts.validate import validate_setup with pytest.raises(SystemExit): validate_setup(checkpoint_dir=checkpoint_dir)