torch-kitsune 0.4.0

Never OOM again. Predict peak GPU memory, max batch size, and cheapest GPU for any PyTorch model — before you run.

Kitsune — GPU Planning for PyTorch

Never OOM again. Know your GPU costs before you run.

Kitsune predicts peak GPU memory, max batch size, and cheapest GPU for any PyTorch model — before you write a single torch.cuda.OutOfMemoryError.

pip install torch-kitsune[predict]

---

Why Kitsune?

Every ML engineer has been here:

→ Start training  →  CUDA OOM  →  Reduce batch size  →  Repeat for 2 hours

Kitsune breaks the loop. It statically analyzes your model and tells you:

• Will it fit? On which GPU, at which batch size
• What's the cheapest GPU? Full rankings by $/hr and $/day
• How to reduce memory? Ranked recommendations (FP16, grad checkpointing, quantization)

No GPU required to run predictions. Works from your laptop.

---

Quick Start

Predict memory for a local model

from kitsune.predict import predict_memory
import torch.nn as nn

model = nn.Sequential(nn.Linear(784, 2048), nn.ReLU(), nn.Linear(2048, 10))

result = predict_memory(model, input_shape=(784,), batch_size=32, gpu="t4")
print(result.summary())
# Peak: 0.14 GB | Tesla T4 (16GB): FITS

HuggingFace models (no weights downloaded)

from kitsune.predict import predict_hf_memory

result = predict_hf_memory("bert-base-uncased", batch_size=32, training=True, gpu="t4")
print(result.summary())
# Peak: 1.89 GB | Tesla T4 (16GB): FITS

Find cheapest GPU

from kitsune.predict import plan_gpu

plan = plan_gpu(model, input_shape=(784,), training=True)
print(plan.cheapest.gpu.name, f"${plan.cheapest.cost_per_hr:.2f}/hr")
# NVIDIA RTX 3080  $0.30/hr

---

CLI

# Single prediction
kitsune predict --hf bert-base-uncased --batch-size 32 --gpu t4

# Sweep batch sizes to find the OOM boundary
kitsune predict --hf bert-large-uncased --sweep --gpu a100-40gb

# GPU cost planner — rank all GPUs by price
kitsune plan --hf bert-base-uncased --training
kitsune plan --hf gpt2 --training --daily-budget 50

# Local model file
kitsune predict model.py:ResNet50 --input-shape 3,224,224 --gpu t4
kitsune plan model.py:ResNet50 --input-shape 3,224,224 --training

# List all GPUs with pricing
kitsune list-gpus

kitsune plan --hf bert-large-uncased --training

╭──────────────────────── Kitsune GPU Plan ────────────────────────────╮
│ GPU                    VRAM    $/hr   $/day  Max Batch  GPU%  Status │
│ NVIDIA RTX 3080        10GB   $0.30     $7      3397    57%  ★ BEST  │
│ NVIDIA RTX 4070        12GB   $0.35     $8      4096    48%   FITS   │
│ Tesla T4               16GB   $0.35     $8      4096    36%   FITS   │
│ NVIDIA RTX 3090        24GB   $0.44    $11      4096    24%   FITS   │
│ NVIDIA A100 80GB       80GB   $1.99    $48      4096     7%   FITS   │
│ NVIDIA H100 SXM        80GB   $3.29    $79      4096     7%   FITS   │
╰──── BertModel | float32 | Training | Needs 5.74 GB ─────────────────╯

  Cheapest: NVIDIA RTX 3080 ($0.30/hr, max batch 3397)
  Prices: lowest on-demand rate (Lambda/RunPod/GCP/AWS, 2025-05).

kitsune predict --hf bert-large-uncased --sweep --gpu a100-40gb

╭────────────────────── Kitsune Batch Size Sweep ──────────────────────╮
│ Batch  Params   Grads    Opt     Act     Total    GPU%   Fit?        │
│     1   1.24    1.24    2.49   0.000   5.71 GB   14.3%    ✓         │
│    16   1.24    1.24    2.49   0.024   5.74 GB   14.3%    ✓         │
│    64   1.24    1.24    2.49   0.097   5.82 GB   14.5%    ✓         │
│   256   1.24    1.24    2.49   0.388   6.08 GB   15.2%    ✓         │
│   512   1.24    1.24    2.49   0.776   6.46 GB   16.2%    ✓         │
╰──── BertModel | float32 | Training | Target: A100 40GB (40GB) ───────╯

  Max safe batch size: 512

---

How It Works

Kitsune statically estimates four components without running the model:

Component	Method	Accuracy
Parameters	sum(p.numel() × element_size)	< 1% error
Gradients	Same as trainable params	< 1% error
Optimizer states	Adam = 2× params in fp32	Exact
Activations	torch.fx trace + meta-device hooks	~15–30%
Safety buffer	+15% overhead margin	Conservative

Validation results (real MPS measurements vs. predictions):

• Average peak error: 16.7%
• Within 20%: 8/9 test cases
• Parameter estimation: < 1% error on all models

---

Supported GPUs

24 GPUs with on-demand pricing from Lambda Labs, RunPod, GCP, and AWS:

Category	GPUs
NVIDIA Data Center	H200, H100 SXM/PCIe, A100 40/80GB, A10G, L40S, L4, V100, T4, A30, A6000
Consumer (cloud-available)	RTX 4090, 4080, 4070 Ti, 4070, 3090, 3080
Apple Silicon (local)	M4 Max, M3 Max, M2 Max, M1 Pro, M1

kitsune list-gpus   # full table with $/hr

---

Python API Reference

from kitsune.predict import (
    predict_memory,        # single prediction
    sweep_batch_sizes,     # batch size sweep table
    plan_gpu,              # GPU cost ranking
    predict_hf_memory,     # HuggingFace model (no weights downloaded)
    sweep_hf_batch_sizes,  # HuggingFace sweep
)

predict_memory()

result = predict_memory(
    model,
    input_shape=(3, 224, 224),   # without batch dim
    batch_size=32,
    dtype="float16",             # float32 | float16 | bfloat16
    optimizer="adam",            # adam | adamw | sgd | none
    training=True,
    gpu="a100-80gb",             # target GPU key
    find_max_batch_size=True,
)

result.total_gb          # 3.71
result.param_gb          # 0.10
result.activation_gb     # 3.31
result.fits_on_gpu       # True
result.headroom_gb       # 4.79
result.max_batch_size    # 128
result.breakdown_dict()  # {"Parameters": 0.10, "Activations": 3.31, ...}
result.recommendations   # ranked list of memory-saving suggestions

plan_gpu()

from kitsune.predict import plan_gpu
from kitsune.predict.cost_calculator import format_plan_table

plan = plan_gpu(
    model,
    input_shape=(784,),
    training=True,
    daily_budget=30.0,    # optional: only show GPUs within budget
    min_batch_size=16,    # optional: exclude GPUs that can't run this batch size
)
print(format_plan_table(plan))

plan.cheapest            # lowest $/hr GPU that fits
plan.best_value          # lowest $/1k-steps GPU
plan.fitting_gpus        # all GPUs that fit, sorted by $/hr

sweep_batch_sizes()

from kitsune.predict import sweep_batch_sizes
from kitsune.predict.sweep import format_sweep_table

sweep = sweep_batch_sizes(
    model,
    input_shape=(3, 224, 224),
    batch_sizes=[1, 2, 4, 8, 16, 32, 64, 128],
    gpu="t4",
)
print(format_sweep_table(sweep))
sweep.max_fitting_batch_size   # 32
sweep.oom_threshold_batch_size # 64

---

Installation

# Core prediction engine (PyTorch only)
pip install torch-kitsune

# + CLI and rich terminal output
pip install torch-kitsune[predict]

# + HuggingFace model support
pip install torch-kitsune[predict,hf]

# Everything
pip install torch-kitsune[all]

---

Memory Optimization Recommendations

Kitsune ranks suggestions by savings automatically:

Technique	Typical Savings	Trade-off
Gradient Checkpointing	~70% of activations	+33% compute
Mixed Precision (FP16)	~50% of activations	Minimal accuracy impact
8-bit Optimizer	~75% of optimizer states	< 0.1% accuracy
4-bit Quantization	~75% of parameters	1–3% accuracy loss
QLoRA	~75% of parameters	Reduced vs full fine-tuning
CPU Offload	~75% of optimizer states	Slower (PCIe transfers)

---

Roadmap

• --verbose per-layer memory breakdown
• FSDP / DDP multi-GPU prediction
• KV-cache estimation for inference servers
• LoRA / QLoRA fine-tuning memory footprint
• CI/CD mode (kitsune predict --ci --budget 16GB, exits 1 if over budget)
• Spot pricing support

---

Contributing

git clone https://github.com/jeeth-kataria/Kitsune_optimization
cd Kitsune_optimization
pip install -e ".[dev,predict,hf]"
pytest tests/unit/test_predict/   # 131 unit tests
pytest tests/integration/         # accuracy validation (requires MPS or CUDA)

---

License

MIT