Metadata-Version: 2.4
Name: torch-kitsune
Version: 0.4.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
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 |

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

---

## Python API Reference

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

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

---

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

- [ ] `--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

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