【ITニュース解説】What we learned fine-tuning our own coding model on a $100 budget

ElderAIという小さなチームが、わずか100ドルの予算で、「ATLAS Code」という独自のコーディングモデルのファインチューニングに取り組んでいる。ファインチューニングとは、既存の大きなAIモデルを、特定の目的に合わせて微調整する作業のことだ。彼らはレンタルした高性能な計算装置であるGPUを使い、この挑戦を通して多くの貴重な教訓を得た。システムエンジニアを目指す初心者にとっても、これは実践的な学びとなるだろう。

まず、彼らが学んだ最も重要な教訓の一つは、「結果を見る前に評価基準（ゲート）を明確に設定する」ことだ。限られた予算で開発していると、少しでも良い結果が出るとそれを「成功」と見なしたくなる誘惑に駆られる。しかし、彼らはこれに抗い、モデルの訓練を開始する前に、どのような状態になれば成功かを定義した小さな「ゲートファイル」を作成した。このファイルは変更できないように管理され、訓練の成否を客観的に判断するための厳格な基準となる。例えば、彼らの最新の訓練では、モデルが元のモデルよりもバイト単位で正確なファイルを生成できるか、一般的なコーディングテストで性能が大きく低下しないか、そしてツール呼び出しのフォーマットが97%以上正しく解析できるか、といった三つの基準が設けられた。これらの基準は、常に同じ環境と方法で元のモデルと比較されるため、公平性が保たれる。この厳格なゲートのおかげで、彼らは都合の良い解釈に惑わされることなく、本当にモデルが改善したかを判断できるようになった。

次に、彼らは「ツール呼び出し能力」と「一般的なコーディング能力」の間にトレードオフ（一方を改善するともう一方が犠牲になる関係）があることを発見した。ATLAS Codeは、コードを編集するためのツールを呼び出すことで、ユーザーのコーディングを支援する。そのため、モデルにツールを使うためのデータを多く与えて訓練したところ、ツール呼び出しのフォーマットは確かに改善した。しかし、その代償として、純粋なコードを生成する能力が少し低下してしまったのだ。これは、モデルがツールを使うことに特化しすぎた結果、汎用的なコーディングスキルが疎かになったためと考えられる。この問題を解決するために、彼らはいくつかの工夫を凝らした。一つは「リハーサルデータ」として、純粋なコード生成の例を訓練データに混ぜることだ。これにより、モデルがコードを書くスキルを忘れないようにする。二つ目は「穏やかな更新」を行うこと。これは、LoRAアダプタという、モデル全体を再訓練せず一部だけを効率的に微調整する技術を使ったり、学習率（モデルが訓練データをどれだけ早く学習するかを示す値）を低く設定したりして、モデルの能力が大きく変動するのを防ぐ方法だ。そして最も効果的だったのは、「40%での中間チェック」である。訓練が40%まで進んだ時点で一度モデルの評価を行い、もし一般的なコーディング能力がすでに低下していれば、その時点で訓練を中止する。これにより、無駄な計算コストを大幅に削減できるようになった。

さらに、コードの「編集」機能の評価方法についても、彼らは深い洞察を得た。当初、彼らはモデルが生成したコードが、最終的な目標のファイルと「バイト単位で完全に一致するか」という非常に厳格な基準で評価していた。しかし、この基準ではほとんどのモデルが失敗してしまうことがわかった。原因を詳しく調べると、実際の開発現場でのコミットメッセージは曖昧な指示が多く、モデルが具体的な変更内容を正確に推測するのは困難だった。また、実際のコミットには、目的の変更とは関係ない空白文字の修正などが含まれていることもあった。モデル自身も、不適切な変更を提案したり、変更対象の文字列がファイル内で複数箇所に存在するために混乱したりする問題があった。これらの発見から、彼らは評価方法を根本的に見直した。新しい評価方法では、「edit_applies」という、モデルの編集指示がファイル内で一意にマッチし、実際にファイルが変更されたかを機械的にチェックする基準を導入した。また、「空白文字正規化後の完全一致」という基準も設け、行末の空白や空行は無視するが、インデントはPythonなどの言語では重要であるため考慮するようにした。さらに、「正確な指示」を与えるテストセットも用意し、モデルに厳密な変更指示を与えた場合の精度を測れるようにした。加えて、モデルが編集ツールに対して複数回呼び出しを試み、その間にエラーが返されても、それを学習して最終的に正しいファイルにたどり着けば評価するという、実際の開発の流れに近い評価も取り入れた。

彼らが遭遇した、頻繁に発生する二つの小さな失敗モードも興味深い。一つは、JSON形式の文字列の中に生のタブ文字が混入してしまう問題だ。モデルがツール呼び出しの引数としてタブでインデントされたコードを引用する際に、そのままタブ文字を書き込んでしまい、それがJSONの仕様違反となることが多かった。この問題が、ツール呼び出し解析の失敗の大部分を占めていたという。彼らは訓練データ側での修正を検討しており、タブでインデントされたファイルやバックスラッシュを多く含むファイルを訓練データに多く含めたり、「間違った呼び出し→実際のエラー→修正された呼び出し」という一連の例を追加したりして、モデルが正しいフォーマットを学習するように促す計画だ。もう一つは、「old_str」の非ユニーク性である。これは、モデルが変更したいコードの断片（old_str）として、ファイル内に複数回出現する文字列を選んでしまい、編集ツールが正しく動作を拒否するという問題だ。この問題に対しては、訓練時に使用するコードの断片が、呼び出し時点でファイル内で一意となるような、最小の完全な行の断片を使用するように訓練データを作成し、すべての訓練データが目標とするファイルを正確に再現できるようにチェックしている。

最後に、限られた予算で効率的に開発を進めるための「コスト管理と監視システム（ウォッチドッグ）」についても言及している。モデルの訓練には、GPUのような高価な計算資源が必要なため、コストが膨らむリスクが常にある。そのため、彼らはそれぞれの訓練実行に対して、複数の監視メカニズムを設けている。例えば、訓練にかかる費用の上限を厳しく設定したり、予測されるコストが上限を超えそうになったら停止させたり、訓練が特定の時間を超えたら強制終了させたりする。さらに、ログの更新が止まったり、GPUが使われずにアイドル状態になったり、訓練中に損失計算がNaN（非数値、計算エラー）になったりした場合も、自動的に訓練を停止させる仕組みが導入されている。訓練終了後には、使用した計算機が確実に削除されたかどうかも確認している。これらのウォッチドッグは非常に有効だが、時には厳しすぎることもある。彼らは、コスト予測の誤差で、健康な訓練が停止してしまった経験から、予測の上限に少し余裕を持たせるように改善した。彼らの最大のコスト削減策は、前述の「40%での中間チェック」で、これにより、もしモデルが期待通りに学習していない場合は、約1.40ドル程度の費用で早期に訓練を中止し、無駄な計算コストを大幅に節約できる。これまでの彼らの訓練にかかったGPU時間は合計で約15ドル程度であり、いかに低予算で効率的な開発を試みているかがわかるだろう。

ElderAIチームは、これらの教訓を活かし、現在も次のステップへと進んでいる。彼らは、より正確な指示を用いた編集データと、事前に設定された厳格な評価基準で、もう一度モデルの訓練を行う予定だ。もしこの訓練でモデルが品質基準をクリアすれば、ATLAS Codeに新しいファインチューニングが適用され、その成果が公開されるだろう。もし失敗したとしても、彼らはその失敗から何を学んだかを正直に共有するとしている。このElderAIの挑戦は、限られたリソースで試行錯誤し、具体的な課題に直面しながら、どのように解決策を見つけていくかという、システムエンジニアを目指す皆さんにとって、非常に実践的な学びの機会となるだろう。