大規模言語モデルを自社でトレーニング＆活用する方法
オンラインIDEを提供しているReplitでは自社で大規模言語モデルをトレーニングしているらしく、そのノウハウがブログ記事にまとめられていたので要約してみました。
なぜ自社で大規模言語モデルをトレーニングするのか？
企業が独自に大規模言語モデル（以下、LLMs）をトレーニングすることを決める理由は、データのプライバシーやセキュリティから、アップデートや改良のコントロールの強化まで様々なものがあるが、Replit社ではカスタマイズ性、依存度の低減、コスト効率に重点を置いている。
カスタマイズ性
カスタムモデルをトレーニングすることで、GPT-4のような汎用モデルやCodexのようなコードに特化したモデルではカバーしきれないプラットフォーム固有の機能、用語、コンテキストなどといった特定のニーズや要件に合わせてモデルを調整することができる。例えば、Replitで人気の高いJavascript React（JSX）やTypescript React（TSX）など、特定のWebベースの言語に対してより良い仕事をするように、Replit社のモデルはトレーニングされている。
依存度の低減
目の前にタスクに最適なモデルを選択する必要はあるものの、一握りのAIプロバイダーへの依存度を下げることには利点がある。これはReplit社だけでなくより広い開発者コミュニティにとっても同じことなので、Replit社ではモデルの一部をオープンソース化する予定だが、これはモデルを訓練する手段がなければできないことである。
コスト効率
コストは今後も下がり続けると予測できるが、グローバルな開発者コミュニティで使用するには、LLMsはまだ法外に高価である。Replit社では、次の10億人のソフトウェアクリエイターをオンライン化することをミッションとしている。インドの学生が携帯電話でコーディングしていても、シリコンバレーのプロの開発者と同じAIにアクセスできるようにすべきで、これを可能にするために、より小さく、より効率的で、コストを大幅に削減したホスティングが可能なカスタムモデルをトレーニングしている。
データパイプライン
LLMsのトレーニングには膨大な量のデータが必要である。LLMsを学習させるためには高度に最適化された堅牢なデータパイプラインを構築する必要があるが、公共データや占有データの新しいソースを容易に取り込むことができる柔軟性も必要である。
The Stack
まずはHugging Faceで公開されているThe Stackを主要なデータソースとして使用する。The StackはBigCodeプロジェクトによって提供されており、データセット構築の詳細はKocetkov et al.（2022）に掲載されている。重複排除後のデータセットのバージョン1.2には、350以上のプログラミング言語で書かれ、寛容にライセンスされたソースコードが約2.7TB含まれている。
Hugging Faceの提供するTransformersライブラリはモデルトレーニングに関する多くの課題を抽象化するために素晴らしい仕事をしてくれるが、Replit社でのプロセスではデータに対する更なる制御と分散処理能力が必要であるため、そのままでは不十分だと判断し、独自のデータパイプラインを構築している。
データ処理
より高度なデータ処理を行うためにDatabricksを使用してパイプラインを構築している。このアプローチであればReplitやStack Overflowのような追加のデータソースを簡単に導入することも可能だ。
最初のステップではHugging Faceから生データをダウンロードしている。Apache Sparkを使用して、各プログラミング言語間でデータセット構築プロセスを並列化している。そしてデータを再分割し、下流処理のために最適化された設定でparquetフォーマットに書き直している。
次にデータクリーニングと前処理を行う。通常、データの重複排除やエンコーディングの問題を解決することが重要であるが、Kocetkov et al.（2022）で述べられている近似重複除去手法を用いてすでにThe Stackがこれを実行している。ただし、Replitデータをパイプラインに導入し始めると重複除去プロセスを再度実行する必要があるため、Databricksを利用してThe Stack、Stack Overflow、Replitデータを大きなデータレイク内の3つのソースとして扱うことでこれを実現している。
Databricksを使用するもう一つの利点は、基礎データに対してスケーラブルで扱いやすい分析を実行できることである。データソースにおけるあらゆる種類のサマリ統計を実行し、ロングテール分布をチェックし、プロセスにおける問題や矛盾を診断する。この作業はすべてDatabricksのノートブックで行われ、MLFlowと統合することで、分析結果を追跡・再現することができる。このステップはデータの定期的なX線検査に相当し、前処理のための様々なステップに役立つ。
トークン化とボキャブラリートレーニング
トークン化の前に、モデルのトレーニングに使用するのと同じデータのランダムなサブサンプルを使用して、独自のカスタムボキャブラリーをトレーニングする。カスタムボキャブラリーはモデルがコードコンテンツをより良く理解し、生成することを可能にする。その結果、モデルのパフォーマンスが向上し、モデルのトレーニングや推論が高速化する。
このステップは、プロセスの3つのステージ（データパイプライン、モデルトレーニング、推論）すべてで使用されるため、プロセスの中でも最も重要な1つである。トークン化については今後の記事でより深く掘り下げる予定である。
カスタムボキャブラリーを学習させた後にデータをトークン化する。最後にトレーニングデータセットを構築し、モデルのトレーニングプロセスに最適化された共有フォーマットに書き出す。
モデルトレーニング
Replit社ではMosaicMLを使用してモデルをトレーニングしている。理由は以下の3点である。
1）複数のクラウドプロバイダー
Mosaicではアカウント設定やインテグレーションのためのオーバーヘッドなしに、様々なクラウドプロバイダーのGPUを活用することができる。
2）LLMトレーニング設定
Composerライブラリには、様々なモデルのトレーニングや、様々なタイプのトレーニングの目的に対応するための、調整された設定が多数用意されている。
3）マネージドインフラ
マネージドインフラはオーケストレーション、効率の最適化、フォールトトレランス（ノード障害からの回復）を提供してくれる。
モデルのパラメータを決定する際は、モデルサイズ、コンテキストウィンドウ、推論時間、メモリフットプリントなどの間の様々なトレードオフを考慮する。一般に大きなモデルは性能が良く、転移学習が可能だが、学習と推論の両方においてより高い計算量を必要とする。Replitはデスクトップのネイティブアプリケーションのようなパフォーマンスを持つクラウドネイティブIDEであるため、コード補完モデルは非常に高速である必要があり、低レイテンシの推論を行うことができる小さなモデルを選択している。
モデルのパラメータに加えて様々な学習目標を選択することができるが、それぞれ独自のメリット／デメリットがある。最も一般的な学習目標は「次のトークンを予測すること」だが、コード補完には有効なものの、文章内の更に下流の文脈を考慮することができない。この問題は「fill-in-the-middle」という学習目標を用いることで軽減することができる。さらに別のアプローチとして、UL2（Unsupervised Latent Language Learning）がある。これは与えられた入力内で欠けている部分シーケンスを回復するように、異なる目的関数をノイズ除去タスクとして扱う。
モデルの構成とトレーニングの目的が決まったら、GPUのマルチノードクラスタでトレーニングを実行する。トレーニングするモデルのサイズや、トレーニングの完了までの時間に応じて、各トレーニング実行に割り当てるノード数を調整することができる。GPUの大規模なクラスタ運用にはコストがかかるため、可能な限り効率的な方法でGPUを活用することが重要である。
Weights & Biasesを使用してリソースの利用状況やトレーニングの進捗状況など、トレーニングプロセスを監視している。トレーニングプロセスの各ステップを通じてモデルが効率的に学習していることを確認するために「損失曲線」を監視している。
また、損失値が突如として上昇する「損失スパイク」も監視している。これはトレーニングデータやモデルのアーキテクチャに問題がある場合に発生する現象である。このような現象が発生したときには更なる調査や調整が必要になることが多いため、このような損失スパイクの原因をより簡単に再現、診断、解決できるようにしている。
評価
モデルをテストするためにChen et al. (2021)で示されているHumanEvalフレームワークのバリエーションを使用する。このモデルを使用し、関数のシグネチャとdocstringが与えられたPythonコードのブロックを生成し、生成された関数に対してテストケースを実行し、生成されたコードブロックが期待通りに動作するかどうかを判断する。複数のサンプルを実行し、対応するPass@Kの数値を分析する。
この方法は、評価器とテストケースがすぐに使えるPythonで最も効果的である。しかし幅広いプログラミング言語に対応するためにはどのようなプログラミング言語でも再現可能な実行環境を構築する必要があったり、テストケースを実行できないプログラミング言語（HTML、CSSなど）では評価が曖昧になってしまう問題がある。
本番環境へのデプロイ
モデルの学習と評価が完了した後、本番環境へデプロイする。
ReplitではNADIAのFasterTransformerとTriton Serverを利用して推論プロセスを高速化している。FasterTransformerは、Transformerベースのニューラルネットワークの推論を高速化するエンジンを実装したライブラリで、Tritonは設定が簡単で安定した高速推論サーバである。この組み合わせにより、Transformerモデルと基礎となるGPUハードウェアの間に高度に最適化された層ができ、大規模モデルの超高速分散推論が可能となっている。
またKubernatesを使用することにより、デマンドに応じての自動スケーリングにも対応している。推論サーバのホスティングにはモデルの重みといったアーティファクトの巨大さ、GPUの数やサイズの違いといった特殊なハードウェア要件のように独特の課題があるが、これに対応できるアーキテクチャを構築している。
実際にユーザーに提供する前に、HumanEvalテスト以外にも自分たちの手で実際に使ってみて、モデルの感触を掴むことも重要視している。
指標として、モデルの性能と利用指標の両方を監視している。モデルの性能は、リクエストの待ち時間やGPUの使用率などの指標を監視している。使用率については、コード提案の受入率を追跡し、プログラミング言語など複数の次元で分類している。これにより様々なモデルのA/Bテストが可能になり、あるモデルと他のモデルの比較のための定量的な指標を得る事ができる。
フィードバックとイテレーション
Replit社のモデルトレーニングプラットフォームは、生データから本番にデプロイされるモデルになるまで1日以内に完了する能力を備えている。しかしそれ以上に重要なのは、デプロイしたモデルに対するフィードバックを集め、フィードバックに基づいて迅速にイテレーションすることができることだ。
次にはReplitそのものを使ってモデルを拡張できるようにプラットフォームを拡張していく予定である。これには人間のフィードバックに基づく強化学習（RLHF）や、Replit Bountiesで収集したデータを使ったインストラクションチューニングなどの技術が含まれる。
所感
未来観として、ゆくゆくは各社で独自のLLMsを所持するような世界になっていくかも知れないなと思ってはいるものの、独自のLLMsを成長させるためのデータパイプラインの構築や推論サーバのホスティングにはかなり高度な技術とランニングコストが要求されるため、一筋縄で用意できるものではないことがReplit社の事例を見ても分かる。
とはいえ、クラウドベンダー各社の動きを見ていると、このデータパイプラインからデリバリまでがサービスとしてサポートされる未来もそう遠くないと感じる。クラウドベンダーにはロックインされてしまうが、そのあたりは事業規模も踏まえたバランスの問題であり、大多数の事業体にとっては落とし所になるのではなかろうか。
現場からは以上です。