Marin、535BパラメータMoEモデルのオープントレーニングを開始
Marinは、公開コード、ログ、スケーリング予測を伴う、18.75兆トークンでの535BパラメータMoEトレーニングを開始した。
一、公開環境で実施される535Bパラメータの実行
Marinは、総パラメータ数約5,350億、各トークンでアクティブになるパラメータ数230億のMixture-of-Experts言語モデル、Marin 535B-A23Bのトレーニングを開始した。このプロジェクトは、約3か月にわたり11基のNVIDIA GB200 NVL72システムを用いて、18.75兆トークンを処理する計画である。
公開された計画では、トークン予算の80%を事前学習に、20%を中間学習に割り当てている。その後にポストトレーニングを実施するが、Marinは最終的なポストトレーニングのレシピやスケジュールをまだ公表していない。プロジェクトは、メイン実行に必要な浮動小数点演算量を約 \(2.7 \times 10^{24}\) と見積もっている。
各GB200 NVL72は、72基のBlackwell GPUと36基のGrace CPUを搭載するラックスケールのシステムである。このようなシステム11基は792基のGPUに相当するハードウェア規模となるが、Marinのエキスパート並列実装レポートでは、各トレーニングラック内に64 GPUのエキスパート並列ドメインがあると説明されている。
この実行はモデルのリリースではない。8月24日時点でトレーニングは進行中であり、評価できる最終ウェイト、ベンチマーク結果、ポストトレーニング評価は存在しない。その直接的な重要性は、成功した結果を選んだ後に技術レポートだけを公表するのではなく、この規模のトレーニングを実施中から公開するという判断にある。
この実行に関するMarinの公開Issueは8月18日に開設された。そこには運用計画、エンジニアリング上のリスク、スケーリング手法、コンテキスト拡張案、コンティンジェンシー手順が含まれている。リンクされたWeights & Biasesレポートは、このプロジェクトのライブ追跡画面を提供している。
発表では18.75兆トークンを明記している一方、GitHub Issueのタイトルでは実行を「18T tokens」と省略している。より正確な数値は公開されたvoyage planで用いられているものであり、プロジェクトは短いタイトルを改訂後の予算として提示していない。
二、スパースモデルの構成
Marin 535B-A23Bは48個のTransformerブロックを使用する。各ブロックは、アテンション分岐と、384個のルーティング対象エキスパートを含むスパースMoE分岐を組み合わせている。ルーターは各トークンに対して8個のエキスパートを選択し、2個の共有エキスパートはローカルに残り、ルーティング経路とは独立してすべてのトークンを処理する。
モデル状態の幅は6,144値である。ルーティングされたアクティベーションをGPU間で交換する前に、潜在射影によって3,072値へ圧縮する。Marinによれば、これによりエキスパート並列all-to-all操作を通過するアクティベーショントラフィックの幅が半減する。出力は共有エキスパート経路と結合される前に、モデル幅へ再射影される。
この転送問題は大きい。64 GPUのエキスパート並列ドメインに384個のルーティング対象エキスパートを分散すると、各GPUは6個のルーティング対象エキスパートをホストする。エキスパートごとに別々のサイズのバッファを送信すると、厳しいメモリ要件と動的な通信パターンが生じる。
そこでMarinは、JAXおよびXLA向けに固定プールウェーブ型all-to-all実装を開発した。送信側は、エキスパートごとに1つのバッファを作るのではなく、宛先GPUごとに1つの固定プールを作成する。転送は同一の配列形状で3つの連続ウェーブとして行われ、エキスパート識別子はアクティベーションペイロードにパックされる。これにより、トークン数やルーティングメタデータを別途交換する必要がなくなる。
この実装は2種類の容量制限を使用する。送信側容量係数1.10は、1つの送信元から1つの宛先へ送られるトラフィック量を制限し、受信側容量係数1.15は、個々のローカルエキスパートに割り当てられる行数を制限する。いずれかの固定バッファを超えた割り当てはドロップされ、別途報告される。
20ステップ、1ラックのゲートはアウト・オブ・メモリ障害なしに完了し、ステップ2から19までの中央値スループットは毎秒250,691トークンだった。Marinは、この短いテストでは最終的なトークンドロップ率は確立できないと明示的に警告している。これはエンジニアリング上の適格性確認結果であり、トレーニング品質のベンチマークでも、11ラック全体の実行の測定でもない。
レポートには失敗した構成も記録されている。直接的な固定エキスパートセル設計では、XLAのメモリ見積もりが192.65 GiBとなり、123.49 GiBのCUDA割り当てで失敗した。6エキスパートの受信バンクもメモリ不足となった一方、処理を3ウェーブに分割すると、そのバンクの一部だけを同時にアクティブに保てた。こうしたネガティブな結果も、プロジェクトの公開された設計記録の一部である。
三、メイン実行に先行したスケーリングラダー
535Bモデルを開始する前に、Marinは4段階のスケーリングラダーをトレーニングした。これは、6,100万パラメータをアクティブにする16億パラメータのMoEを480億トークンでトレーニングするところから始まった。最大の段階は総パラメータ数277億、アクティブパラメータ数12億で、9,260億トークンを処理した。
このラダーは、予測と診断の基準点の両方として機能する。Marinはメイン実行の損失、勾配ノルム、トークンドロップ、評価の軌跡を、小規模で観測されたパターンと比較できる。大きな乖離があれば、フル実行が数か月分の計算資源を消費する前に調査を開始できる。
プロジェクトによれば、このラダーのコストはメイン実行の計算量の約1%である。Marinはこの支出をリスク管理の一形態と見なしている。すなわち、選択したアーキテクチャ、データ混合、オプティマイザ設定、トレーニング期間が、モデルサイズの増加に伴って一貫して振る舞うかをテストするものだ。
プロジェクトによると、以前のラダーでは、トークンの学習期間を延長するにつれて勾配ノルムが4を超えて増加することが明らかになった。この発見を受け、logit z-lossが使用されるようになった。その後のアブレーションでは、一部の高バッチ構成がそうでなければトレーニング中に発散する可能性があることが示されたとされる。
小規模な実行は、535Bモデルがその予測に従うことを保証するものではない。外挿は依然としてこの実験の中心的な不確実性の一つである。その実用的な価値は、約100日に及ぶ実行中に行われる介入を外部の研究者が評価できる基準を、Marinが公開したことにある。
Marinは、インフラ遅延や想定を下回るモデルFLOP利用率に対するコンティンジェンシーも記録している。トークン予算のおよそ最初の4分の1では、デフォルトの対応としてトークンの学習期間を短縮し、データ混合を調整し、改訂後の終点でも線形学習率減衰がピークレートの5%に到達するようタイミングを再設定する。このことは、発表されたスケジュールが不変の仕様ではなく、運用計画であることを意味する。
四、長いコンテキストはトークンドロップの解決に依存する
モデルは、4,096トークンのシーケンス長で事前学習を開始する。Marinの以前の大規模実行は8Kで開始し、1兆トークンにわたって65Kへ拡張され、その後262Kへの拡張が予定されていた。4Kへ戻すと、各バッチに8Kの場合の2倍の個別シーケンスが含まれるため、トークンがエキスパート間でより均等に分配されるはずである。
この選択は、現行のエキスパート並列実装の弱点に対処するものである。以前のテストでは、トークンドロップは4Kコンテキストで約7%だったものが、65Kでは約40%まで増加した。Marinは、新しいプールウェーブ設計では4Kで約3%がドロップすると報告しているが、65Kでの率は依然として過大になる可能性があると見込んでいる。
ドロップされた割り当ては、必ずしもトークン全体が消えることを意味しない。2個の共有エキスパートはすべてのトークンを引き続き処理し、選択される8個のルーティング対象エキスパートは、その共有経路に追加されるものである。Marinは、ルーティングされた割り当てが利用可能な容量を超えた場合に共有エキスパートがより密なバックボーンを提供すると述べているが、高いドロップ率はスパースエキスパートの利点を減らし、トレーニングの挙動を変える可能性がある。
プロジェクトは、実行開始から約10〜20日後に、1〜2日間の早期クールダウンを計画している。この分岐は、強化学習実験向けのフルスケールチェックポイントを提供し、メインのトレーニング軌跡を変更せずに、より長いコンテキストがルーティングへ与える影響を検証することを目的としている。
この実験が安定すれば、暫定的なコンテキストスケジュールは、トレーニングの中間で4Kから8Kへ、約95%時点で8Kから65Kへ、そして終盤近くで目標とする262Kフェーズへ移行する。これらは条件付きの目標であり、完成したモデルで確認済みの能力ではない。
Marinは、トークンドロップが高すぎるままである場合の代替案を3つ挙げている。ドロップレスなragged all-to-all実装を採用する、容量係数を増やして付随するメモリコストを受け入れる、またはシーケンスレベルのバランシングを導入する、の3つである。最後の選択肢ではトレーニング中にエキスパートの再専門化を強いる可能性があるため、プロジェクトは現時点でフォールバックとして扱っている。
五、この実行をオープン開発の節目にするもの
Marinは、オープン開発をトレーニング後のウェイト公開とは区別している。その標準ワークフローは、実験の仮説と目標を記録するGitHub Issueから始まる。実装はレビュー可能なコードとして提出され、実行は公開テレメトリーへリンクされ、失敗した試行を含む分析がIssueへ戻される。
535B実行については、公開記録にはすでにhero-run Issue、ライブ追跡レポート、ソースリポジトリ、そしてエキスパート転送実装の詳細な説明が含まれている。転送レポートは、測定結果とエンジニアリング上の判断を分け、どの結論が短いプロファイリング実行だけに基づくかを特定している。
この記録により、研究者はスケーリング予測がフルモデルでの実際の結果に耐えるかを評価できる。また、最終モデルカードでは数行に圧縮されがちな判断も明らかにする。なぜ実行が4Kコンテキストで始まるのか、どこでトークン割り当てがドロップされるのか、メモリ制約が転送設計をどう変えたのか、トレーニングが予定より遅れた場合にチームが何をする予定なのか、といった判断である。
公開ドキュメントによって、メイン実験の再現コストが低くなるわけではない。11ラックの実行を再現することは、依然としてほとんどの独立研究者のリソースを超えている。より利用しやすい成果物は、コード、小規模なスケーリング実行、失敗分析、構成上の選択、そして縮小スケールで検査またはテストできるライブ測定値である。
したがって、プロジェクトの状況は限定的に説明する必要がある。Marinは、オープンに文書化されたフロンティア規模のトレーニング実行を開始した。最終モデルの品質、長文コンテキストの挙動、ポストトレーニング性能、最終的な成果物のリリースは、まだ実証されていない。
よくある質問
「535B-A23B」は何を意味しますか？
このモデルは総パラメータ数が約5,350億で、各トークンでは約230億がアクティブになる。各Transformerブロックでは、スパースルーティングが384個のルーティング対象エキスパートのうち8個を選択する。
完成したモデルは今すぐダウンロードできますか？
いいえ。トレーニングはまだ進行中であり、Marinはこの実行の最終ウェイトやポストトレーニング評価を公開していない。
なぜトレーニングは4Kコンテキストだけで始まるのですか？
4Kのシーケンス長では、各バッチにより多くの個別シーケンスが含まれ、エキスパートのバランシングが改善される。Marinは、以前の実装で65K時にトークンドロップが急増したため、より長いコンテキストをテストしている。
研究者はどのように実行を追跡できますか？
Marinは、公開GitHub Issue、ソースリポジトリ、エキスパート並列の設計レポート、リンクされたWeights & Biases追跡ページを提供している。
262Kのコンテキスト長は保証されていますか？
いいえ。8K、65K、262Kへの拡張案は、早期クールダウンおよびトークンドロップ実験の結果に依存する。
参考ソース
Share