ML Training Recipes

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

references/biomedical.md

references/scaling-and-selection.md

Scaling Laws

Chinchilla rule (Hoffmann et al., 2022)

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

lr * (d_model / 768)^(-0.5)

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

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

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

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

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

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)

references/experiment-loop.md

Evaluation metrics by domain