Metadata-Version: 2.4
Name: vram-check
Version: 0.1.1
Summary: Estimate GPU VRAM requirements for PyTorch models before running them
License: MIT
Project-URL: Homepage, https://github.com/harshavardhanreddyseethagari-sjsu2103/vram-check
Project-URL: Issues, https://github.com/harshavardhanreddyseethagari-sjsu2103/vram-check/issues
Keywords: pytorch,gpu,vram,memory,machine-learning,deep-learning,profiling
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: click>=8.0
Provides-Extra: torch
Requires-Dist: torch>=2.0; extra == "torch"

# 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.

```bash
$ 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

```bash
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

```bash
# 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

```bash
# 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

```bash
vram-check gpus
```

### Python API

```python
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
