jax-training-recipe

JAX のトレーニングレシピ

JAX は、構文が異なるだけの PyTorch ではありません。関数は純粋で、状態は明示的に扱われ、変換（jit、vmap、grad、pmap）は、関数の境界が明確な場合にのみ合成できます。

具体的な設定やチェックリストについては、references/training-recipe-details.md を参照してください。

要件の確認

以下を指定します。

• タスク、データセット、損失、評価指標；
• 対象デバイス：CPU、単一 GPU、複数 GPU、TPU；
• モデルフレームワーク：Flax/Linen、Equinox、生の JAX；
• 乱数の用途：ドロップアウト、サンプリング、データ拡張；
• 状態の要件：バッチ正規化、オプティマイザー状態、EMA、チェックポイント。

ワークフロー

1. 実験と最小限のデータバッチを定義します。
2. PRNG キーの所有権と分割方法を計画します。
3. モデルと状態の境界を定義します。
4. モデルとオプティマイザーを初期化します。
5. 純粋な train_step を記述します。
6. jit なしでステップが動作することを確認してから、jit を追加します。
7. 単一デバイスでの動作が正しいことを確認してから、vmap／pmap を追加します。
8. チェックポイントと復元のテストを追加します。
9. 初回実行時の安全性チェックを実施します。
10. 小さなバッチから全体のトレーニングへと段階的に展開します。

JAX 固有のルール

• PRNG キーを再利用しないでください。
• 状態は明示的に扱い、クロージャー内に可変状態を隠さないでください。
• jit 内で Python の副作用を避けてください。
• jit 下での形状変更は、設計上の判断として扱ってください。
• まず jit なしでデバッグし、その後、変換を一度に 1 つずつ再有効化してください。
• 複数デバイスでのトレーニングでは、状態を意図的に複製してください。

よくある失敗モード

• PRNG キーの再利用により、ドロップアウトやサンプリングが繰り返される。
• Python の制御フローや副作用によるトレーサーエラー。
• バッチ正規化／ドロップアウトの状態がパラメーターと分離されていない。
• vmap または pmap 下での形状不一致。
• コンパイル済みループによって NaN が見えなくなる。
• チェックポイントにオプティマイザー状態またはバッチ統計状態が含まれていない。

出力

以下を返します。

• 実験の概要；
• PRNG 戦略；
• モデル／状態の設計；
• トレーニング／評価ステップの計画；
• 最小限の設定；
• 初回実行時のチェック項目；
• JAX 固有の失敗モード；
• スケールアップ計画。