Metadata-Version: 2.4
Name: trainkit-vp
Version: 1.1.0
Summary: PyTorch training utilities: cosine LR with warmup, VRAM profiling, automatic batch size and gradient accumulation planning
Author: F000NK, Voluntas Progressus
License-Expression: MIT
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.14
Classifier: Operating System :: OS Independent
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.14
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0.0
Dynamic: license-file

# trainkit-vp

PyTorch training utilities: LR scheduling, VRAM profiling, and automatic batch planning.

Part of the [MorphFormer](https://pypi.org/project/morphoformer/) project by Voluntas Progressus.

## Installation

```bash
pip install trainkit-vp
```

Requires Python >= 3.14 and PyTorch >= 2.0.

## Features

- **Cosine LR scheduler** with linear warmup (`create_cosine_schedule`)
- **Memory profiling** — detect system RAM and GPU VRAM across CUDA, ROCm, XPU, MPS backends
- **Automatic training plan** — selects batch size, gradient accumulation, gradient checkpointing, and AMP based on available VRAM
- **VRAM estimation** heuristics for Transformer models (parameters + activations + optimizer state)

## Quick Start

```python
from trainkit import detect_memory, plan_training, create_cosine_schedule, estimate_model_memory

# Profile system memory
profile = detect_memory()
profile.print_info()

# Auto-plan training based on available VRAM
model_bytes = estimate_model_memory(model)
plan = plan_training(profile, model_bytes, batch_size=64, max_len=96, d_model=512, total_layers=14)
print(plan.reason)

# Create LR schedule with warmup
scheduler = create_cosine_schedule(optimizer, warmup_steps=4000, total_steps=100000)
```

## API

| Function / Class | Description |
|---|---|
| `create_cosine_schedule(optimizer, warmup, total)` | Cosine annealing with linear warmup |
| `detect_memory()` | Returns `MemoryProfile` with RAM/VRAM info |
| `plan_training(profile, ...)` | Returns `TrainingPlan` with recommended settings |
| `estimate_model_memory(model)` | Estimate model parameter memory in bytes |
| `estimate_training_vram(...)` | Estimate total VRAM needed for training |
| `suggest_batch_size(...)` | Suggest optimal batch size for available VRAM |
| `format_bytes(n)` | Human-readable byte formatting |
| `MemoryProfile` | Dataclass with system memory info |
| `TrainingPlan` | Dataclass with recommended training config |

## License

MIT
