Metadata-Version: 2.4
Name: optifine-torch
Version: 0.1.0
Summary: Como el mod de Minecraft, pero para tu training loop de PyTorch en vez de tus FPS.
Requires-Python: >=3.9
Description-Content-Type: text/markdown
Requires-Dist: torch>=2.0.0

# optifine

Como el mod de Minecraft, pero en vez de FPS optimiza tu training loop de PyTorch.

Pensado para GPUs chicas (T4, T4x2 en Kaggle) entrenando modelos chicos-medianos
desde cero. No es magia, son las optimizaciones estandar (mixed precision, fused
optimizer, flash attention, gradient accumulation) empaquetadas para no reescribir
el mismo boilerplate en cada proyecto nuevo.

## Instalar

```bash
pip install -e .
```

## Uso rápido

```python
import torch
import torch.nn as nn
from optifine import FastTrainer, OptifineConfig, FastAttention, make_optimizer, try_compile

# 1. Reemplaza tu attention manual por FastAttention (usa el kernel fused de torch)
class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = FastAttention(n_embd, n_head, block_size, dropout=0.1)
        # ... resto del block igual que siempre

# 2. Optimizer fused (gratis en GPUs modernas)
optimizer = make_optimizer(model.parameters(), lr=3e-4, weight_decay=0.01)

# 3. (opcional) torch.compile con fallback seguro si falla
model = try_compile(model)

# 4. Trainer con AMP + accumulation + checkpointing automático
config = OptifineConfig(
    grad_accum_steps=2,
    amp_dtype="bf16",          # o "fp16" si tu GPU no soporta bf16
    checkpoint_dir="/kaggle/working/checkpoints",
    checkpoint_every=500,
)
trainer = FastTrainer(model, optimizer, config, run_name="normal")

# resume automático si hay checkpoint previo
trainer.load_checkpoint()

def step_fn(batch):
    x, y = batch
    logits, loss = model(x, y)
    return loss

for batch in dataloader:
    stats = trainer.step(step_fn, batch, extra_upload=mi_funcion_subir_a_hf)
    if stats["stepped"] and stats["iter"] % 50 == 0:
        print(f"iter {stats['iter']} | loss {stats['loss']:.4f}")
```

## Qué incluye

- **`FastAttention`**: drop-in de self-attention causal usando
  `F.scaled_dot_product_attention` (flash attention built-in de PyTorch 2.x,
  no hay que instalar nada aparte).
- **`make_optimizer`**: AdamW con `fused=True` cuando está disponible, con
  fallback transparente si la GPU/torch no lo soporta.
- **`try_compile`**: envuelve `torch.compile` con try/except — si falla
  (pasa seguido en Kaggle por versiones de triton), sigue sin compilar en
  vez de tirar abajo el training.
- **`FastTrainer`**: loop de un paso con autocast (bf16/fp16), gradient
  accumulation, gradient clipping, y checkpointing periódico con resume
  automático — mismo patrón que usarías a mano, pero sin repetir el
  boilerplate cada vez.

## Notas honestas

- `bf16` es preferible a `fp16` cuando la GPU lo soporta (T4 la soporta):
  no necesita `GradScaler`, es más estable, sin drama.
- Ganancia esperada: mixed precision solo ya suele dar 2-3x en GPUs con
  tensor cores. `torch.compile` puede sumar otro 20-40%, pero varía mucho
  según arquitectura y a veces no vale la pena en modelos muy chicos (el
  overhead de compilar puede no amortizarse en pocas iteraciones).
- Esto no reemplaza pensar el tamaño de modelo, dataset o batch size — es
  optimización de "cuánto tarda cada iteración", no de "cuántas
  iteraciones necesitás".
