Claude CodeでPyTorchの再現性を固定する — seedと決定的アルゴリズム

PyTorchの学習を同じ設定で2回回すと、損失や精度が少しずつ違う。この「ずれ」は、乱数の種(seed)、アルゴリズムの選び方、DataLoaderのワーカーという3系統に分けると、どこを固定すればよいかが見えてきます。Claude Codeに任せる場合も、いきなり「再現性を直して」と頼むより、ずれを測るテストを先に書かせるほうが確実です。

この記事では、同じ設定で2回学習して結果を比べるテストを軸に、torch.manual_seed・torch.use_deterministic_algorithms・worker_init_fn を1つずつ足して差が縮むかを見る手順を示します。

ずれの原因は3系統に分けて考える

PyTorchのReproducibilityのページは、再現性の確保を「乱数源の制御」と「非決定的アルゴリズムの回避」の2段で説明しています。ここにDataLoaderの話を加えると、次の3系統になります。

先に断っておくと、完全な再現は保証されていません。PyTorchのリリース間、コミット間、プラットフォーム間では、同じseedでも結果が一致するとは限りません。CPUとGPUの間でも同様です。固定の目標は、解析パイプライン全体の再現性(Snakemakeの再現性の確認のような話)ではなく、学習1回分に絞ります。具体的には「同じマシン・同じPyTorchのバージョンで、2回の実行が一致すること」です。

手順1: ずれを測るテストを先に書かせる

修正の前に、ずれを数値で見えるようにします。乱数と日時の固定をテストの前提に置く考え方は、Claude Codeで特性テストを書くと同じです。Claude Codeへの指示は、直したい動作ではなく、失敗するテストの形で渡します。

学習スクリプト train.py を同じ設定で2回実行し、最終損失と全パラメータが
一致するかを確かめる pytest を tests/test_repro.py に書いてください。
- 学習は CPU・小さなモデル・20ステップ程度に縮める
- 比較は torch.equal(完全一致)。許容誤差で逃げない
- まだ train.py は変更しない。テストが落ちることを確認して報告する

比較に許容誤差を入れないのが要点です。誤差を許すと、小さなずれが見えなくなり、固定が効いたかどうかも判定できません。

テストの骨格は次のような形になります。モデルとデータは説明用の最小例です。

import torch
from torch.utils.data import DataLoader, TensorDataset
 
def train_once(seed: int) -> torch.Tensor:
    torch.manual_seed(seed)
    x = torch.randn(256, 8)
    y = torch.randn(256, 1)
    loader = DataLoader(TensorDataset(x, y), batch_size=32,
                        shuffle=True, num_workers=2)
    model = torch.nn.Linear(8, 1)
    opt = torch.optim.SGD(model.parameters(), lr=0.1)
    for _ in range(2):
        for xb, yb in loader:
            opt.zero_grad()
            torch.nn.functional.mse_loss(model(xb), yb).backward()
            opt.step()
    return torch.cat([p.detach().flatten() for p in model.parameters()])
 
def test_two_runs_match():
    assert torch.equal(train_once(0), train_once(0))
 
def test_different_seed_differs():
    assert not torch.equal(train_once(0), train_once(1))

2つ目のテストは、固定しすぎて何も変わらなくなっていないかを見る保険です。seedを変えたら結果が変わることまで確かめておきます。

手順2: seedを1か所にまとめる

乱数の種は、PyTorch・Python標準・NumPyの3つが別々の生成器を持っています。torch.manual_seed() はCPUとCUDAの全デバイスの乱数生成器の種を決めます。独自演算子などでは random.seed() も必要で、NumPyのグローバル生成器は np.random.seed() で固定します。NumPyの Generator オブジェクトを使うコードは、グローバルの種では固定されない点にも注意が必要です。

import random
import numpy as np
import torch
 
def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

Claude Codeには「学習の入口で一度だけ seed_everything を呼ぶ。スクリプト内に散らばった manual_seed は整理する」と依頼します。散らばったseedは、後から呼ばれたものが先のものを上書きし、何が効いているのか分からなくなる原因になります。

手順3: 決定的アルゴリズムを有効にする

seedを固定してもずれる場合は、アルゴリズムの側を疑います。torch.use_deterministic_algorithms(True) を呼ぶと、PyTorchは決定的な実装がある演算ではそちらを使い、決定的な実装が無い演算では RuntimeError を投げます。

torch.use_deterministic_algorithms(True)
torch.backends.cudnn.benchmark = False

cudnn.benchmark は別の設定です。cuDNNの畳み込みはサイズごとに複数のアルゴリズムを試して最速のものを選ぶ機能を持ち、ベンチマークのノイズや環境の違いで実行ごとに選ばれるアルゴリズムが変わりえます。False にすると選択が決定的になりますが、性能が落ちる可能性があります。

もう1つ、torch.backends.cudnn.deterministic = True は畳み込みの挙動だけを制御する設定です。use_deterministic_algorithms は畳み込み以外の演算にも効くので、迷ったら後者を使います。

起きるRuntimeErrorの読み方

非決定的な演算を含むモデルでは、有効化した時点でエラーになることがあります。ドキュメントの例では、GPU上の index_add_ が次の形で失敗します。

RuntimeError: index_add_cuda_ does not have a deterministic implementation,
but you set 'torch.use_deterministic_algorithms(True)'. ...

この手のエラーは、固定が足りないのではなく、どの演算が非決定的かを教えてくれています。Claude Codeには、エラー全文を渡して「どの層・どの演算が原因か」を特定させ、決定的な代替(別の演算への置き換えなど)を提案させます。置き換えが難しい場合の逃げ道が warn_only=True です。エラーの代わりに警告を出して先へ進みますが、その演算の非決定性は残ります。

手順4: DataLoaderのワーカーを固定する

num_workers を1以上にすると、データ読み込みは別プロセスで動きます。DataLoaderは「マルチプロセスのデータ読み込みの乱数」のアルゴリズムに従ってワーカーの種を決め直します。再現性を保つ手段は worker_init_fn と generator です。例は次のとおりです。

import random
import numpy as np
import torch
from torch.utils.data import DataLoader
 
def seed_worker(worker_id):
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)
 
g = torch.Generator()
g.manual_seed(0)
 
loader = DataLoader(
    train_dataset,
    batch_size=batch_size,
    num_workers=num_workers,
    worker_init_fn=seed_worker,
    generator=g,
)

ワーカー内でNumPyやPython標準の乱数を使うデータ拡張(ランダムクロップなど)を書いているなら、この2行が効きます。ワーカー内でtorchの乱数だけを使うなら、generator の固定が主役になります。どちらに当たるかは、データセットの __getitem__ を読めば分かります。

手順5: 1つずつ足して差が縮むかを見る

ここまでの3つを一度に入れると、どれが効いたのか分かりません。テストを「落ちる状態」から始めて、次の順で足し、毎回 pytest を回します。

逆向きの確認も有効です。緑になった状態から固定を1つ外し、テストが赤に戻るなら、その固定は実際に効いています。外しても緑のままの固定は、その学習では不要か、テストの規模が小さくてずれを拾えていないかのどちらかです。

Claude Codeに毎回守らせる

固定の方針は、セッションをまたいで効かせたいので、CLAUDE.mdに短く書いておきます。CLAUDE.mdは毎セッションの開始時に読み込まれ、1ファイル200行以内が目安とされています。学習コード専用の数行にとどめます。

## 学習コードの再現性
- 学習の入口で seed_everything(seed) を1回だけ呼ぶ。manual_seed を散らさない
- torch.use_deterministic_algorithms(True) を前提にする。RuntimeError は握りつぶさず原因の演算を報告する
- DataLoader に num_workers>=1 を指定するときは worker_init_fn と generator を付ける
- 変更後は pytest tests/test_repro.py を実行し、結果を報告する

CLAUDE.mdは助言であって強制ではありません。確実に走らせたいテストは、フックに載せます。次はClaudeが応答を終える時点で再現性テストを走らせ、落ちていたら続行させるStopフックの例です。

{
  "hooks": {
    "Stop": [
      {
        "hooks": [
          { "type": "command", "command": ".claude/hooks/check-repro.sh" }
        ]
      }
    ]
  }
}

#!/usr/bin/env bash
# .claude/hooks/check-repro.sh
input=$(cat)
# すでにフックで続行中なら、無限ループを避けて止める
if [ "$(echo "$input" | jq -r '.stop_hook_active')" = "true" ]; then
  exit 0
fi
if ! out=$(python -m pytest tests/test_repro.py -x -q 2>&1); then
  echo "$out" | tail -20 >&2
  exit 2
fi

Stopフックは終了コード2で、Claudeを止めずに作業へ戻します。stderrの内容が、続行すべき理由としてClaudeに渡ります。stop_hook_active は、すでにStopフックの結果で続行している最中に true になる値で、確認しないと解消しない条件でブロックし続けかねません。連続8回の続行で上限がかかる保護もあります。

つまずきやすい点

GPUで固定したのにCPUと結果が合わない

CPUとGPUの実行間でも結果が一致しない場合があります。比較は同じデバイスの中で行います。テストを実GPUに広げる場合も同じで、2回の実行は同じGPU上で比べます。cudnn.deterministic は畳み込みだけを縛る設定なので、テストの前提には use_deterministic_algorithms を置きます。

決定的にしたら遅くなった

決定的な演算は非決定的な演算より遅いことが多い点は、ドキュメントにも警告があります。開発中やデバッグ・回帰テストでは決定的に、本番の学習では外す、という切り替えを環境変数で用意する手があります。次の形なら、REPRO=1 のときだけ決定的にできます。

import os
 
if os.environ.get("REPRO") == "1":
    torch.use_deterministic_algorithms(True)
    torch.backends.cudnn.benchmark = False

フックと回帰テストの側で REPRO=1 を付けて走らせ、本番の学習コマンドには付けません。テストが見るのは「決定的にした状態で2回が一致するか」であり、本番の速度とは切り離せます。

scaled_dot_product_attentionの逆伝播だけずれる

ドキュメントの表では、FLASH_ATTENTION・EFFICIENT_ATTENTION・CUDNN_ATTENTIONの各バックエンドは逆伝播が非決定的です。use_deterministic_algorithms(True, warn_only=False) で決定的な逆伝播になるのはFlashとEfficientで、cuDNNのバックエンドは無効化されます。warn_only=True では警告が出るだけで、非決定的な経路のまま動きます。固定を目的に warn_only=True を使うと、効いていないのに安心してしまう点に注意が必要です。

初期化されていないメモリを読んでいる

torch.empty などが返す未初期化のメモリを入力に使うと、結果は非決定的になります。use_deterministic_algorithms(True) のとき、既定ではその領域を既知の値で埋める設定(torch.utils.deterministic.fill_uninitialized_memory)が働きます。性能のために切ることもできますが、未初期化メモリを使っていないと確信できるときに限ります。

torch.compileを使っている

use_deterministic_algorithms を有効にするとInductorの決定的モードも自動で有効になり、数値に影響するオンデバイスのベンチマークや自動チューニングを避けます。

まとめ

ずれの調査は、固定の手段を探す前に、ずれを測るテストを置くところから始まります。完全一致で比べるテストを赤にし、seed・決定的アルゴリズム・worker seedを1つずつ足して緑にする。この順なら、どの固定が効いたのかが記録として残ります。学習したモデルの出力を別の形式と突き合わせる段階では、PyTorchモデルをONNXに書き出して出力の一致を検証する手順が続きになります。そのテストをフックに載せておけば、モデルを書き換えた次のセッションでも、再現性が崩れた時点で気づけます。