# Training Loop

## Mixed Precision Training

Use bfloat16 on A100 for 2x speedup:

```python
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for step in range(max_steps):
    # Update learning rate
    lr = get_lr(step, warmup_steps, max_steps, max_lr, min_lr)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

    # Forward pass with mixed precision
    with autocast(dtype=torch.bfloat16):
        _, loss = model(x, y)

    # Backward pass
    scaler.scale(loss).backward()

    # Gradient clipping
    scaler.unscale_(optimizer)
    grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

    # Optimizer step
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad(set_to_none=True)
```

## Complete Training Function

```python
def train_model(model, train_dataset, val_dataset, config, device="cuda"):
    model = model.to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=6e-4, weight_decay=0.1)
    scaler = GradScaler()

    # Training config
    max_steps = config.get("max_steps", 2000)
    warmup_steps = config.get("warmup_steps", 200)
    max_lr = config.get("max_lr", 6e-4)
    min_lr = config.get("min_lr", 6e-5)
    batch_size = config.get("batch_size", 32)
    eval_interval = config.get("eval_interval", 100)

    # Tracking
    train_losses = []
    val_losses = []
    grad_norms = []

    for step in range(max_steps):
        model.train()

        # Update LR
        lr = get_lr(step, warmup_steps, max_steps, max_lr, min_lr)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

        # Get batch
        x, y = train_dataset.get_batch(batch_size, device)

        # Forward/backward
        with autocast(dtype=torch.bfloat16):
            _, loss = model(x, y)

        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad(set_to_none=True)

        # Track metrics
        train_losses.append(loss.item())
        grad_norms.append(grad_norm.item())

        # Evaluation
        if step % eval_interval == 0 or step == max_steps - 1:
            model.eval()
            with torch.no_grad():
                x_val, y_val = val_dataset.get_batch(batch_size, device)
                with autocast(dtype=torch.bfloat16):
                    _, val_loss = model(x_val, y_val)
                val_losses.append(val_loss.item())
                print(f"Step {step}: train_loss={loss.item():.4f}, val_loss={val_loss.item():.4f}, grad_norm={grad_norm.item():.4f}")

    return {
        "final_val_loss": val_losses[-1],
        "train_losses": train_losses,
        "val_losses": val_losses,
        "grad_norms": grad_norms,
        "grad_norm_std": float(np.std(grad_norms)),
        "max_grad_norm": float(max(grad_norms)),
    }
```

## Evaluation

```python
@torch.no_grad()
def evaluate(model, dataset, batch_size=32, num_batches=10, device="cuda"):
    model.eval()
    losses = []
    for _ in range(num_batches):
        x, y = dataset.get_batch(batch_size, device)
        with autocast(dtype=torch.bfloat16):
            _, loss = model(x, y)
        losses.append(loss.item())
    return sum(losses) / len(losses)
```

## Progress Logging

```python
# Inside training loop
if step % 100 == 0:
    print(f"Step {step}/{max_steps}")
    print(f"  Train Loss: {loss.item():.4f}")
    print(f"  Val Loss: {val_loss.item():.4f}")
    print(f"  Grad Norm: {grad_norm.item():.4f}")
    print(f"  LR: {lr:.2e}")
```