vram-check 0.1.0

Estimate GPU VRAM requirements for PyTorch models before running them

vram-check

Estimate GPU VRAM requirements for PyTorch models — before you run them.

To avoid wasting GPU hours on jobs that OOM 2 hours in, vram-check statically analyzes your model architecture and computes memory requirements for parameters, gradients, optimizer state, and activations without executing any training code.

$ vram-check inspect torchvision.models:resnet50 --batch-size 32 --dtype fp16 --gpu A100-40GB

  vram-check  torchvision.models:resnet50
  Training · FP16 · batch=32 · optimizer=adam

  Parameters:        25,557,032
  Trainable:         25,557,032

  Parameters:       0.048 GB
  Gradients:        0.048 GB
  Optimizer:        0.190 GB
  Activations:      3.046 GB  (estimated)
  ──────────────────────────────────────
  Total:            3.331 GB

  ✓ NVIDIA A100 40GB (40.0 GB): fits with margin

  Top 5 layers by parameter count:
  1. layer4.0.conv2 (Conv2d): 2,359,296 params  (9.0 MB)
  2. layer4.1.conv2 (Conv2d): 2,359,296 params  (9.0 MB)
  3. layer4.2.conv2 (Conv2d): 2,359,296 params  (9.0 MB)
  4. layer4.0.downsample.0 (Conv2d): 2,097,152 params  (8.0 MB)
  5. fc (Linear): 2,049,000 params  (7.8 MB)

Why is this needed?

PyTorch Profiler tells you memory usage after a run — which means you already crashed (or wasted allocation time on a shared cluster) to find out. vram-check answers the question before you submit the job.

This is especially useful on shared HPC clusters (Slurm, PBS) where a failed job wastes allocation that could have gone to someone else.

Installation

pip install vram-check

PyTorch is optional — required only for inspect (passing a real model object). The quick command works without it.

Usage

Inspect a real PyTorch model

# By import path (module:ClassName or module:factory_function)
vram-check inspect torchvision.models:resnet50
vram-check inspect torchvision.models:vgg16 --batch-size 32 --dtype fp16
vram-check inspect myproject.models:MyTransformer --gpu A100-40GB --inference

# Common options
--batch-size, -b    Batch size (default: 1)
--dtype, -d         fp32 | fp16 | bf16 | int8 (default: fp32)
--optimizer, -o     adam | adamw | sgd | none (default: adam)
--inference         Estimate for inference only (no gradients/optimizer)
--gpu, -g           Check against a specific GPU (e.g. A100-40GB, RTX3090)
--json              Output as JSON for scripting

Quick estimate by parameter count

# No model file needed — useful for back-of-envelope calculations
vram-check quick --params 7000000000 --dtype fp16 --gpu A100-40GB
#  LLaMA-7B in fp16, Adam training: ~104 GB total — does not fit on A100-40GB

vram-check quick --params 7000000000 --dtype fp16 --optimizer none --inference --gpu A100-80GB
#  LLaMA-7B fp16 inference (no optimizer): fits on A100-80GB

List known GPUs

vram-check gpus

Python API

from vram_check import analyze, estimate

# Full analysis from a model object
import torchvision.models as models
model = models.resnet50()

report = analyze(model, batch_size=16, dtype="fp16", optimizer="adam")
print(report["estimate"].total_gb)       # 1.809
print(report["estimate"].breakdown_gb)   # {'parameters': 0.048, ...}
print(report["estimate"].fits_on(24.0))  # True

# Raw estimate from a parameter count
result = estimate(num_params=117_000_000, batch_size=1, dtype="fp32")
print(result)

How the math works

VRAM usage for a training run has four components:

Component	Formula	Notes
Parameters	num_params × bytes_per_dtype	fp32=4B, fp16/bf16=2B, int8=1B
Gradients	Same as parameters	Only during training; zero for inference
Optimizer state	num_params × 4B × multiplier	Adam=2× (m+v), SGD=1× (momentum), none=0
Activations	Heuristic: ~2× params × batch_size	Approximate — hardest to estimate statically

Important: optimizer state is always stored in fp32 regardless of model dtype, for numerical stability. This is why training a fp16 model uses significantly more memory than its parameter size alone suggests.

Activations are approximate. The tool overestimates slightly on purpose — better to see "3.2 GB" and get 2.8 GB than to see "2.5 GB" and OOM. For convolutional models especially, real activation memory may be notably lower than our heuristic.

Supported GPUs

Run vram-check gpus to see the full list. Includes H100, A100, V100, RTX 4090/3090/3080, T4, P100, MI250X, and more.

Limitations

• Activations are estimated, not exact. Static analysis can't perfectly predict activation memory without knowing the full computation graph and input shapes. Treat the activation number as a conservative upper bound.
• PyTorch only for now. HuggingFace Transformers config support is planned for v0.2.0.
• No custom CUDA kernels or quantization schemes. Estimates assume standard PyTorch memory layout.

Contributing

Bug reports and PRs welcome. If you find a model where the estimate is significantly wrong, please open an issue with the model architecture and actual measured VRAM usage.

License

MIT