Know before you train — VRAM estimation for LLM fine-tuning.
fitcheck predicts GPU memory usage from first principles. Given a model, GPU, and training method, it tells you whether your config will fit — before you spend an hour discovering it won't.
Fine-tuning LLMs means guessing at batch sizes and hoping you don't OOM. The feedback loop is brutal: pick a config, wait for the run to start, crash 2 minutes in, adjust, repeat.
fitcheck collapses that loop. It computes each VRAM component — model weights, optimizer states, gradients, activations, the logits buffer, eval KV-cache spikes — and produces a breakdown with confidence bounds.
| Component | What it is | Why it matters |
|---|---|---|
| Model weights | Base params in training dtype (bf16/NF4) | 4.2 GB for QLoRA 8B, 16 GB for full bf16 |
| Optimizer states | AdamW momentum + variance per trainable param | Dominates full fine-tune (~60 GB for 8B) |
| Gradients | One gradient per trainable param | Small for LoRA, huge for full FT |
| Activations | Per-layer stored tensors for backward pass | Flash-attention-aware, scales with batch × seq |
| Logits buffer | batch × seq × vocab × 4 bytes (float32) | The surprise OOM — 2 GB at bs=4 with 128k vocab |
| Eval KV-cache | Spike during evaluation steps | Can exceed training steady-state |
pip install -r requirements.txt
# Run tests
pytestfrom fitcheck.hub.resolver import resolve_from_config
from fitcheck.hardware.registry import get_hardware
from fitcheck.profilers.vram.engine import VRAMEstimator
from fitcheck.models.profiles import TrainingMethod, LoRAConfig
# QLoRA Llama 8B on an RTX 3090
estimator = VRAMEstimator()
breakdown = estimator.estimate(
model=resolve_from_config("meta-llama/Llama-3.1-8B", config),
hardware=get_hardware("3090"),
method=TrainingMethod.QLORA,
batch_size=4,
seq_len=1024,
lora_config=LoRAConfig(rank=16),
)
print(f"Steady-state: {breakdown.steady_state_gb:.1f} GB")
print(f"Usable VRAM: {get_hardware('3090').usable_vram_gb} GB")
# Steady-state: 16.6 GB
# Usable VRAM: 22.8 GB ← fits with 6 GB headroompytest # Run all tests (51)
pytest tests/fitcheck/profilers/test_estimator.py -v # End-to-end estimator tests
ruff format fitcheck tests # Format codeMIT