コードとしてのモデルトレーニング

研究 Michael Barlow 2026/05/22 コードとしてのモデルトレーニング TL;DR: モデルトレーニングは非常に複雑になり、多くの専門的なステージとチームが必要とされており、それらの間の手作業による調整ではスケールしない。Aleph Alphaでは、トレーニングパイプライン全体をコードで実装し、モデルトレーニングを協調的なソフトウェアプロジェクトに変えるモデルファクトリーSavannaを構築した。Savannaでは、エンドツーエンドのトレーニング実行は密閉型であり、ワンクリックで起動できる。この記事では、Savanna、それが必要とされる理由、そしてその有効性を可能にするエンジニアリング文化について説明する。 はじめに モデルトレーニングは急速に進展している。新しいステージがパイプラインに次々と加わり、既存のものはより複雑になり、モデルトレーニングは3つの重要な理由からエンジニアリングの課題となっている。第一に、複雑さが増すということはエラーの余地が増えるということだ。データ、コード、または設定のバグや不整合は、トレーニング実行全体を失敗させたり、発散させたりする可能性がある。第二に、失敗のコストが増加し続けている。モデルは大規模化し、GPU価格は上昇し、1回の実行あたりに処理されるデータ量も増え続けている。何千ものGPU時間を消費しているときに、「おっと」というのは高価な言葉だ。第三の、そして最も困難な課題は組織的なものだ。その複雑さは単一の頭脳の能力をはるかに超えており、そのため私たちのラボのような場所では、大規模で専門化されたチームを構築してきた。問題は、それらのチームを調整することになる。メンバーは自分の専門分野の最新の研究を自律的に探索しながら、変更を本番パイプラインに統合し、それを壊したり、互いの作業を妨害したりすることなく行うにはどうすればよいのか？そして、個々のステージの改善が最終的に優れたモデルにつながることをどのように保証するのか？従来の、手動によるモデルトレーニングプロセスでは、これらの問いに対する良い答えがない。 手動によるモデルトレーニングの隠れたコスト 従来の、手動によるプロセスを分解してみよう。大まかに言えば、モデルトレーニングはシンプルに見える。インターネットを吸収するための事前トレーニングの後、指示に従うことを学習するための事後トレーニングが行われる。実際には、最初の試行で良いモデルをトレーニングすることはできない。良いデータミックス、アーキテクチャ、トレーニングレシピに到達するには、評価に導かれる反復的で計算量の多いプロセスであり、以下の各矢印は相対的なGPUコストによって重み付けされている。 これらの各コンポーネントは、複数の専門チームを必要とするほど複雑だ。例えば、最新の事後トレーニングは、教師ありファインチューニング（SFT）ステージとそれに続く強化学習（RL）ステージで構成される。SFTとRLは異なるスキルセットとツールを必要とするが、モデルをトレーニングするためには統合されなければならない。手動のラボで、1つのモデルがパイプラインを通過する道のりがどのようなものかを考えてみよう。 データチームは新しいミックスを完成させ、データベースパスをSlackで事前トレーニングチームに送り、事前トレーニングチームは数週間にわたる実行を開始する。2週間後、ストレージクォータがいっぱいになり、トレーニング実行がクラッシュする。ファイルシステムは手動で管理されているため、「do_not_delete」というファイル名を持つ30TBのデータセットを削除しても安全かどうか誰も確信できず、事前トレーニングチームがこれを解決する間、GPUはアイドル状態になる。ようやく再起動する際、彼らは記憶とSlackのスレッドから元の設定を再構築し、フラグを設定し忘れていないことを願う。これが最初の隠れたコストだ。すべての手動ステップはヒューマンエラーの機会となる。 事前トレーニングが最終的に完了すると、事前トレーニングチームはチェックポイントをSFTチームに引き渡す。良いレシピを見つけるために、SFTチームは異なる設定とデータミックスで並列トレーニングのスイープを手動で開始する。チェックポイントが届くと、チームはそれぞれの評価スクリプトを実行し、結果と分析をSlackで共有する。いくつかのレシピは有望に見え、他のものはそうではないため、数週間このプロセスを繰り返し、良いものに絞り込む。気付かないうちに、チームは数ヶ月前の以前の事前トレーニング済みチェックポイントのためにすでに完了したいくつかの実験を繰り返してしまった。これが第二の隠れたコストだ。チームはその学習を忘れてしまう。ハイパーパラメータの現在の値の背後にある推論の永続的な記録はなく、データミックスとその構成データセットとの間に正式なリンクもなく、モデルとそれを生成したトレーニングレシピとの間に明確な帰属もない。手動のラボでは、この系譜はSlack、ファイルシステム、実験マネージャー、さまざまなWikiページに散らばっており、時間の経過とともに容易に失われる。 SFTチームはチェックポイントをRLチームに引き渡し、RLチームはこれをベースとしてトレーニング実行を開始する。最終的なモデルのパフォーマンスが低い。RLレシピは先月のSFTチェックポイントに過剰適合しているのか、それともSFTチェックポイント自体に問題があるのか？2週間のデバッグの後、RLチームは後者であることを確認する。どちらのチームも相手のステージを実行できないため、各チームはパイプラインの自分たちの部分ではなく、最終的なモデル全体のために最適化していた。そして、統合が手動の引き渡しであるため、めったに行われず、毎回1ヶ月分の乖離を調整することになる。これが第三の隠れたコストだ。手動で、頻繁でない引き渡しはチームの所有権を分断する。手動のモデルトレーニングがスケールしないことは明らかであり、自動化だけが不足しているわけではない。これらの隠れたコストはすべて、同じ根本的な問題に起因している。パイプラインがチームの心の中にあるのではなく、共有され、永続的なアーティファクトの中に存在しないことだ。スケールするには、パイプライン自体がチームが協力できる扱いやすいものである必要がある。 紹介：コードとしてのモデルトレーニング 私たちのモデルファクトリー、コードネームSavannaは、モデルトレーニングパイプラインとプロセス全体を命令型コードで実装している。これをコードとしてのモデルトレーニング（MTaC）と呼ぶ。Savannaのシンプルな事後トレーニングパイプラインは擬似コードで次のようになる。 ```python async post_train(config: PostTrainConfig) -> PostTrainEvaluation: sft_checkpoint = await sft(config.sft) sft_eval = spawn evaluate(config.eval, sft_checkpoint) rl_checkpoint = await rl(config.rl, sft_checkpoint) rl_eval = spawn evaluate(config.eval, rl_checkpoint) return PostTrainEvaluation(await sft_eval, await rl_eval) ``` パイプラインをコードにすることで、3つのものが得られる。構成可能性、合意、そして来歴だ。 構成可能性: 手動のステップを型付きの入力と出力を持つ関数として表現することで、抽象化を構築し、それらを組み合わせてワンクリックで実行されるエンドツーエンドのパイプラインを作成できることから得られる。パイプラインの変更は関数の編集と同じくらい簡単になり、中間チェックポイントの評価のような反復作業はforループで自動化でき、異なるパラメータ化でパイプラインのサブセットやスケールダウンしたバージョンを実行できるため、テストも簡単になる。 合意: バージョン管理から得られる。メインブランチは、モデルをトレーニングする方法に関するチームの集合的な最良の理解を表す。コードには完全なトレーニングレシピが含まれているため、トレーニング実行を起動する際に再構築するセットアップや忘れるべきフラグはない。 来歴: コードコメントとコミット履歴から得られる。これらは、メインに至る学習と決定をエンコードする。過去のトレーニング実行は再現可能であり続ける。なぜなら、それらを生成したコードは、チェックアウトして再実行できるコミットに固定されているからだ。どのチームでもフルパイプラインを自分で起動できるとき、彼らは他のチームのステージを実行し、自分の部分だけでなくモデル全体を反復的に改善できる。 ラボは通常、スケールするときにモデルトレーニングを時間的に分解し、各チームがパイプラインの1つのステージを所有する。MTaCは、チームが代わりに多言語性などのモデルの振る舞いをエンドツーエンドで所有する、能力ベースの分解の扉を開く。 少しずつ統合し、頻繁に統合する パイプラインがコードにある場合、標準的なコードコラボレーションのベストプラクティスが適用され、MTaCを最大限に活用するにはそれらに従う必要がある。最も重要なのはトランクベース開発であり、変更は可能な限り早く、小さな増分でメインに投入されるため、チームはできるだけ早く互いの作業に基づいて構築し、アプローチが間違っている場合はすぐに失敗できる。代わりに、長期間生きるブランチに変更を蓄積すると、以前と同じ統合債務を支払うことになる。Savannaの使用 SavannaはGitHub上に存在し、そのCIがモデルトレーニングのエントリポイントとなる。