ML Training Recipes

What it does

When your agent uses it

• Choosing an architecture and training recipe for a new dataset
• Debugging loss spikes, out-of-memory errors or slow GPU throughput
• Estimating a token budget from model size using scaling laws
• Setting up AdamW or Muon with a learning-rate schedule

Example prompts

• “Pick an architecture for a small labeled medical imaging dataset and outline the training recipe.”
• “My loss spikes after a few thousand steps, so help me find the cause.”
• “Set up Muon plus AdamW with warmup and decay for my transformer.”
• “How many tokens should a 1B model train on to be compute-optimal?”

Requirements

• Python with PyTorch
• A GPU for the training examples

Workflow steps

6 steps, taken from the first numbered list in SKILL.md.

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

What it can do on your machine

Read from SKILL.md and the folder at commit 773a529. It shows what the files ask for, not the result of running them.

• Tool permissions

  Pre-approves nothing: there is no allowed-tools line, so your agent's usual permission prompts apply.

  From allowed-tools in the SKILL.md frontmatter.

• Runs code

  No scripts in the folder and no shell commands in SKILL.md (its code samples are python).

  From the folder's file list and the shell code blocks in SKILL.md.

• Network

  No URLs in SKILL.md.

  From URLs in SKILL.md, links to its own repository left out.

• Credentials

  Names no API keys, tokens, secrets or passwords.

  From names ending in _API_KEY, _TOKEN, _SECRET, _KEY or _PASSWORD in SKILL.md.

Context cost

ML Training Recipes loads about 2.8k tokens when it runs, and up to ~24k if it reads all its reference files. Until then it costs about 109 tokens; SKILL.md has 931 words of instructions outside code blocks.

Always · name and description, kept in context so the agent knows when to use it

~109

When it runs · the whole SKILL.md, loaded when a task matches

~2.8k

With references · SKILL.md plus every file in references/, read only if the agent opens them

~24k

Estimates: characters ÷ 4, the usual rule of thumb; real counts depend on the model's tokenizer. Scripts and assets cost tokens only if the agent reads them.

Safety

Automated static check — not a guarantee. Review scripts before installing. It scans the text of SKILL.md for risky patterns (piping downloads into a shell, reading credential files, hidden Unicode, destructive commands); files beside SKILL.md are not scanned.

SKILL.md

The full file from Orchestra-Research/AI-Research-SKILLs at commit 773a529, republished under its MIT licence (© Orchestra-Research). 931 words, ~2,772 tokens.

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

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

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

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

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)

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

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

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)

© Orchestra-Research, MIT. Rendered from Markdown: HTML in the file is shown as text, images as links, and headings moved down two levels. Raw file

Files

SKILL.md and 6 other files (references) in 10-optimization/ml-training-recipes of Orchestra-Research/AI-Research-SKILLs.

• SKILL.md
• references/architecture.md
• references/biomedical.md
• references/domain-specific.md
• references/experiment-loop.md
• references/optimizers.md
• references/scaling-and-selection.md

Open the folder on GitHubat commit 773a529

Used in 1 other repository

We found 1 copy of this SKILL.md (exact, near-identical or edited) in other folders, from 1 other GitHub owner. This page covers the copy in Orchestra-Research/AI-Research-SKILLs, which our catalogue first saw on October 7, 2026.

• OpenRaiser/NanoResearchskills/vendor-ai-research/ml-training-recipes