trainkit-vp 1.1.0

PyTorch training utilities: cosine LR with warmup, VRAM profiling, automatic batch size and gradient accumulation planning

trainkit-vp

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

Part of the MorphFormer project by Voluntas Progressus.

Installation

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

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