@r1z4x/tezgah 1.0.0

package/skills/ai-research/10-optimization/ml-training-recipes/references/experiment-loop.md

@@ -0,0 +1,129 @@
+<!-- vendored from orchestra-research/AI-research-SKILLs@773a52944ba4747a18bd4ae9ade53fff041adcbc 10-optimization/ml-training-recipes/references/experiment-loop.md; MIT (c) Orchestra Research -->
+# Autonomous Experiment Loop (autoresearch pattern)
+
+A systematic workflow for rapid ML experimentation, drawn from Karpathy's autoresearch project.
+Use this when iterating on architecture or hyperparameters and you want to run many quick experiments.
+
+## Core idea
+
+Run every experiment with a **fixed time budget** (e.g., 5 minutes) so results are directly comparable.
+This enables ~12 experiments/hour or ~100 overnight. The key insight: wall-clock time is a better
+budget unit than steps or epochs because it naturally accounts for throughput differences between configs.
+
+## The experiment loop
+
+```
+1. Read current state (results.tsv, train.py)
+2. Decide what to try next (one change at a time)
+3. Modify train.py
+4. git commit -m "description of change"
+5. Run training (with timeout)
+6. Parse results from stdout
+7. Decision:
+   - If val_bpb improved → KEEP (advance branch)
+   - If val_bpb worsened → DISCARD (git reset --hard HEAD~1)
+   - If crashed → FIX trivial bugs and retry, or LOG and move on
+8. Append result to results.tsv
+9. Repeat
+```
+
+## Results tracking
+
+```
+commit    val_bpb   memory_gb  status   description
+a1b2c3d   0.9979    44.0       keep     baseline
+b2c3d4e   0.9932    44.2       keep     increase matrix LR to 0.04
+c3d4e5f   1.0050    44.0       discard  switch to GeLU activation
+d4e5f6g   0.0000    0.0        crash    double model width (OOM)
+```
+
+## Key principles
+
+### Single-file constraint
+Confine all changes to one file (e.g., `train.py`). This makes diffs reviewable and rollbacks clean.
+Everything — model, optimizer, data loading, evaluation — lives in one file during experimentation.
+Refactor into modules only after the experiment phase.
+
+### Keep/discard discipline
+- **Keep**: val metric improved (or equal with less memory/time)
+- **Discard**: val metric worsened, regardless of how clever the idea was
+- **The simplicity criterion**: all else being equal, simpler is better. Removing something and
+  getting equal results is a great outcome — it means the removed thing was dead weight.
+
+### Crash recovery
+- **Trivial crash** (typo, shape mismatch): fix and retry the same experiment
+- **Fundamental crash** (OOM, numerical instability): log as `crash`, move on
+- **Timeout** (>2x budget): kill the process, log as `timeout`
+
+### Fixed budget comparison
+```python
+import time
+
+TIME_BUDGET = 300  # 5 minutes
+t_start = time.time()
+
+for step in range(max_steps):
+    # ... training step ...
+    elapsed = time.time() - t_start
+    if elapsed >= TIME_BUDGET:
+        break
+```
+
+## Tokenizer training
+
+When training from scratch, train a BPE tokenizer on your data:
+
+```python
+import rustbpe
+
+# GPT-4 split pattern (handles code, numbers, whitespace well)
+SPLIT_PATTERN = r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+"""
+
+# Train tokenizer
+tokenizer = rustbpe.Tokenizer()
+tokenizer.train(
+    text_iterator,        # yields text chunks
+    vocab_size=8192,      # small vocab for quick experiments; 32K+ for production
+    split_pattern=SPLIT_PATTERN,
+    special_tokens=["<|bos|>"]
+)
+
+# Build token_bytes lookup for BPB evaluation
+token_bytes = torch.zeros(vocab_size, dtype=torch.long)
+for i in range(vocab_size):
+    token_bytes[i] = len(tokenizer.decode([i]).encode("utf-8"))
+```
+
+### Vocab size tradeoffs
+
+| Vocab Size | Use Case | Notes |
+|-----------|----------|-------|
+| 4K-8K | Quick experiments, small models | Faster tokenizer training, more tokens per doc |
+| 32K | Standard LLM pretraining | Good balance of compression and vocab coverage |
+| 64K-128K | Multilingual, code-heavy | Better compression but larger embedding table |
+
+## Data preparation
+
+### Shard-based train/val split
+```python
+# Use last shard as validation (always the same data for consistent eval)
+shard_files = sorted(glob("data/shard_*.bin"))
+val_shard = shard_files[-1]       # pinned validation
+train_shards = shard_files[:-1]   # everything else
+```
+
+Split by shard, not by random sampling — this ensures no data leakage and makes
+the val set deterministic across experiments.
+
+## Environment setup
+
+```python
+import os
+os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"  # BEFORE torch import
+os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"               # clean logs
+
+import torch
+```
+
+Setting `PYTORCH_ALLOC_CONF` before importing torch is important — it configures the
+CUDA allocator at initialization time.

package/skills/ai-research/10-optimization/ml-training-recipes/references/optimizers.md

@@ -0,0 +1,330 @@
+<!-- vendored from orchestra-research/AI-research-SKILLs@773a52944ba4747a18bd4ae9ade53fff041adcbc 10-optimization/ml-training-recipes/references/optimizers.md; MIT (c) Orchestra Research -->
+# Optimizer Patterns Reference
+
+Deep dive into optimizer configurations for modern LLM training. Referenced from the main SKILL.md.
+
+## Table of Contents
+
+1. [AdamW Best Practices](#adamw-best-practices)
+2. [Muon Optimizer](#muon-optimizer)
+3. [Hybrid MuonAdamW](#hybrid-muonadamw)
+4. [Per-Parameter-Group Configuration](#per-parameter-group-configuration)
+5. [LR Scaling Rules](#lr-scaling-rules)
+6. [Weight Decay Strategies](#weight-decay-strategies)
+7. [Momentum Scheduling](#momentum-scheduling)
+8. [Compiled Optimizer Steps](#compiled-optimizer-steps)
+
+---
+
+## AdamW Best Practices
+
+AdamW (decoupled weight decay) is the baseline optimizer for everything that isn't a 2D matrix
+in modern LLM training.
+
+```python
+# Typical hyperparameters for LLM pretraining
+optimizer = torch.optim.AdamW(
+    params,
+    lr=3e-4,
+    betas=(0.9, 0.95),    # β1=0.9, β2=0.95 (not the default 0.999)
+    eps=1e-8,
+    weight_decay=0.1,
+)
+```
+
+### Key differences from default PyTorch AdamW
+
+- **β2 = 0.95** (not 0.999): Faster adaptation to changing gradient statistics. The default 0.999
+  has a ~1000-step memory, too slow for the rapidly changing loss landscape of LLM training.
+- **β1 = 0.8-0.9**: Some modern recipes use 0.8 for faster momentum.
+- **eps = 1e-10** (not 1e-8): Smaller epsilon for bf16 training where gradients can be very small. autoresearch uses 1e-10; 1e-8 can cause stale updates when gradient second moments are tiny.
+
+### Fused step (for torch.compile)
+
+To avoid recompilation when hyperparameters change, use 0-D CPU tensors:
+
+```python
+# Create once at init
+self._lr_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+
+# Fill before each step (no recompile)
+self._lr_t.fill_(group['lr'])
+
+@torch.compile(dynamic=False, fullgraph=True)
+def adamw_step_fused(p, grad, exp_avg, exp_avg_sq, step_t, lr_t, beta1_t, beta2_t, eps_t, wd_t):
+    p.mul_(1 - lr_t * wd_t)
+    exp_avg.lerp_(grad, 1 - beta1_t)
+    exp_avg_sq.lerp_(grad.square(), 1 - beta2_t)
+    bias1 = 1 - beta1_t ** step_t
+    bias2 = 1 - beta2_t ** step_t
+    denom = (exp_avg_sq / bias2).sqrt() + eps_t
+    p.add_(exp_avg / denom, alpha=-lr_t / bias1)
+```
+
+---
+
+## Muon Optimizer
+
+Muon is designed for 2D matrix (weight) parameters. It uses Nesterov momentum followed by
+"Polar Express" orthogonalization — a fast Newton-Schulz iteration that approximates the
+matrix polar decomposition (finding the nearest orthogonal matrix to the gradient).
+
+### Why orthogonalize gradients?
+
+Standard gradient descent updates can create rank-deficient weight matrices over time.
+Orthogonalizing the update direction encourages diverse feature learning and prevents
+mode collapse in the weight space. Think of it as giving every update direction "equal voice."
+
+### Core algorithm
+
+1. **Nesterov momentum**: Standard momentum with look-ahead
+2. **Polar Express**: Newton-Schulz iterations to orthogonalize the gradient matrix
+3. **NorMuon**: Variance reduction that normalizes per-row or per-column
+4. **Cautious update**: Only update weights where the gradient agrees with the parameter sign
+
+```python
+@torch.compile(dynamic=False, fullgraph=True)
+def muon_step_fused(grads, params, momentum_buf, second_momentum_buf,
+                    momentum, lr, wd, beta2, ns_steps, red_dim):
+    # 1. Nesterov momentum
+    momentum_buf.lerp_(grads, 1 - momentum)
+    g = grads.lerp_(momentum_buf, momentum)
+
+    # 2. Polar Express (Newton-Schulz orthogonalization)
+    X = g.bfloat16()
+    X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.02 + 1e-6)
+    coeffs = [  # Pre-computed optimal coefficients
+        (8.16, -22.48, 15.88),
+        (4.04, -2.81, 0.50),
+        (3.89, -2.77, 0.51),
+        (3.29, -2.37, 0.46),
+        (2.35, -1.71, 0.42),
+    ]
+    # Choose which dimension to contract based on matrix shape
+    if g.size(-2) > g.size(-1):  # tall matrix
+        for a, b, c in coeffs[:ns_steps]:
+            A = X.mT @ X
+            B = b * A + c * (A @ A)
+            X = a * X + X @ B
+    else:  # wide matrix
+        for a, b, c in coeffs[:ns_steps]:
+            A = X @ X.mT
+            B = b * A + c * (A @ A)
+            X = a * X + B @ X
+    g = X
+
+    # 3. NorMuon variance reduction
+    v_mean = g.float().square().mean(dim=red_dim, keepdim=True)
+    second_momentum_buf.lerp_(v_mean, 1 - beta2)
+    step_size = second_momentum_buf.clamp_min(1e-10).rsqrt()
+    # Normalize so total norm is preserved
+    ...
+
+    # 4. Cautious weight decay + update
+    mask = (g * params) >= 0  # only decay where gradient agrees
+    params.sub_(lr * g + lr * wd * params * mask)
+```
+
+### Muon hyperparameters
+
+| Parameter | Typical Value | Notes |
+|-----------|--------------|-------|
+| lr | 0.02-0.04 | Scaled by `max(1, rows/cols)^0.5` for non-square matrices |
+| momentum | 0.95 | Warm up from 0.85 over first 300 steps |
+| ns_steps | 5 | Number of Newton-Schulz iterations (more = better approx, slower) |
+| beta2 | 0.95 | For second moment tracking in NorMuon |
+| weight_decay | 0.1-0.2 | Cautious (only where gradient agrees with param) |
+
+---
+
+## Hybrid MuonAdamW
+
+The key insight: different parameter types benefit from different optimization strategies.
+
+| Parameter Type | Optimizer | Why |
+|---------------|-----------|-----|
+| 2D weight matrices (attention, MLP) | Muon | Benefits from orthogonalization |
+| Token embeddings | AdamW | Sparse updates, not a matrix transform |
+| Unembedding (lm_head) | AdamW | Needs lower LR for stability |
+| Per-layer scalars | AdamW | Too small for matrix methods |
+| Value embeddings | AdamW | Same as token embeddings |
+
+```python
+class MuonAdamW(torch.optim.Optimizer):
+    def step(self):
+        for group in self.param_groups:
+            if group['kind'] == 'adamw':
+                self._step_adamw(group)
+            elif group['kind'] == 'muon':
+                self._step_muon(group)
+```
+
+### Grouping Muon parameters
+
+Group Muon parameters by shape for efficient stacked updates:
+
+```python
+# Group same-shape params together
+for shape in sorted({p.shape for p in matrix_params}):
+    group_params = [p for p in matrix_params if p.shape == shape]
+    param_groups.append({
+        'kind': 'muon',
+        'params': group_params,
+        'lr': matrix_lr,
+        'momentum': 0.95,
+        'ns_steps': 5,
+    })
+```
+
+This enables `torch.stack` for vectorized Newton-Schulz across all params of the same shape.
+
+---
+
+## Per-Parameter-Group Configuration
+
+A complete optimizer setup for modern LLM training:
+
+```python
+def setup_optimizer(model, d_model=768):
+    lr_scale = (d_model / 768) ** -0.5
+
+    param_groups = [
+        # Unembedding: low LR, no weight decay
+        {
+            'kind': 'adamw',
+            'params': list(model.lm_head.parameters()),
+            'lr': 0.004 * lr_scale,
+            'betas': (0.8, 0.95),
+            'eps': 1e-10,
+            'weight_decay': 0.0,
+        },
+        # Token embeddings: higher LR (sparse updates need bigger steps)
+        {
+            'kind': 'adamw',
+            'params': list(model.wte.parameters()),
+            'lr': 0.6 * lr_scale,
+            'betas': (0.8, 0.95),
+            'eps': 1e-10,
+            'weight_decay': 0.0,
+        },
+        # Transformer matrices: Muon
+        {
+            'kind': 'muon',
+            'params': list(model.transformer.h.parameters()),
+            'lr': 0.04,
+            'momentum': 0.95,
+            'ns_steps': 5,
+            'beta2': 0.95,
+            'weight_decay': 0.2,
+        },
+        # Per-layer scalars: separate AdamW
+        {
+            'kind': 'adamw',
+            'params': [model.resid_lambdas],
+            'lr': 0.005 * lr_scale,
+            'betas': (0.8, 0.95),
+            'eps': 1e-10,
+            'weight_decay': 0.0,
+        },
+    ]
+
+    # Store initial LR for scheduling
+    optimizer = MuonAdamW(param_groups)
+    for group in optimizer.param_groups:
+        group["initial_lr"] = group["lr"]
+    return optimizer
+```
+
+---
+
+## LR Scaling Rules
+
+### By model dimension
+
+As models get wider, per-parameter learning rates should decrease:
+
+```
+lr_effective = lr_base * (d_model / d_reference) ^ (-0.5)
+```
+
+This comes from the observation that larger matrices amplify gradient norms. Scaling by `1/√d`
+keeps the effective step size constant across model sizes.
+
+### By matrix shape (Muon specific)
+
+Non-square matrices need LR adjustment:
+
+```python
+effective_lr = lr * max(1.0, rows / cols) ** 0.5
+```
+
+This compensates for the asymmetry in the orthogonalization process.
+
+---
+
+## Weight Decay Strategies
+
+### Linear decay to zero
+
+```python
+def get_weight_decay(progress):
+    return base_wd * (1 - progress)
+```
+
+Rationale: early in training, regularization prevents overfitting to initial data distribution.
+Late in training, we want the model to fully commit to learned features.
+
+### Cautious weight decay (Muon)
+
+Only apply weight decay where the gradient and parameter have the same sign:
+
+```python
+mask = (gradient * parameter) >= 0
+parameter -= lr * (gradient + wd * parameter * mask)
+```
+
+This prevents weight decay from fighting the gradient — if the gradient says "increase this weight"
+but weight decay says "decrease it", cautious WD skips the decay for that element.
+
+### What to weight-decay
+
+- **Yes**: Transformer weight matrices (attention projections, MLP weights)
+- **No**: Embeddings, biases, layer norm parameters, per-layer scalars
+
+---
+
+## Momentum Scheduling
+
+Warm up momentum over the first few hundred steps:
+
+```python
+def get_muon_momentum(step, warmup_steps=300):
+    frac = min(step / warmup_steps, 1.0)
+    return 0.85 + frac * (0.95 - 0.85)  # 0.85 → 0.95
+```
+
+Lower momentum early in training allows faster adaptation when the loss landscape is changing
+rapidly. As training stabilizes, higher momentum smooths the updates.
+
+---
+
+## Compiled Optimizer Steps
+
+When using `torch.compile`, avoid recompilation from changing scalar values by using 0-D tensors:
+
+```python
+class CompiledOptimizer:
+    def __init__(self):
+        # 0-D CPU tensors: changing their values doesn't trigger recompile
+        self._lr = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._wd = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+
+    def step(self, group):
+        self._lr.fill_(group['lr'])        # update value
+        self._wd.fill_(group['weight_decay'])
+        compiled_step(params, grads, self._lr, self._wd)  # no recompile
+```
+
+This is critical for training loops where LR changes every step — without this pattern,
+`torch.compile` would recompile the optimizer step function every time the LR changes,
+defeating the purpose of compilation.