@r1z4x/tezgah 1.0.0

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

@@ -0,0 +1,319 @@
+---
+name: ml-training-recipes
+description: Battle-tested PyTorch training recipes for all domains — LLMs, vision, diffusion, medical imaging, protein/drug discovery, spatial omics, genomics. Covers training loops, optimizer selection (AdamW, Muon), LR scheduling, mixed precision, debugging, and systematic experimentation. Use when training or fine-tuning neural networks, debugging loss spikes or OOM, choosing architectures, or optimizing GPU throughput.
+version: 1.0.0
+author: dailycafi
+license: MIT
+tags: [PyTorch, Training, Optimization, LLM, Vision, Diffusion, Biomedical, Muon, AdamW, Debugging]
+dependencies: [torch>=2.0.0]
+---
+
+# ML Training Recipes
+
+Battle-tested patterns for PyTorch training across domains. Drawn from production codebases
+(Karpathy's autoresearch/nanochat, torchvision, HuggingFace) and modern training practice.
+
+## Reference files (read when needed)
+
+- `references/architecture.md` — Transformer/LLM architecture code patterns, weight init
+- `references/optimizers.md` — Muon, AdamW hybrid, per-group LR, compiled optimizer steps
+- `references/domain-specific.md` — Vision, diffusion, contrastive, distributed, checkpointing, data loading
+- `references/scaling-and-selection.md` — Scaling laws, compute budget tables, decision trees, DGX Spark
+- `references/biomedical.md` — Drug discovery, protein models, medical imaging, genomics, clinical NLP
+- `references/experiment-loop.md` — Autonomous experiment loop (autoresearch keep/discard/revert)
+
+---
+
+## Architecture Selection
+
+Pick the right model by **data type** and **data scale**:
+
+| Data Type | < 10K samples | 10K-100K | > 100K |
+|-----------|--------------|----------|--------|
+| **Images** | Pretrained CNN + fine-tune | Fine-tune ViT or CNN | ViT from scratch |
+| **Text (gen)** | Few-shot prompting | Fine-tune GPT/LLaMA (LoRA) | Pretrain from scratch |
+| **Tabular** | XGBoost/LightGBM | Still XGBoost | Neural viable |
+| **Audio** | Pretrained Whisper | Fine-tune AST | Train from scratch |
+| **Molecules** | Pretrained GNN | Fine-tune molecular LM | Train GNN from scratch |
+| **Proteins** | ESM-2 embeddings + head | Fine-tune ESM-2 | Train protein LM |
+| **Medical img** | Pretrained CNN | nnU-Net (auto-config) | Swin-UNETR / MedSAM |
+
+**Key principle**: architecture matters less than training recipe at equal compute. A well-tuned
+ResNet beats a poorly-tuned ViT (ref: "ResNet Strikes Back", Wightman 2021).
+
+For biomedical domains, see `references/biomedical.md`.
+For sequence model selection and compute planning, see `references/scaling-and-selection.md`.
+
+---
+
+## Scaling Laws
+
+### Chinchilla rule (Hoffmann et al., 2022)
+
+Compute-optimal training: **~20 tokens per parameter**.
+
+| Model Size | Compute-Optimal | Inference-Optimal (100×) |
+|-----------|----------------|--------------------------|
+| 125M | 2.5B tokens | 12.5B tokens |
+| 1B | 20B tokens | 100B tokens |
+| 7B | 140B tokens | 700B tokens |
+
+**FLOPs ≈ 6 × N × D** (N=params, D=tokens). Data repetition limit: ~4 epochs before diminishing returns.
+
+---
+
+## Training Loop
+
+```python
+import gc, time, torch
+
+torch.manual_seed(42)
+torch.set_float32_matmul_precision("high")  # TF32 on Ampere+
+autocast_ctx = torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16)
+
+grad_accum_steps = total_batch_size // (batch_size * seq_len)
+step = 0
+
+while not done:
+    t0 = time.time()
+    for micro_step in range(grad_accum_steps):
+        with autocast_ctx:
+            loss = model(x, y)
+        (loss / grad_accum_steps).backward()
+        x, y = next(train_loader)
+
+    update_lr(optimizer, progress)
+    optimizer.step()
+    model.zero_grad(set_to_none=True)  # frees memory vs zeroing
+
+    if loss.item() > 100:  # fast-fail on divergence
+        print("FAIL: loss exploded"); exit(1)
+
+    torch.cuda.synchronize()
+    if step == 0:
+        gc.collect(); gc.freeze(); gc.disable()  # avoid ~500ms GC stalls
+    step += 1
+```
+
+### Key principles
+
+- **Gradient clipping**: `clip_grad_norm_(params, 1.0)` — near-universal for Transformers.
+  Exception: Muon optimizer normalizes updates via orthogonalization, so clipping is optional.
+- **Tensor Core alignment**: batch size, hidden dims should be multiples of 8 (bf16) or 64 (A100).
+- **Time-based budgets** make experiments comparable across hardware.
+- **`cudnn.benchmark = True`** for fixed-size vision inputs.
+
+---
+
+## Optimizer Configuration
+
+Modern LLM training uses different optimizers per parameter group:
+
+| Parameter Type | Optimizer | LR (base) | Weight Decay |
+|---------------|-----------|-----------|--------------|
+| 2D weight matrices | Muon | 0.04 | 0.2 |
+| Token embeddings | AdamW | 0.6 × scale | 0.0 |
+| Unembedding (lm_head) | AdamW | 0.004 × scale | 0.0 |
+| Per-layer scalars | AdamW | 0.005 × scale | 0.0 |
+
+**LR scaling by dimension**: `lr * (d_model / 768)^(-0.5)` — keeps dynamics stable across sizes.
+
+### Rules of thumb
+
+- Embeddings need higher LR (sparse updates). Never weight-decay embeddings.
+- Weight decay scheduling: linearly decay WD to 0 over training.
+- AdamW defaults: β1=0.9, β2=0.95, eps=1e-10 (not default 1e-8 — prevents stale updates in bf16).
+
+For Muon details (polar express orthogonalization, NorMuon), see `references/optimizers.md`.
+
+---
+
+## Learning Rate Scheduling
+
+### Time-based (autoresearch style)
+
+```python
+def get_lr_multiplier(progress):  # progress = elapsed_time / time_budget
+    if progress < warmup_ratio:
+        return progress / warmup_ratio
+    elif progress < 1.0 - warmdown_ratio:
+        return 1.0
+    else:
+        cooldown = (1.0 - progress) / warmdown_ratio
+        return cooldown + (1 - cooldown) * final_lr_frac
+```
+
+### Cosine decay
+
+```python
+def get_lr(step, total_steps, max_lr, min_lr, warmup_steps):
+    if step < warmup_steps:
+        return max_lr * step / warmup_steps
+    progress = (step - warmup_steps) / (total_steps - warmup_steps)
+    return min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * progress))
+```
+
+**WSD (Warmup-Stable-Decay)**: gaining traction — easier to resume training mid-run.
+
+### Guidance
+
+- **Warmup**: 1-5% of training. Zero warmup valid with Muon (autoresearch uses `WARMUP_RATIO=0.0`).
+- **Warmdown**: 30-50% of training in LR decay. Matters more than warmup for final quality.
+- **Final LR**: 0 or ~10% of peak. Zero is simpler.
+
+---
+
+## Mixed Precision & Compilation
+
+```python
+import os
+os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"  # before torch import
+
+import torch
+torch.set_float32_matmul_precision("high")
+autocast_ctx = torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16)
+model = torch.compile(model, dynamic=False)
+```
+
+- **bf16** (Ampere+): same exponent as fp32, no loss scaling needed. Preferred over fp16.
+- **fp16**: needs GradScaler. Use only on V100 or older.
+- `dynamic=False` enables max optimization. Add `fullgraph=True` if no graph breaks.
+- First steps are slow (JIT) — exclude from timing.
+
+---
+
+## Memory & Performance
+
+### Meta device init (large models)
+
+```python
+with torch.device("meta"):
+    model = GPT(config)          # zero memory
+model.to_empty(device="cuda")
+model.init_weights()
+```
+
+### MFU (Model FLOPs Utilization)
+
+```python
+achieved_flops = model_flops_per_token * batch_tokens / step_time
+mfu = achieved_flops / gpu_peak_flops
+# H100 SXM: 989.5 TFLOPS | A100: 312 | RTX 4090: 165
+```
+
+Good targets: >30% decent, >40% good, >50% excellent (single-GPU).
+
+### OOM solutions (in order)
+
+1. Reduce `DEVICE_BATCH_SIZE`, increase `grad_accum_steps`
+2. `PYTORCH_ALLOC_CONF=expandable_segments:True`
+3. `model.zero_grad(set_to_none=True)`
+4. Meta device init → `to_empty`
+5. Activation checkpointing: `torch.utils.checkpoint.checkpoint()`
+6. 8-bit optimizer (bitsandbytes): ~30% savings on optimizer states
+
+---
+
+## Hyperparameter Search
+
+### Priority order (tune first → last)
+
+1. **Learning rate** — most impactful. Always tune first.
+2. **Batch size** — largest that fits. Speed knob, not quality knob.
+3. **Weight decay** — 0.01-0.1 for AdamW.
+4. **Warmup steps** — 1-5% of training.
+
+### The 2025 default recipe
+
+| Setting | Value |
+|---------|-------|
+| Optimizer | AdamW (β1=0.9, β2=0.95, eps=1e-10) |
+| Weight decay | 0.1 |
+| LR schedule | Cosine decay or WSD |
+| Peak LR | 3e-4 (scale down for larger models) |
+| Precision | bf16 |
+| Grad clipping | max_norm=1.0 |
+| Normalization | RMSNorm (pre-norm) |
+| Activation | SwiGLU |
+| Position encoding | RoPE |
+| Attention | Flash Attention, optionally GQA |
+
+---
+
+## Debugging Checklist
+
+### Karpathy's recipe (still canonical)
+
+1. **Become one with the data** — visualize, check distributions, verify labels
+2. **Get end-to-end running first** — verify on a trivial case
+3. **Overfit one batch** — if you can't, you have a bug
+4. **Then regularize** — add regularization only after overfitting works
+5. **Tune hyperparameters** — start with known defaults
+
+### Loss exploding / NaN
+
+1. Reduce LR (3-10× smaller)
+2. Add gradient clipping: `clip_grad_norm_(params, 1.0)`
+3. Check for inf/nan in inputs
+4. Add logit soft capping: `softcap * tanh(logits / softcap)`
+5. Add QK-norm in attention
+6. Verify weight init (zero-init output projections?)
+7. Check loss reduction with gradient accumulation (`loss / grad_accum_steps`)
+
+### Slow training / Low MFU
+
+1. Verify `torch.compile` is active
+2. Check `torch.set_float32_matmul_precision("high")`
+3. Pin memory + non_blocking transfers
+4. Profile with `torch.profiler`
+5. GC stalls? `gc.freeze(); gc.disable()`
+6. Tensor Core alignment: dims multiples of 8/64
+
+### Loss plateau / Slow convergence
+
+1. LR too low — try 2-5× larger
+2. Warmup too long
+3. Weight decay too high
+4. Verify LR schedule is actually applied (print each step)
+5. Model too small for task
+
+### Silent failures
+
+1. **Data leakage** between train/val
+2. **Wrong preprocessing at inference** — augmentation mismatch
+3. **Label errors** — use cleanlab to detect
+4. **Shuffling bugs** — correlated batches
+5. **Tokenizer mismatch** with pretrained model
+
+### What to monitor
+
+- **Gradient norms** — spike precedes loss spike
+- **Per-layer activation stats** — reveals exploding/vanishing
+- **Dead neurons** — >50% zero ReLU = dying ReLU problem
+- **Learning rate** — verify schedule applied (common silent bug)
+
+---
+
+## Experiment Management
+
+Track experiments in TSV for easy comparison:
+
+```
+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 (worse)
+```
+
+**Simplicity criterion**: all else equal, simpler is better. Removing something and getting equal
+results is a great outcome. For systematic agent-driven experimentation, see `references/experiment-loop.md`.
+
+### Evaluation metrics by domain
+
+| Domain | Primary Metric | Notes |
+|--------|---------------|-------|
+| LLM | BPB (bits per byte) | Vocab-size-independent |
+| Classification | Accuracy / F1 | Macro-F1 for imbalanced |
+| Segmentation | mIoU / Dice | Per-class IoU reveals weak spots |
+| Generation | FID | Needs >10k samples |
+| Regression | RMSE / MAE | Log-transform skewed targets |

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

@@ -0,0 +1,320 @@
+<!-- vendored from orchestra-research/AI-research-SKILLs@773a52944ba4747a18bd4ae9ade53fff041adcbc 10-optimization/ml-training-recipes/references/architecture.md; MIT (c) Orchestra Research -->
+# Architecture Patterns Reference
+
+Detailed code patterns for modern transformer architectures. Referenced from the main SKILL.md.
+
+## Table of Contents
+
+1. [RMSNorm](#rmsnorm)
+2. [Rotary Position Embeddings (RoPE)](#rotary-position-embeddings-rope)
+3. [Flash Attention with Sliding Window](#flash-attention-with-sliding-window)
+4. [Grouped Query Attention (GQA)](#grouped-query-attention-gqa)
+5. [Value Embedding (ResFormer)](#value-embedding-resformer)
+6. [Activation Functions](#activation-functions)
+7. [Residual Scaling](#residual-scaling)
+8. [Logit Soft Capping](#logit-soft-capping)
+9. [Full Transformer Block](#full-transformer-block)
+10. [Model Configuration Pattern](#model-configuration-pattern)
+
+---
+
+## RMSNorm
+
+Root Mean Square Layer Normalization — drops the mean-centering of LayerNorm, keeping only the
+variance normalization. ~15% faster with equivalent quality for transformers.
+
+```python
+def norm(x):
+    return F.rms_norm(x, (x.size(-1),))
+```
+
+Apply pre-norm (before attention and MLP), not post-norm:
+```python
+class Block(nn.Module):
+    def forward(self, x):
+        x = x + self.attn(norm(x))   # pre-norm
+        x = x + self.mlp(norm(x))    # pre-norm
+        return x
+```
+
+Also normalize the final output before the unembedding layer:
+```python
+x = norm(x)
+logits = self.lm_head(x)
+```
+
+---
+
+## Rotary Position Embeddings (RoPE)
+
+RoPE encodes position through rotation of query/key pairs. It's relative (only depends on
+distance between tokens) and naturally handles varying sequence lengths.
+
+### Precomputation
+
+Compute cos/sin tables once at model init, not every forward pass:
+
+```python
+def precompute_rotary(seq_len, head_dim, base=10000, device=None):
+    """Precompute RoPE cos/sin for positions [0, seq_len)."""
+    channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+    inv_freq = 1.0 / (base ** (channel_range / head_dim))
+    t = torch.arange(seq_len, dtype=torch.float32, device=device)
+    freqs = torch.outer(t, inv_freq)
+    cos, sin = freqs.cos().bfloat16(), freqs.sin().bfloat16()
+    # Shape: [1, seq_len, 1, head_dim//2] for broadcasting
+    return cos[None, :, None, :], sin[None, :, None, :]
+```
+
+Register as non-persistent buffers (not saved in state_dict, but moved with `.to(device)`):
+```python
+self.register_buffer("cos", cos, persistent=False)
+self.register_buffer("sin", sin, persistent=False)
+```
+
+### Application
+
+```python
+def apply_rotary_emb(x, cos, sin):
+    """Apply RoPE to query or key tensor. x shape: [B, T, H, D]."""
+    d = x.shape[3] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    return torch.cat([y1, y2], dim=3)
+```
+
+### Tips
+- Pre-allocate for `seq_len * 10` (or max expected length) to avoid recomputation
+- Apply RoPE **after** splitting into heads but **before** attention
+- Normalize q and k **after** RoPE: `q, k = norm(q), norm(k)` (QK-norm stabilizes training)
+
+---
+
+## Flash Attention with Sliding Window
+
+Flash Attention computes exact attention in O(N) memory instead of O(N^2), and is significantly
+faster due to IO-awareness.
+
+### Sliding Window Pattern
+
+Use a repeating pattern like `SSSL` — most layers use short (local) windows, with periodic long
+(global) windows. The last layer always gets full context.
+
+```python
+def compute_window_sizes(config):
+    pattern = config.window_pattern.upper()  # e.g., "SSSL"
+    long_window = config.sequence_len
+    short_window = long_window // 2  # half context
+
+    window_sizes = []
+    for layer_idx in range(config.n_layer):
+        char = pattern[layer_idx % len(pattern)]
+        if char == "L":
+            window_sizes.append((long_window, 0))
+        else:
+            window_sizes.append((short_window, 0))
+
+    # Last layer always gets full context
+    window_sizes[-1] = (long_window, 0)
+    return window_sizes
+```
+
+This saves ~25% attention compute while maintaining quality — most layers only need local context,
+and information propagates through the occasional global layer.
+
+### Integration
+
+```python
+# Using Flash Attention 3
+from kernels import get_kernel
+fa3 = get_kernel("kernels-community/flash-attn3").flash_attn_interface
+
+y = fa3.flash_attn_func(q, k, v, causal=True, window_size=window_size)
+
+# Or using PyTorch native (2.0+)
+y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+```
+
+---
+
+## Grouped Query Attention (GQA)
+
+Use fewer KV heads than query heads. Saves memory/compute with minimal quality loss.
+
+```python
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.n_head = config.n_head        # e.g., 12
+        self.n_kv_head = config.n_kv_head  # e.g., 4 (GQA) or 1 (MQA)
+        self.head_dim = config.n_embd // config.n_head
+
+        assert config.n_head % config.n_kv_head == 0
+
+        self.c_q = nn.Linear(config.n_embd, self.n_head * self.head_dim, bias=False)
+        self.c_k = nn.Linear(config.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_v = nn.Linear(config.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
+```
+
+Common ratios:
+- **MHA** (multi-head): `n_kv_head = n_head` — full quality, most memory
+- **GQA**: `n_kv_head = n_head / 4` — good tradeoff
+- **MQA** (multi-query): `n_kv_head = 1` — most memory savings, slight quality loss
+
+---
+
+## Value Embedding (ResFormer)
+
+Alternating layers receive value embeddings — learned per-token vectors added to the V projection
+with an input-dependent gate. This creates a "value residual stream" parallel to the main residual.
+
+```python
+def has_ve(layer_idx, n_layer):
+    """Alternating layers get value embeddings, last layer always included."""
+    return layer_idx % 2 == (n_layer - 1) % 2
+
+# In attention forward:
+if ve is not None:
+    ve = ve.view(B, T, self.n_kv_head, self.head_dim)
+    # Input-dependent gate: sigmoid output scaled by 2 (neutral at init)
+    gate = 2 * torch.sigmoid(self.ve_gate(x[..., :gate_channels]))
+    v = v + gate.unsqueeze(-1) * ve
+```
+
+Initialize gate weights to zero so `sigmoid(0) = 0.5`, scaled by 2 = 1.0 (neutral start):
+```python
+nn.init.zeros_(block.attn.ve_gate.weight)
+```
+
+---
+
+## Activation Functions
+
+### ReluSquared (recommended for simplicity)
+```python
+def forward(self, x):
+    x = self.c_fc(x)
+    x = F.relu(x).square()  # sparse + smooth
+    x = self.c_proj(x)
+    return x
+```
+Benefits: naturally sparse (ReLU zeros + squaring), simple, good performance.
+
+### SwiGLU (recommended for quality)
+```python
+class SwiGLUMLP(nn.Module):
+    def __init__(self, config):
+        hidden = int(config.n_embd * 8 / 3)  # ~2.67x, compensate for gate
+        hidden = ((hidden + 63) // 64) * 64   # round to 64 for efficiency
+        self.w1 = nn.Linear(config.n_embd, hidden, bias=False)
+        self.w2 = nn.Linear(hidden, config.n_embd, bias=False)
+        self.w3 = nn.Linear(config.n_embd, hidden, bias=False)  # gate
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+```
+
+### GELU (safe default)
+```python
+x = F.gelu(self.c_fc(x))
+```
+
+---
+
+## Residual Scaling
+
+Learnable per-layer residual scaling stabilizes deep networks:
+
+```python
+class GPT(nn.Module):
+    def __init__(self, config):
+        self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))
+        self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))
+
+    def forward(self, idx):
+        x = norm(self.wte(idx))
+        x0 = x  # save initial representation
+
+        for i, block in enumerate(self.transformer.h):
+            # x0 skip connection: mix in initial representation
+            x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
+            x = block(x, ...)
+
+        return norm(x)
+```
+
+Initialize: `resid_lambdas = 1.0` (normal residual), `x0_lambdas = 0.1` (small initial skip).
+
+This helps because:
+- Deep networks can have vanishing/exploding residual norms
+- x0 skip connections let gradients flow directly to the embedding layer
+- Learnable scaling lets the network decide how much skip vs. residual per layer
+
+---
+
+## Logit Soft Capping
+
+Prevents extreme logit values that cause training instability:
+
+```python
+softcap = 15
+logits = self.lm_head(x).float()  # compute in fp32 for stability
+logits = softcap * torch.tanh(logits / softcap)
+```
+
+This smoothly clamps logits to [-softcap, +softcap]. Values in the normal range (much smaller
+than softcap) pass through nearly unchanged; extreme values are compressed.
+
+---
+
+## Model Configuration Pattern
+
+Use a dataclass for clean configuration:
+
+```python
+@dataclass
+class GPTConfig:
+    sequence_len: int = 2048
+    vocab_size: int = 32768
+    n_layer: int = 12
+    n_head: int = 6
+    n_kv_head: int = 6
+    n_embd: int = 768
+    window_pattern: str = "SSSL"
+
+def build_config(depth, aspect_ratio=64, head_dim=128):
+    """Derive model dimensions from depth using aspect ratio."""
+    base_dim = depth * aspect_ratio
+    model_dim = ((base_dim + head_dim - 1) // head_dim) * head_dim  # round to head_dim
+    num_heads = model_dim // head_dim
+    return GPTConfig(n_layer=depth, n_head=num_heads, n_kv_head=num_heads, n_embd=model_dim)
+```
+
+The aspect ratio pattern (`d_model = depth * ratio`) keeps width proportional to depth,
+which empirical research shows is more compute-efficient than scaling width alone.
+
+---
+
+## FLOPs Estimation
+
+For monitoring MFU, estimate FLOPs per token:
+
+```python
+def estimate_flops_per_token(model):
+    """Forward + backward FLOPs per token (approx 6 * params + attention)."""
+    # Main rule: 6 * N (2 for fwd matmuls, 4 for bwd matmuls per param)
+    # Exclude embeddings (sparse lookups, not matmuls)
+    nparams_dense = sum(p.numel() for p in model.parameters())
+    nparams_dense -= model.wte.weight.numel()      # token embedding
+    nparams_dense -= model.lm_head.weight.numel()   # if tied, already counted
+
+    # Attention FLOPs: 2 * n_heads * head_dim * seq_len per layer (Q@K + attn@V)
+    attn_flops = 0
+    for window in model.window_sizes:
+        effective_seq = min(window[0], model.config.sequence_len)
+        attn_flops += 12 * model.config.n_head * head_dim * effective_seq
+
+    return 6 * nparams_dense + attn_flops
+```