vramwatch 0.1.0

GPU memory profiler for PyTorch training on shared HPC environments. Record during training, visualize after.

vramwatch

GPU memory profiler for PyTorch training on shared HPC environments.

Record during training. Visualize after. No live servers during training.

vramwatch/
  record on HPC node (zero network access) ──► .vram file
                                                    │
                          ┌─────────────────────────┤
                          ▼                         ▼
               HTML report (offline)      Local web server
               Jupyter inline charts     (with live tail)

---

Installation

pip install vramwatch

# With all optional extras:
pip install "vramwatch[all]"

---

Quick Start

1. Record during training

from vramwatch import VRAMWatch

watch = VRAMWatch("run.vram", device=0, flush_every=10)

for step, batch in enumerate(dataloader):
    with watch.forward():
        output = model(batch)
        loss = criterion(output, batch["labels"])

    with watch.backward():
        loss.backward()

    with watch.optimizer():
        optimizer.step()
        optimizer.zero_grad()

    watch.step(step=step, loss=loss.detach().item(), epoch=epoch)

watch.save()

As a context manager:

with VRAMWatch("run.vram", device=0) as watch:
    for step, batch in enumerate(dataloader):
        with watch.forward():
            ...
        watch.step(step=step, loss=loss.item())

The .vram file is flushed every flush_every steps — if your job is killed, data up to the last flush is preserved.

---

2. Visualize

Option A — Standalone HTML report (offline-capable)

vramwatch analyze run.vram
# → run.vram.html (self-contained, works offline after generation)

# With AI recommendations:
vramwatch analyze run.vram --api-key sk-ant-...

Option B — Local web server (with live tail)

vramwatch serve run.vram
# Opens http://localhost:7842 in your browser
# Live-tails the .vram file if training is still running

vramwatch serve run.vram --port 8080 --api-key sk-ant-...

Option C — Jupyter inline

from vramwatch import analyze
analyze("run.vram")

# With AI analysis:
analyze("run.vram", anthropic_api_key="sk-ant-...")

Quick terminal stats (useful on HPC login nodes)

vramwatch info run.vram

---

CLI Reference

vramwatch analyze <path.vram> [--output path.html] [--api-key sk-...]
  Generate a self-contained HTML report. Prints summary to terminal.

vramwatch serve <path.vram> [--port 7842] [--no-browser] [--api-key sk-...]
  Start local web server. Live-tails file if still being written.

vramwatch info <path.vram>
  Print metadata + summary stats as a table. No charts.
  Designed for HPC login nodes.

vramwatch generate-sample [--steps 200] [--output sample.vram] [--seed 42]
  Generate a realistic fake .vram file for testing/demo.

---

.vram File Format

JSONL (one JSON object per line). The first line is a metadata header:

{"__meta__": true, "gpu_name": "NVIDIA A100-SXM4-40GB", "total_vram_mb": 40960, "cuda_version": "12.1", "torch_version": "2.3.0", "hostname": "hpc-node-42", "created_at": "2026-05-03T14:32:00Z", "vramwatch_version": "0.1.0"}

Each subsequent line is an event:

{"step": 42, "epoch": 1, "phase": "backward", "loss": 0.3421, "ts": 1746123456.789, "allocated_mb": 18432, "reserved_mb": 20480, "peak_mb": 21504, "params_mb": 4512, "gradients_mb": 4512, "activations_mb": 5120, "optimizer_mb": 9024, "free_mb": 18944, "total_vram_mb": 40960}

Phases: forward | forward_start | backward | backward_start | optimizer | optimizer_start | step

The format is forward-compatible: unknown keys are preserved and ignored.

---

Python API

from vramwatch.reader import VRAMFile
from vramwatch.analyzer import compute_stats, build_recommendations

# Load and parse
vf = VRAMFile.load("run.vram")
print(vf.gpu_name)           # "NVIDIA A100-SXM4-40GB"
print(vf.total_vram_mb)      # 40960.0
print(len(vf.events))        # number of recorded events

# Filter by phase
fwd_events = vf.filter_phase("forward")
step_events = vf.step_events()

# Pandas DataFrame (requires pandas)
df = vf.to_df()

# Stats and recommendations
stats = compute_stats(vf)
recs  = build_recommendations(stats)

# HTML report
from vramwatch.exporters.html import generate_html
generate_html(vf, "report.html", anthropic_api_key="sk-ant-...")

# Local server
from vramwatch.exporters.server import serve
serve(vf, port=7842, anthropic_api_key="sk-ant-...")

---

What it measures

Field	Source
allocated_mb	torch.cuda.memory_allocated()
reserved_mb	torch.cuda.memory_reserved()
peak_mb	torch.cuda.max_memory_allocated()
params_mb	memory_stats()["active_bytes.all.current"]
gradients_mb	Reserved − active (backward phase delta)
activations_mb	Forward peak − params baseline
optimizer_mb	Allocated during optimizer phase − params
free_mb	total_vram − reserved

If torch.cuda.memory_stats() is unavailable (PyTorch < 1.8), fine-grained breakdown fields will be null with a one-time warning. Core metrics always work.

---

Rule-based Recommendations

vramwatch generates recommendations without any AI required:

Condition	Recommendation
Efficiency < 60%	Increase batch size (with safe multiplier estimate)
Efficiency > 90%	Enable gradient checkpointing or reduce batch size
Monotonic memory increase	Memory leak — check .detach(), empty_cache()
Activations > Params	Enable gradient checkpointing
Optimizer > 2× Params	Switch to bf16 or 8-bit Adam (bitsandbytes)
Headroom > 8 GB	Reduce gradient accumulation steps

---

HPC Design

• recorder.py requires zero network access — suitable for isolated compute nodes
• Crashes/OOM: data up to last auto-flush is preserved (configurable interval)
• No torch import required for viewer — install on machines without CUDA
• Graceful degradation if torch.cuda is unavailable (writes null for memory fields)

---

Testing

pip install "vramwatch[dev]"
pytest tests/ -v

Generate the sample fixture:

vramwatch generate-sample --steps 200 --output tests/fixtures/sample.vram

---

License

MIT