Ml Training Recipes

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 | 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

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)

def get_lr_multiplier(progress):  # progress = elapsed_time / time_budget
    if progress 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 valbpb memorygb 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 |

## Source & license

This open-source skill is cataloged on AgentStack and links to its original source — we do not rehost the code.

- **Author:** [Orchestra-Research](https://github.com/Orchestra-Research)
- **Source:** [Orchestra-Research/AI-Research-SKILLs](https://github.com/Orchestra-Research/AI-Research-SKILLs)
- **License:** MIT
- **Homepage:** http://orchestra-research.com

Install and usage instructions live in the source repository linked above.