torch-kitsune 0.7.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 (now includes spot pricing)
kitsune plan --hf bert-base-uncased --training
kitsune plan --hf gpt2 --training --daily-budget 50

# KV-cache estimation for inference servers
kitsune kv-cache --hf gpt2 --seq-len 1024 --batch-size 8 --gpu t4
kitsune kv-cache --hf facebook/opt-125m --sweep --gpu a100-40gb   # find max context

# Multi-GPU: DDP vs FSDP ZeRO-1/2/3 comparison
kitsune multi-gpu --hf bert-large-uncased --gpu a100-40gb
kitsune multi-gpu --hf bert-large-uncased --strategy zero3 --gpus 4 --gpu a100-40gb

# LoRA / QLoRA fine-tuning memory
kitsune lora --hf bert-base-uncased --rank 16 --gpu t4
kitsune lora --hf bert-large-uncased --rank 64 --qlora --gpu a100-40gb

# CI/CD budget gate (exits 1 if over budget)
kitsune predict --hf bert-base-uncased --budget 8GB --ci
kitsune predict --hf bert-large-uncased --inference --budget 4GB --ci --format json

# 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
kitsune lora model.py:ResNet50 --input-shape 3,224,224 --rank 16

# List all GPUs with pricing (on-demand + spot)
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
    estimate_lora_memory,  # LoRA / QLoRA fine-tuning memory
    estimate_kv_cache,     # KV-cache for LLM inference
    sweep_kv_seq_lens,     # context length sweep
    estimate_multi_gpu,    # DDP / FSDP ZeRO per-GPU memory
    plan_multi_gpu,        # compare all strategies × GPU counts
    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

estimate_lora_memory()

from kitsune.predict import estimate_lora_memory, LoRAConfig

config = LoRAConfig(
    rank=16,
    target_modules=("q_proj", "v_proj", "query", "value"),  # LLaMA or BERT naming
    use_qlora=False,   # True = 4-bit quantized base (QLoRA)
    optimizer="adamw",
)

pred = estimate_lora_memory(model, input_shape=(512,), batch_size=4, config=config)
print(pred.summary())
# LoRA r=16: 0.48 GB total | 589,824 trainable params (24 layers)

pred.total_gb                       # 0.480
pred.base_gb                        # 0.408  (frozen base)
pred.adapter_gb                     # 0.008  (adapters + grads + optimizer states)
pred.savings_vs_full_finetune_gb    # 1.21   (memory saved vs full fine-tuning)
pred.num_adapter_params             # 589824
pred.num_targeted_layers            # 24
pred.breakdown_dict()               # {"Base model (frozen)": 0.408, ...}

estimate_multi_gpu() and plan_multi_gpu()

from kitsune.predict import estimate_multi_gpu, plan_multi_gpu, Strategy

# Single strategy
result = estimate_multi_gpu(
    model,
    input_shape=(784,),
    total_batch_size=32,
    strategy=Strategy.ZERO3,   # "ddp", "zero1", "zero2", "zero3"
    world_size=4,
    gpu="a100-40gb",
)
print(result.summary())
# FSDP ZeRO-3 × 4 GPUs: 0.71 GB/GPU (saved 1.85 GB vs single-GPU) | A100 40GB: FITS

result.total_gb                       # 0.707  (per GPU)
result.param_bytes                    # params/4 (sharded)
result.savings_vs_single_gpu_gb       # 1.85

# Compare ALL strategies × world sizes at once
plan = plan_multi_gpu(model, input_shape=(784,), gpu="t4", world_sizes=[1, 2, 4, 8])
print(f"Min GPUs for ZeRO-3 to fit: {plan.min_gpus_needed}")
print(f"Best: {plan.best_fit.summary()}")

estimate_kv_cache() and sweep_kv_seq_lens()

from kitsune.predict import estimate_kv_cache, sweep_kv_seq_lens, KVCacheConfig

# From a HuggingFace model (config auto-detected)
from kitsune.predict.hf_integration import load_hf_model
model, _ = load_hf_model("gpt2")

result = estimate_kv_cache(model, seq_len=1024, batch_size=8, dtype="float16", gpu="t4")
print(result.summary())
# seq=1024 bs=8: 0.31 GB total (KV=0.28 GB, 90%) | Tesla T4 (16GB): FITS

result.kv_cache_gb          # 0.281  (just the KV tensors)
result.weight_gb            # 0.231  (model weights)
result.kv_pct_of_total      # 90.2   (KV dominates at long context)
result.fits                 # True

# Or pass config directly (no model needed)
config = KVCacheConfig(
    num_layers=32,
    num_kv_heads=8,       # GQA: LLaMA-2 style
    head_dim=128,
    max_seq_len=4096,
    num_attention_heads=32,
)
result = estimate_kv_cache(config, seq_len=4096, batch_size=4, gpu="a100-40gb")

# Sweep context lengths to find the OOM boundary
sweep = sweep_kv_seq_lens(model, batch_size=8, gpu="t4")
sweep.max_fitting_seq_len      # 1024
sweep.oom_threshold_seq_len    # 2048

---

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
• CI/CD budget gate (kitsune predict --ci --budget 16GB)
• Spot pricing in kitsune list-gpus and kitsune plan
• LoRA / QLoRA fine-tuning memory (kitsune lora)
• KV-cache estimation for LLM inference servers (kitsune kv-cache)
• FSDP / DDP multi-GPU prediction — ZeRO-1/2/3 comparison (kitsune multi-gpu)

---

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