ml-training-recipes

ML トレーニングレシピ

さまざまな領域で実証済みの、PyTorch トレーニングパターンです。実運用のコードベース （Karpathy の autoresearch/nanochat、torchvision、HuggingFace）と、最新のトレーニング手法を参考にしています。

参考ファイル（必要に応じて参照）

• references/architecture.md — Transformer/LLM のアーキテクチャコードパターン、重みの初期化
• references/optimizers.md — Muon、AdamW のハイブリッド、グループごとの LR、コンパイル済みオプティマイザーステップ
• references/domain-specific.md — ビジョン、拡散、コントラスト学習、分散処理、チェックポイント、データ読み込み
• references/scaling-and-selection.md — スケーリング則、計算予算の表、判断ツリー、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 を参照してください。

---

スケーリング則

Chinchilla の法則（Hoffmann ら、2022 年）

計算量を最適化したトレーニングでは、パラメーターあたり約 20 トークンを使用します。

モデルサイズ	計算量最適	推論最適（100×）
125M	2.5B トークン	12.5B トークン
1B	20B トークン	100B トークン
7B	140B トークン	700B トークン

FLOPs ≈ 6 × N × D（N=パラメーター数、D=トークン数）。データの繰り返しは、効果が薄れ始める約 4 エポックまでに制限します。

---

トレーニングループ

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 トレーニングでは、パラメーターグループごとに異なるオプティマイザーを使います。

パラメーターの種類	オプティマイザー	LR（基本値）	重み減衰
2D 重み行列	Muon	0.04	0.2
トークン埋め込み	AdamW	0.6 × scale	0.0
アン埋め込み（lm_head）	AdamW	0.004 × scale	0.0
層ごとのスカラー	AdamW	0.005 × scale	0.0

次元に応じた LR のスケーリング：lr * (d_model / 768)^(-0.5) — モデルサイズが異なっても学習ダイナミクスを安定させます。

経験則

• 埋め込みにはより高い LR が必要です（更新が疎なため）。埋め込みには重み減衰を適用しないでください。
• 重み減衰のスケジューリング：トレーニング中に 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以前でのみ使用してください。
• dynamic=Falseにすると、最大限の最適化が有効になります。グラフブレークがなければ、fullgraph=Trueを追加してください。
• 最初のステップはJITにより遅くなります。計測時間から除外してください。

---

メモリとパフォーマンス

Metaデバイスによる初期化（大規模モデル）

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デバイスで初期化してから、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（pre-norm）
活性化関数	SwiGLU
位置エンコーディング	RoPE
アテンション	Flash Attention、必要に応じてGQA

---

デバッグチェックリスト

Karpathyのレシピ（現在も標準的）

1. データを徹底的に理解する — 可視化し、分布を確認し、ラベルを検証します
2. まずエンドツーエンドで動作させる — 単純なケースで検証します
3. 1つのバッチに過学習させる — できない場合は、バグがあります
4. その後で正則化する — 過学習が機能してから、正則化を追加します
5. ハイパーパラメーターを調整する — 既知のデフォルト値から始めます

損失の発散／NaN

1. 学習率を下げます（3～10分の1）
2. 勾配クリッピングを追加します：clip_grad_norm_(params, 1.0)
3. 入力にinfやnanが含まれていないか確認します
4. ロジットにソフトキャッピングを追加します：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. メモリを固定し、non_blocking転送を使用します
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. シャッフルのバグ — バッチ間に相関がある
5. 事前学習済みモデルとのトークナイザーの不一致

監視項目

• 勾配ノルム — 損失が急増する前に勾配ノルムが急増します
• レイヤーごとのアクティベーション統計 — 発散や消失を明らかにします
• 非活性ニューロン — ReLUの50％超がゼロの場合は、ReLUの不活性化の問題です
• 学習率 — スケジュールが適用されているか確認します（よくある見落としやすいバグ）

---

実験管理

比較しやすいよう、実験を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)

シンプルさの基準：ほかの条件がすべて同じなら、シンプルなほうが優れています。何かを取り除いて同等の結果が得られれば、それは大きな成果です。エージェント主導の体系的な実験については、references/experiment-loop.mdを参照してください。

ドメイン別の評価指標

ドメイン	主要な指標	備考
LLM	BPB（ビット／バイト）	語彙サイズに依存しません
分類	Accuracy／F1	クラス不均衡の場合はMacro-F1を使用します
セグメンテーション	mIoU／Dice	クラスごとのIoUで弱点を特定できます
生成	FID	10,000件を超えるサンプルが必要です
回帰	RMSE／MAE	偏りのあるターゲットは対数変換します

関連スキル

• serving-llms-vllm：vLLMのPagedAttentionと連続バッチ処理を使用して、高スループットでLLMを提供します。
• llama-factory：LLaMA-Factoryを使用した大規模言語モデルのファインチューニングについて、コード不要のWebUI、100以上のサポート対象モデルのトレーニングなどを解説します…
• peft-fine-tuning：Hugging Face PEFTを使用し、LoRA、QLoRA、IA3、AdaLoRA、プレフィックスチューニング、プロンプトチューニングによって、パラメーターの1％未満を対象にLLMをファインチューニングします…