Checkpoint Health Checker

A single-file tool for inspecting .pt checkpoints from tiny language model training runs. Tells you if your model is healthy, degenerate, or broken — before you waste time uploading or evaluating it.

What it checks

Check	What it looks for
Architecture	vocab size, d_model, layers, heads, tied embeddings (inferred from state_dict keys + config)
Parameter count	Total parameters (sum of all tensor numel)
Weight health	NaN, Inf, extreme values (>100 abs) in any weight tensor
Generation sanity	EOS collapse, degenerate loops, low diversity, too-short output

Requirements

Python 3.8+
PyTorch (pip install torch)
A tokenizer (HuggingFace tokenizers library) for the generation check — optional, the weight checks work without it

Usage

# Basic: weight health only (no tokenizer needed)
python3 checkpoint_health.py model.pt

# With generation check
python3 checkpoint_health.py model.pt --tokenizer tokenizer.json --prompt "Once upon a time"

# JSON output (for scripting)
python3 checkpoint_health.py model.pt --tokenizer tokenizer.json --json

Exit codes

Code	Meaning
0	Healthy — weights clean, generation looks reasonable
1	Warning — degenerate generation or extreme weight values
2	Critical — NaN/Inf weights detected (model is broken)

Example output

Healthy model (12.6M params, TinyStories-trained)

=== Checkpoint Health Report ===
File: story10m/final.pt  (20.6 MB)

[ARCHITECTURE]
  vocab_size:  256
  d_model:     256
  n_layers:    6
  n_heads:     4
  n_params:    12,603,648
  tied_embeddings: yes

[WEIGHT HEALTH]
  NaN tensors:  0
  Inf tensors:  0
  Max |val|:    3.214 (wte.weight)
  Status: CLEAN

[GENERATION]
  Prompt: "Once upon a time"
  Output: "Once upon a time, there was a little girl who lived in a small house..."
  New tokens: 42
  EOS hits: 0
  Trigram repetition: 4.2%
  Char diversity: 0.412
  Status: OK

VERDICT: HEALTHY

Degenerate model (1.6M params, undertrained)

=== Checkpoint Health Report ===
File: jokeclaude/final.pt  (6.3 MB)

[ARCHITECTURE]
  vocab_size:  50257
  d_model:     128
  n_layers:    2
  n_params:    1,638,400
  tied_embeddings: yes

[WEIGHT HEALTH]
  NaN tensors:  0
  Inf tensors:  0
  Max |val|:    1.847 (wte.weight)
  Status: CLEAN

[GENERATION]
  Prompt: "The cat sat on the mat and"
  Output: "The cat sat on the mat and the the the the the the the the the..."
  New tokens: 100
  EOS hits: 0
  Trigram repetition: 78.3%
  Char diversity: 0.089
  Status: LOW_DIVERSITY

VERDICT: WARNING — Low character diversity (0.089). Model may be stuck on a character.

Limitations

RoPE/GQA architectures: The generation check requires instantiating a model class. If the architecture uses RoPE or grouped query attention (common in 2024+ tiny LMs), the generation check is skipped with an explanation. Weight health and architecture inference still work.
Tokenizer required for generation: Without --tokenizer, only weight health is reported.
Single-file only: Expects a .pt file with either a raw state_dict or a dict with model/model_state_dict key.

When to use it

Right after a training run completes, before uploading
When a model "works" in your training loop but generates garbage at inference
To quickly triage a checkpoint you found on the Hub
To confirm a training run didn't produce NaN weights (silent failure mode)

Downloads last month: -; Downloads are not tracked for this model. How to track