ml-training-recipes

ML 训练方法

适用于各领域 PyTorch 训练的实战模式。内容源自生产级代码库（Karpathy 的 autoresearch/nanochat、torchvision、HuggingFace）和现代训练实践。

参考文件（按需阅读）

• references/architecture.md — Transformer/LLM 架构代码模式、权重初始化
• references/optimizers.md — Muon、AdamW 混合优化器、按组设置学习率、编译优化器步骤
• references/domain-specific.md — 视觉、扩散、对比学习、分布式、检查点、数据加载
• references/scaling-and-selection.md — Scaling laws、计算预算表、决策树、DGX Spark
• references/biomedical.md — 药物发现、蛋白质模型、医学影像、基因组学、临床 NLP
• references/experiment-loop.md — 自主实验循环（autoresearch 保留/丢弃/回退）

---

架构选择

根据数据类型和数据规模选择合适的模型：

数据类型	< 10K 个样本	10K-100K	> 100K
图像	预训练 CNN + 微调	微调 ViT 或 CNN	从头训练 ViT
文本（生成）	少样本提示	微调 GPT/LLaMA（LoRA）	从头预训练
表格数据	XGBoost/LightGBM	仍然使用 XGBoost	可使用神经网络
音频	预训练 Whisper	微调 AST	从头训练
分子	预训练 GNN	微调分子 LM	从头训练 GNN
蛋白质	ESM-2 嵌入 + 预测头	微调 ESM-2	训练蛋白质 LM
医学影像	预训练 CNN	nnU-Net（自动配置）	Swin-UNETR / MedSAM

核心原则：在计算量相同时，训练方案的重要性高于架构。调优良好的 ResNet 优于调优不佳的 ViT（参考：“ResNet Strikes Back”，Wightman，2021）。

生物医学领域请参阅 references/biomedical.md。 序列模型选择和计算规划请参阅 references/scaling-and-selection.md。

---

Scaling Laws

Chinchilla 规则（Hoffmann 等，2022）

计算最优训练：每个参数约对应 20 个 token。

模型规模	计算最优	推理最优（100×）
125M	2.5B 个 token	12.5B 个 token
1B	20B 个 token	100B 个 token
7B	140B 个 token	700B 个 token

FLOPs ≈ 6 × N × D（N=参数，D=token）。数据重复上限：约 4 个 epoch 后收益递减。

---

训练循环

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

核心原则

• 梯度裁剪：clip_grad_norm_(params, 1.0) — 几乎是 Transformer 的通用做法。 例外：Muon 优化器通过正交化来归一化更新，因此裁剪可选。
• Tensor Core 对齐：批次大小和隐藏维度应为 8（bf16）或 64（A100）的倍数。
• 基于时间的预算可让不同硬件上的实验具有可比性。
• 对固定尺寸的视觉输入，使用 cudnn.benchmark = True。

---

优化器配置

现代 LLM 训练会针对不同参数组使用不同优化器：

参数类型	优化器	学习率（基础值）	权重衰减
2D 权重矩阵	Muon	0.04	0.2
Token 嵌入	AdamW	0.6 × scale	0.0
反嵌入（lm_head）	AdamW	0.004 × scale	0.0
每层标量	AdamW	0.005 × scale	0.0

按维度缩放学习率：lr * (d_model / 768)^(-0.5) — 确保不同规模下的动态稳定。

经验法则

• 嵌入需要较高的学习率（更新稀疏）。切勿对嵌入使用权重衰减。
• 权重衰减调度：在训练过程中将 WD 线性衰减至 0。
• AdamW 默认值：β1=0.9、β2=0.95、eps=1e-10（而不是默认的 1e-8——可避免 bf16 中的过时更新）。

Muon 的详细信息（polar express 正交化、NorMuon）请参阅 references/optimizers.md。

---

学习率调度

基于时间（autoresearch 风格）

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

WSD（预热-稳定-衰减）：正逐渐普及，更便于在训练过程中途恢复训练。

指导原则

• 预热：占训练量的 1-5%。使用 Muon 时可不预热（autoresearch 使用 WARMUP_RATIO=0.0）。
• 预热衰减：将训练量的 30-50% 用于学习率衰减。与预热相比，这对最终质量更重要。
• 最终学习率：设为 0 或峰值的约 10%。设为 0 更简单。

---

混合精度与编译

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 及更新架构）：指数范围与 fp32 相同，无需损失缩放。优先于 fp16。
• fp16：需要 GradScaler。仅在 V100 或更旧的 GPU 上使用。
• dynamic=False 可启用最大优化。若没有图中断，可添加 fullgraph=True。
• 开始的几步会比较慢（JIT）——计时应排除这些步骤。

---

内存与性能

Meta device 初始化（大型模型）

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

MFU（模型 FLOPs 利用率）

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

理想目标：>30% 尚可，>40% 良好，>50% 出色（单 GPU）。

OOM 解决方案（按顺序尝试）

1. 减小 DEVICE_BATCH_SIZE，增大 grad_accum_steps
2. PYTORCH_ALLOC_CONF=expandable_segments:True
3. model.zero_grad(set_to_none=True)
4. Meta device 初始化 → to_empty
5. 激活检查点：torch.utils.checkpoint.checkpoint()
6. 8 位优化器（bitsandbytes）：优化器状态可节省约 30% 内存

---

超参数搜索

优先级顺序（从先调到后调）

1. 学习率 — 影响最大。始终优先调节。
2. 批大小 — 选择能容纳的最大值。它影响速度，不影响质量。
3. 权重衰减 — AdamW 使用 0.01-0.1。
4. 预热步数 — 占训练量的 1-5%。

2025 年默认方案

设置	值
优化器	AdamW (β1=0.9, β2=0.95, eps=1e-10)
权重衰减	0.1
学习率调度	余弦衰减或 WSD
峰值学习率	3e-4（更大的模型应相应调低）
精度	bf16
梯度裁剪	max_norm=1.0
归一化	RMSNorm（预归一化）
激活函数	SwiGLU
位置编码	RoPE
注意力	Flash Attention，可选用 GQA

---

调试清单

Karpathy 的方案（仍是权威做法）

1. 深入了解数据 — 可视化数据、检查分布、验证标签
2. 先确保端到端流程可运行 — 用一个简单案例验证
3. 让模型过拟合一个批次 — 如果做不到，说明存在 bug
4. 然后再正则化 — 只有过拟合成功后才添加正则化
5. 调节超参数 — 从已知的默认值开始

损失爆炸 / NaN

1. 降低学习率（缩小 3-10 倍）
2. 添加梯度裁剪：clip_grad_norm_(params, 1.0)
3. 检查输入中是否存在 inf/nan
4. 添加 logit 软封顶：softcap * tanh(logits / softcap)
5. 在注意力中添加 QK-norm
6. 验证权重初始化（输出投影是否应零初始化？）
7. 检查梯度累积时的损失缩减（loss / grad_accum_steps）

训练缓慢 / MFU 低

1. 确认 torch.compile 已启用
2. 检查 torch.set_float32_matmul_precision("high")
3. 固定内存 + 非阻塞传输
4. 使用 torch.profiler 分析性能
5. 是否有 GC 停顿？gc.freeze(); gc.disable()
6. Tensor Core 对齐：维度为 8/64 的倍数

损失停滞 / 收敛缓慢

1. 学习率太低 — 尝试增大 2-5 倍
2. 预热时间过长
3. 权重衰减过高
4. 确认学习率调度确实已应用（打印每一步的学习率）
5. 模型对于任务来说太小

隐性故障

1. 训练集与验证集之间存在数据泄漏
2. 推理时预处理错误 — 数据增强不匹配
3. 标签错误 — 使用 cleanlab 检测
4. 打乱数据的 bug — 批次之间存在相关性
5. 与预训练模型的分词器不匹配

需要监控的内容

• 梯度范数 — 梯度尖峰通常先于损失尖峰出现
• 各层激活统计信息 — 可揭示梯度爆炸或消失
• 死亡神经元 — ReLU 输出为零的比例 >50% 表示 ReLU 神经元死亡问题
• 学习率 — 验证调度是否已应用（常见的隐性 bug）

---

实验管理

使用 TSV 跟踪实验，便于比较：

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)

简单性准则：其他条件相同时，越简单越好。移除某项后结果不变，就是很好的结果。如需了解系统化的 Agent 驱动实验，请参阅 references/experiment-loop.md。

按领域划分的评估指标

领域	主要指标	备注
大模型	BPB（每字节位数）	不受词表大小影响
分类	准确率 / F1	对于类别不平衡的数据，使用 Macro-F1
分割	mIoU / Dice	按类别计算的 IoU 可揭示薄弱环节
生成	FID	需要超过 10k 个样本
回归	RMSE / MAE	对偏斜目标进行对数变换