# Optimizers

## AdamW (Standard)

Standard choice for transformer training:

```python
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=6e-4,
    betas=(0.9, 0.95),
    weight_decay=0.1,
)
```

## Muon Optimizer

Muon uses Newton-Schulz iterations to orthogonalize momentum, providing better conditioning:

```python
def newton_schulz_iteration(G, num_iters=5):
    """Orthogonalize a matrix using Newton-Schulz iteration."""
    a, b, c = (3.4445, -4.7750, 2.0315)
    X = G / (G.norm() + 1e-7)
    for _ in range(num_iters):
        A = X @ X.T
        X = a * X + b * A @ X + c * A @ A @ X
    return X


class Muon(torch.optim.Optimizer):
    """Muon optimizer with orthogonalized momentum."""

    def __init__(self, params, lr=0.02, momentum=0.95, nesterov=True):
        defaults = dict(lr=lr, momentum=momentum, nesterov=nesterov)
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self):
        for group in self.param_groups:
            lr, momentum = group['lr'], group['momentum']
            for p in group['params']:
                if p.grad is None:
                    continue

                g = p.grad
                state = self.state[p]

                if len(state) == 0:
                    state['momentum_buffer'] = torch.zeros_like(g)

                buf = state['momentum_buffer']
                buf.mul_(momentum).add_(g)

                if g.ndim >= 2:
                    g_orth = newton_schulz_iteration(buf.reshape(g.shape[0], -1))
                    g_orth = g_orth.reshape(g.shape)
                else:
                    g_orth = buf

                if group['nesterov']:
                    g_orth = g_orth.mul_(momentum).add_(g)

                p.add_(g_orth, alpha=-lr)
```

## Learning Rate Schedule

Use linear warmup followed by cosine decay:

```python
import math

def get_lr(step, warmup_steps, max_steps, max_lr, min_lr):
    if step < warmup_steps:
        return max_lr * (step + 1) / warmup_steps
    if step >= max_steps:
        return min_lr
    decay_ratio = (step - warmup_steps) / (max_steps - warmup_steps)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coeff * (max_lr - min_lr)
```

## Gradient Clipping

Always clip gradients to prevent explosions:

```python
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
```

## Optimizer Comparison

| Optimizer | LR | Momentum | Use Case |
|-----------|-----|----------|----------|
| AdamW | 6e-4 | β1=0.9, β2=0.95 | Standard training |
| Muon | 0.02 | 0.95 | Better conditioning |

## Parameter Groups (Advanced)

Different learning rates for different parameter types:

```python
# Separate embedding and other parameters
embed_params = [p for n, p in model.named_parameters() if 'wte' in n or 'lm_head' in n]
other_params = [p for n, p in model.named_parameters() if 'wte' not in n and 'lm_head' not in n]

optimizer = torch.optim.AdamW([
    {'params': embed_params, 'lr': 6e-4, 'weight_decay': 0.0},
    {'params': other_params, 'lr': 6e-4, 'weight_decay': 0.1},
])
```