ml-training-recipes

A field guide to training neural networks in PyTorch: which model to start from for your kind and amount of data, optimizer and learning-rate settings, and step-by-step fixes when loss explodes, training stalls or the GPU runs out of memory. It covers language, vision and diffusion models, and biomedical work such as proteins and medical images.

WHEN TO USE

SCOPE

USE FOR

Pick a starting model for your data type and dataset size, from images to proteins.
Set learning rate, warmup and weight decay from a tested default recipe.
Work through a fixed checklist when the loss turns into NaN or explodes.
Free GPU memory step by step when training runs out of it.
Check how much of the GPU's capacity a training run actually uses.
Keep a simple table of experiments to see which changes to keep.

This is a saved copy of the skill's SKILL.md instructions. It may differ from the upstream file or the version you install.

In this document (12)

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

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

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

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)

with torch.device("meta"):
    model = GPT(config)          # zero memory
model.to_empty(device="cuda")
model.init_weights()

MFU (Model FLOPs Utilization)

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)

Reduce DEVICE_BATCH_SIZE, increase grad_accum_steps
PYTORCH_ALLOC_CONF=expandable_segments:True
model.zero_grad(set_to_none=True)
Meta device init → to_empty
Activation checkpointing: torch.utils.checkpoint.checkpoint()
8-bit optimizer (bitsandbytes): ~30% savings on optimizer states

Hyperparameter Search

Priority order (tune first → last)

Learning rate — most impactful. Always tune first.
Batch size — largest that fits. Speed knob, not quality knob.
Weight decay — 0.01-0.1 for AdamW.
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)

Become one with the data — visualize, check distributions, verify labels
Get end-to-end running first — verify on a trivial case
Overfit one batch — if you can't, you have a bug
Then regularize — add regularization only after overfitting works
Tune hyperparameters — start with known defaults

Loss exploding / NaN

Reduce LR (3-10× smaller)
Add gradient clipping: clip_grad_norm_(params, 1.0)
Check for inf/nan in inputs
Add logit soft capping: softcap * tanh(logits / softcap)
Add QK-norm in attention
Verify weight init (zero-init output projections?)
Check loss reduction with gradient accumulation (loss / grad_accum_steps)

Slow training / Low MFU

Verify torch.compile is active
Check torch.set_float32_matmul_precision("high")
Pin memory + non_blocking transfers
Profile with torch.profiler
GC stalls? gc.freeze(); gc.disable()
Tensor Core alignment: dims multiples of 8/64

Loss plateau / Slow convergence

LR too low — try 2-5× larger
Warmup too long
Weight decay too high
Verify LR schedule is actually applied (print each step)
Model too small for task

Silent failures

Data leakage between train/val
Wrong preprocessing at inference — augmentation mismatch
Label errors — use cleanlab to detect
Shuffling bugs — correlated batches
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

This skill ships inside the AI Research Skills plugin.

Run this outside your AI tool, using Node.js 22.20.0 or newer. Skills CLI 1.7.2 skips skill selection; prompts cancel installation without an interactive terminal. If choosing Project when asked about installation scope, open your project's folder before running this.

Install checklist

Check selected tools if a picker appears; Claude Code may be unchecked under Additional agents. ↑/↓ moves, Space toggles selections, Enter confirms.
Without a tool picker, check Installing to: for your intended tool. Stop if it names a different tool.
If asked, choose Symlink (Recommended) for shared files or Copy to all agents for separate copies; review the summary before confirming installation.
Install the find-skills skill? defaults to Yes, adding another skill for all projects. Choose No if you don't want that extra installation.

Skills CLI 1.7.2 sends install data to skills.sh by default, including the source, skill names, target tools and file details; GitHub sources are reported only when GitHub confirms the repository is public, while other remote sources can also be reported. Confirmed-public GitHub sources also trigger a security-audit request. Before installing, set the environment variable DO_NOT_TRACK=1 or DISABLE_TELEMETRY=1 to disable both.

npx skills add orchestra-research/ai-research-skills --skill "ml-training-recipes" --full-depth

After installing for Claude Code, open it in the installation project if you chose Project; Global makes the skill available across projects on this computer. Once Claude Code loads the skill, its description helps Claude decide when to use it.

If installation created a top-level skills folder that did not exist when your Claude Code session started, run /reload-skills in that session. Run it again after later changes in that folder.

What's inside

Files included with this skill.

View on GitHub