Metadata-Version: 2.4
Name: torch-kitsune
Version: 0.6.0
Summary: Never OOM again. Predict peak GPU memory, max batch size, and cheapest GPU for any PyTorch model — before you run.
Author-email: Jeeth Kataria <jeethkataria9798@icloud.com>
Maintainer-email: Jeeth Kataria <jeethkataria9798@icloud.com>
License: MIT
Project-URL: Homepage, https://github.com/jeeth-kataria/Kitsune_optimization
Project-URL: Documentation, https://jeeth-kataria.github.io/Kitsune_optimization
Project-URL: Repository, https://github.com/jeeth-kataria/Kitsune_optimization
Project-URL: Changelog, https://github.com/jeeth-kataria/Kitsune_optimization/blob/main/CHANGELOG.md
Project-URL: Bug Tracker, https://github.com/jeeth-kataria/Kitsune_optimization/issues
Project-URL: Release Notes, https://github.com/jeeth-kataria/Kitsune_optimization/releases
Keywords: pytorch,cuda,optimization,inference,gpu,deep-learning,apple-silicon,mps,t4,rtx,jit,torch-compile,fp16,mixed-precision,memory-pooling,performance
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Classifier: Topic :: System :: Distributed Computing
Classifier: Environment :: GPU :: NVIDIA CUDA
Classifier: Environment :: MacOS X
Classifier: Operating System :: OS Independent
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0.0
Requires-Dist: numpy>=1.21.0
Requires-Dist: packaging>=21.0
Provides-Extra: memory
Requires-Dist: supermemory>=3.0.0; extra == "memory"
Provides-Extra: triton
Requires-Dist: triton>=2.1.0; sys_platform == "linux" and extra == "triton"
Provides-Extra: tensorrt
Requires-Dist: tensorrt>=8.6.0; sys_platform == "linux" and extra == "tensorrt"
Requires-Dist: pycuda>=2022.1; sys_platform == "linux" and extra == "tensorrt"
Requires-Dist: onnx>=1.14.0; extra == "tensorrt"
Requires-Dist: onnxsim>=0.4.33; extra == "tensorrt"
Provides-Extra: dev
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: pytest-benchmark>=4.0.0; extra == "dev"
Requires-Dist: black>=23.0.0; extra == "dev"
Requires-Dist: isort>=5.12.0; extra == "dev"
Requires-Dist: mypy>=1.0.0; extra == "dev"
Requires-Dist: ruff>=0.1.0; extra == "dev"
Requires-Dist: build>=1.0.0; extra == "dev"
Requires-Dist: twine>=4.0.0; extra == "dev"
Provides-Extra: viz
Requires-Dist: matplotlib>=3.5.0; extra == "viz"
Requires-Dist: tensorboard>=2.10.0; extra == "viz"
Provides-Extra: docs
Requires-Dist: mkdocs>=1.5.0; extra == "docs"
Requires-Dist: mkdocs-material>=9.4.0; extra == "docs"
Requires-Dist: mkdocstrings[python]>=0.24.0; extra == "docs"
Requires-Dist: mkdocs-git-revision-date-localized-plugin>=1.2.0; extra == "docs"
Provides-Extra: predict
Requires-Dist: click>=8.0; extra == "predict"
Requires-Dist: rich>=13.0; extra == "predict"
Provides-Extra: hf
Requires-Dist: transformers>=4.30.0; extra == "hf"
Provides-Extra: all
Requires-Dist: torch-kitsune[dev,docs,hf,memory,predict,triton,viz]; extra == "all"
Dynamic: license-file

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

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

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

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

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

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

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

```bash
kitsune list-gpus   # full table with $/hr
```

---

## Python API Reference

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

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

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

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

```python
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_kv_cache()` and `sweep_kv_seq_lens()`

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

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

- [x] `--verbose` per-layer memory breakdown
- [x] CI/CD budget gate (`kitsune predict --ci --budget 16GB`)
- [x] Spot pricing in `kitsune list-gpus` and `kitsune plan`
- [x] LoRA / QLoRA fine-tuning memory (`kitsune lora`)
- [x] KV-cache estimation for LLM inference servers (`kitsune kv-cache`)
- [ ] FSDP / DDP multi-GPU prediction (ZeRO stages)

---

## Contributing

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