4月24日、私たちはSnowflake Arcticをリリースしました。その主な目標は、真にオープンであることです。この目標に沿って、Snowflake AI調査チームは、Arcticなどの大規模な専門家混合（MoE）の事前トレーニング、ファインチューニング、評価、サービスの提供方法を説明する一連の料理本を執筆中です。Arcticモデルをトレーニングした過程、トレーニング前データの調達と構成、MoEアーキテクチャの設計、トレーニングと推論のシステムを考慮したモデルの共同設計、モデルの指示調整と評価の方法に関する知見を共有します。 

シリーズ全体については、Snowflake Arctic cookbookカタログを参照してください。

ここ数か月は、目を見張るようなオープンソースLLMと、LLMをファインチューニングして提供する革新的な方法で満たされ、これらのモデルがますます利用しやすくなりました。この話題にもかかわらず、トレーニングMoEモデルの特殊性について一般に共有されることはほとんどありません。このブログ記事では、(1)インストラクションチューニングの概要、(2)モデルの機能、(3)データセットミックスのファインチューニングの内訳、(4)このようなスパースモデルを大勢のエキスパートでファインチューニングする際に直面したいくつかの問題の解決方法について説明します。

インストラクションチューニングの目標は、プレトレーニングで基本的な言語スキルをすでに習得しているLLMに、指示を正確に実行し、さまざまなテキストタスクを実行する方法を教えることです。通常、命令調整済みモデルは1つまたは2つのステージを経て、

1. 一般的な指示と推論に関連するさまざまなユースケースをカバーするプロンプト完了形式のデータセットのファインチューニング（SFT）を監督。より狭い焦点がある場合は、一般的なSFTに続いて、チャットやコード生成などのより特殊なユースケースに焦点を当てた別のSFTを行うこともできます。
2. オプションとして、嗜好データとPPOやDPOなどのアルゴリズムを使用したアライメントトレーニングがあります。

ファインチューニングにはよりシンプルなアプローチを採用し、一般的な指示に従うタスク、コード、数学、チャットデータをLLMのエンタープライズアプリケーションの要件を念頭に置いて組み合わせて使用しました。Arctic-Instructはマルチターンの会話を行うことができますが、そのような設定に明示的に最適化されているわけではありません。

リリースブログで取り上げられているように、Arctic-Instructはほとんどの非生成および学術的なベンチマークで優れた性能を発揮し、企業固有のタスクを得意としています。Arctic-Instructのエンタープライズベンチマークの平均パフォーマンスは、DBRX、Mixtral 8x7B、Llama 2 70Bを上回っており、はるかに多くのコンピュートでトレーニングされている（Llama 3 70BはArcticの16倍のコンピュートで事前トレーニングされているなど）Llama 3 70BとMixtral 8x22Bの射程圏内にあります。私たちのモデルは、一般的な常識的推論ベンチマークでも競争力のある性能を発揮します。 

Arctic-Instructは会話機能を備えていますが、シングルターンインタラクション向けに最適化されています。さらに、RLHF/DPOやアライメントトレーニングは一切行っていません。コミュニティにぜひご参加いただき、モデルのファインチューニングを簡単に実行するためのツールをリリースする予定です。選好トレーニングを行わなくても、Arctic-InstructはMTBenchのスコア7.95、ターン1のスコア8.31を達成し、Helpful, Honest, & Harmless（HHH）アライメントデータセットで競争力のあるパフォーマンスを発揮しています。

 8億個のトークンからなる約140万個のサンプルを含む多様なデータセットに対して、Arcticをファインチューニングしました。類似しきい値0.7のミンハッシュ重複排除を使用して、このデータを重複排除しました。

データセットのより詳細な内訳は次のとおりです。

• 一般：データセットのこの部分は、Chain-of- Thought（CoT； 思考連鎖）推論の有無にかかわらず、さまざまな命令応答ペアで構成されます。
• コード：コーディングは、プログラミングの問題と、プログラミングトピックに関する概念的な質問で構成されています。GitHubで最も人気のある言語に重点が置かれており、例文の50%以上がPythonで、Java、JavaScript、Goはかなり少数です。
• 数学：このデータの大部分は、説明のトレースを含む単語問題の形式です。これらの問題は、算術や単純な代数、たとえば線形方程式を解く能力で解決できます。
• IDの質問：これらは、モデルの名前やトレーニング方法などを正しく応答させるためのトレーニングに使用される1k個の例です。これらの例は、2つの異なる方法で生成しました。テンプレートから派生したものもあれば、まず手動でシードセットにラベルを付け、それを言い換えて得られたものもあります。

このモデルのサイズのため、私たちは次の2つの段階でデータセットミックスを仕上げるアブレーション実験を行いました。

1. 私たちは、事前にトレーニングしたより小さなMoEモデル（最初は2Bのアクティブパラメータ、次に7Bのアクティブパラメータ）で、チャット、数学、プログラミングデータの割合に焦点を当てたより多くの実験を実行しました。この段階の結果として、いくつかのデータセット候補「ビーム」が得られました。
2. 続いて、これらの候補から最終的なデータセットに収束する、より大きなモデル（17Bのアクティブパラメータ）についての小規模な実験を行いました。

興味深い点をいくつかご紹介します。

• モデルサイズが小さいほど、より多くの数学データが変換され、数学に加えてプログラミングベンチマークのスコアも向上します。このコード改善は、おそらくベースモデルの非常に強力なコーディング能力のおかげで、より大きなモデルサイズで蒸発しました。
• どちらのモデルサイズでも、数学データセットとその他のミックス（コーディング、チャットなど）を組み合わせた場合、単独でファインチューニングした場合よりも、数学ベンチマークではるかに優れた結果が確認されました。これは、私たちのモデルが、数学データセット内の推論パターンだけにさらされていると、すぐに適合しすぎることを示しています。

2e-5の学習率を用いて、フルパラメータでモデルを2つのエポックにファインチューニングし、学習ステップの最初の5%で学習率がウォームアップし、その後コサイン減衰しました。トレーニングは、エキスパート並列処理とZeRO-2を使用して32ノード（256台のH100 GPU）で行われました。デバイスあたりのバッチサイズを1とし、グローバルバッチサイズを256のシーケンスパック例にしました。それぞれの例で、完了トークンにのみ適用される標準的なクロスエントロピー損失関数を使用しました。

インストラクションチューニングArcticでの経験を通じて、大規模なMoEモデルのトレーニングに特有の課題がいくつか見つかりました。このセクションでは、私たちが直面したいくつかの問題とその克服を試みました。

問題1：専門家による安定性。1つの課題は、128人のエキスパートを擁する非常にスパースなMoEアーキテクチャ自体にあります。ファインチューニングの例ごとに、ミニバッチ内の各トークンが2人のエキスパートに送信されます。しかし、中央値の場合、完了トークンの数は多くても数百個であるため、各エキスパートは1つの例に対して1個未満しか見ません。その結果、専門家のトレーニングが不安定になる可能性があります。

問題2：トレーニング効率。Arctic-Instructモデルのパラメータは480Bであり、命令チューニングには数千時間のGPU時間がかかる場合があります。ファインチューニングへの甘いアプローチは、いくつかの非効率性をもたらします。特に、私たちのトレーニングシステムは、すべての例を同じ長さ（この例では4096）にパディングしてファインチューニングすると最も効率的です。そのため、パディングで大量のコンピュートが無駄になり、実際のトレーニング進捗に寄与しません。

注目すべきは、これらの課題はいずれもファインチューニング特有のものです。なぜなら、プレトレーニングではすべてのトークンに対してトレーニングを行い、パディングなしでシーケンスを最大長に詰め込むのは簡単なためです。その結果、私たちはプレトレーニングのアイデアを借りました。

解決策1：複数の例を1つのシーケンスに詰め込む。インストラクションチューニングで再利用したプレトレーニングからの標準的なプラクティスは、シーケンスパッキングです。長さ4096の各シーケンスに1つの（プロンプト、完了）ペアだけでなく、複数の例を（prompt_1、completion_1、…、prompt_n、completion_n）という形式の各シーケンスに「パック」しました。シーケンスパッキングにより、各バッチにより多くの例が収まるため、トレーニングのスループットが向上します。

解決策2：会話を詰め込むと、1つの例になります。マルチターンの例をトレーニングする最も簡単な方法は、アシスタントの完了ごとに個別のトレーニングポイントを設けて、例をロール解除することです。各会話を単一のトレーニング例として扱い、すべてのアシスタントの完了について共同でトレーニングを行います。これにより、会話ごとに必要なトレーニングコンピュートが減るだけでなく、各バッチのトレーニング可能なトークンの割合が増えるため、エキスパートの勾配スパース問題が軽減されます。

解決策3：トレーニング中のMoEレイヤーの凍結。トレーニング効率の改善を図ったもう1つのソリューションは、トレーニング中にすべてのMoEレイヤーをフリーズすることです。これは、私たちのモデルアーキテクチャに特化して調整された効率技術であり（Arcticの高密度レイヤーが残っているため）、プレーンなLoRAよりも私たちのアーキテクチャに適していると考えられます。これにより、トレーニング中のグラデーションのコンピュートとサイズが大幅に減少し、はるかに少ないGPUで行えるようになります。これにより、一般的な非生成ベンチマーク+GSM8Kは劣化しないことが確認されました。しかし、HumanEvalなどの他の生成ベンチマークは、大きな影響を受け、5～6ポイント低下しました。最終的なレシピには使いませんでしたが、異なるデータ構成の繰り返しを行うときに役立つ近似でした。

MoEモデルをトレーニングする際のもう1つの重要な考慮事項は、トークンがエキスパート間で均等にルーティングされることです。ほとんどのMoEアーキテクチャでは、1人のエキスパートの計算負荷を最小限に抑えるために、レイヤー内のエキスパートごとにトークン容量が固定されています。MoE レイヤのルーティング機能がエキスパートに割り当てたトークンが多すぎる場合、超過分のトークンはドロップされ、残余を介して次のレイヤに渡されるだけです。したがって、モデルが可能な限りロードバランシングを学習することが重要です。ファインチューニングの初期段階では、ファインチューニング中と事前トレーニング中に使用する効果的なトレーニングトークンが小さかったため、非常に妄想的になっていました。 

ファインチューニングプロセスを通じて、MoEロードバランシングの損失を監視し、理論上の最小値に近づけることで、トークンルーティングが確かに均一に近くなることを確認しました。具体的には、GShardエキスパート損失の正規化バリアントを使用し、レイヤーごとに

ここで、kは定数（0.01）、Eはエキスパートの数、Sはバッチ内のトークンの数、 c_eはエキスパートeにルーティングされたトークンの数、m_eはすべてのSトークンでエキスパートeに割り当てられた平均ゲーティングウェイトです。この損失関数は、ルーティング関数がすべてのトークンについてすべてのエキスパートに1/Eのゲーティングウェイトを割り当てるときに最小化されます。この場合、平均ゲーティングウェイトは1/Eで、各エキスパートはS/Eトークンを中心に割り当てられます。

この損失は35層のトランスレイヤーで合計されるため、この損失の最小値は0.35となります。ファインチューニングを通じて、バランス損失は約0.35にとどまっており、私たちのモデルがエキスパート間で入力トークンのバランスをほぼ完璧に維持していることを示しています。 

SnowflakeのArcticベースモデルは、一般的なテキストタスクで優れた性能を発揮し、SQLやコーディングなどのエンタープライズ関連のタスクを得意とするようにインストラクションチューニングされています。今後も、インストラクションチューニングに関するより多くの知見を共有していくとともに、オープンソースコミュニティと協力し、より多くの人が強力なLLMを利用できるようになることを楽しみにしています。今後もSnowflake Arctic cookbookカタログに最新情報が追加される予定ですので、ご期待ください。