スタンフォード大学の Marin 基盤モデル: JAX を使用して開発された初の完全オープンモデル

現在の AI 時代におけるエキサイティングな要素として、強力な基盤モデルがオープンに共有され、あらゆる人のイノベーションを加速している点が挙げられます。この進歩から、「オープン性の次のステップは何なのか？」という問いが見えてきます。Marin プロジェクトでは、「オープン」の定義をモデルの背後にある科学的プロセス全体にまで拡張する機会を見出しています。

スタンフォード大学 Center for Research on Foundation Models（CRFM）の Marin プロジェクトは、「オープンラボ」として設計されており、モデルを共有するだけでなく、コード、データセット、データ手法、テスト、ハイパーパラメータ、トレーニング ログなど、開発のすべての過程を公開することを目標としています。このレベルの透明性は、研究者が開発中のモデルを精査、構築し、信頼性を高められるようにするユニークで完全に再現可能なリソースを提供することで、既存のエコシステムを補完します。スタンフォード大学の Marin プロジェクトは、今後の基盤モデル研究をより透明でアクセスしやすいものにすることを目指しています。

AI モデルのオープン性の範囲

このオープンラボから最初に公開されたモデルは、Marin-8B-Base と Marin-8B-Instruct です。プロジェクトの原則に従って、モデル、データ、コード、トークナイザーはすべて許容性の高い Apache 2.0 ライセンスの下で公開されています。完全な再現性を目指すこの取り組みには、極めて困難なエンジニアリング上の課題があります。大規模に分散されたシステムにおいて、あらゆる変動要因を制御する必要があります。このプロジェクトの成功は、この再現性を大規模に実現し、優れたコスト パフォーマンスで基盤モデルのトレーニング効率を最大化できるテクノロジー スタックにかかっています。

オープンな基盤モデルの構築における主な課題

Marin プロジェクトにおいて、真にオープンでスケーラブルかつ再現可能な基盤モデルの作成に成功するために、CRFM チームはいくつかのエンジニアリング上の課題を解決しなければなりませんでした。チームが JAX を基盤として選んだのは、その設計原理がこれらの課題に対する直接的な解決策を提供していたためです。そして、JAX の力を活用するための新しいフレームワークであるLevanter（下記を参照）を構築しました。以下に課題とその解決策の例をいくつかご紹介します。

1 つのアクセラレータで最大速度を達成

課題: コアとなるトレーニング ループは何十億回も実行されるため、Python のようなインタプリタ言語に起因するオーバーヘッドが、パフォーマンスに対する大きなボトルネックとなります。オペレーションが段階的にディスパッチされる場合、ループによってメモリ トラフィックとオーバーヘッドが過剰に生じる可能性もあります。特に TPU のようなハードウェアでは、スループットが融合されたオペレーションの効率性に依存しているため、この問題が顕著になります。

解決策:

• インタプリタのオーバーヘッドを排除するため、Levanter は複数段階のトレーニング ステップ（フォワードパス、損失計算、誤差逆伝播、更新）全体を単一の関数にカプセル化し、@jax.jit デコレータを使用します。JAX の XLA コンパイラは、このプロセス全体を高度に最適化された単一のマシンコード カーネルに変換し、オペレーションを融合させて大規模にハードウェアを最大限活用できるようにします。

• 冗長な計算を避けるため、jax.value_and_grad を使用して 1 つのパスで損失とその勾配の両方を計算します。また、JAX を使用すると、勾配チェックポインティングなどの高度な手法を簡単に使用でき、メモリを節約して、オーバーヘッドをほとんど発生させることなくより大きなバッチサイズを使用できます。

• Levanter はまた、JAX の強力な Pallas ベースの Splash Attention カーネルを使用しています。これは、ほぼすべての大規模言語モデルの中核をなす最も重要なオペレーションの一つである Dot Product Attention の高度に最適化された実装です。

大規模な並列処理の複雑さを克服

課題: 最先端のモデルをトレーニングするには、何千ものアクセラレータ チップにスケールアウトする必要があります。モデルとデータの分割方法やデバイス間の通信方法を手動で管理するのは非常に複雑で、コードの読み取り、デバッグ、適応がすぐに困難になります。

解決策:

• JAX の @jax.jit デコレータは、基盤となるデータのシャーディングと通信を自動化する「Single-Program, Multiple-Data 」（SPMD）並列化もシームレスにサポートしています。XLA コンパイラはアクセラレータ間の通信を自動的にスケジュールして、ネットワーク上の待機時間を最小限に抑え、計算に費やす時間を最大化します。

• jit の機能をさらに安全で使いやすくするために、Levanter は名前付きテンソルのライブラリである Haliax を開発しました。テンソル軸を、位置インデックスではなく人間が読み取り可能な名前（「embed」や「batch」など）で参照することにより、コードの自己文書化が進み堅牢になります。

• この抽象化により、モデルコードに触れることなく構成ファイルの数行を変更するだけで、Fully Sharded Data Parallelism（FSDP）やテンソル並列処理などの高度なシャーディング戦略を定義し、変更できます。

レジリエントで費用対効果の高いコンピューティング クラスタの構築と管理

課題: 大規模なトレーニングには、大規模なコンピューティング クラスタへの柔軟なアクセスが不可欠です。コスト管理のためにプリエンプティブル TPU インスタンスに大きく依存しているので、多くの小規模で分散した TPU スライスを 1 つの論理クラスタに簡単に結合し、頻繁な中断にも耐えられる仕組みが必要です。

解決策:

• Google Cloud TPU マルチスライスは複数の TPU スライスを 1 つの大きなシステムとして使用できるテクノロジーです。これを活用することで、多数の小規模なプリエンプティブル TPU スライスを 1 つの強力なコンピューティング クラスタに簡単に結合し、トレーニングに利用できます。

• Levanter は、Ray を使用してこのプロセスをオーケストレートし、トレーニング ジョブ中に TPU スライスの数をシームレスに増減させます。さらに重要な点として、1 つのスライスがプリエンプトされた場合でも確実にジョブのレジリエンスを維持できるようにしました。

• JAX と XLA のおかげで、Levanter と Marin は GPU 上でも同様の高性能な結果を達成しました。

完璧な再現性で科学的信頼を醸成

課題: Marin プロジェクトの中核的な目標は、検証可能な科学を実現することです。そのためには、トレーニングが中断、再開、または異なるマシン構成に移行した場合でも再現可能な結果を達成する必要があります。これは大きな技術的ハードルでした。

解決策:

• これは Levanter の設計推進における基本的な要件でした。JAX を選んだのは、決定論的擬似乱数生成器（PRNG）をデフォルトで使用するなど、強力な再現性の保証を備えていたためです。

• この選択については Marin-8B のトレーニング中に検証されました。トレーニングではさまざまな TPU スライスとハードウェア タイプ間を移行しながら、プリエンプション間でビット単位での再現性を完全に維持することに成功しました。

• Levanter には、Google の TensorStore ライブラリを基盤とした堅牢なデータ読み込みシステムも組み込まれています。Levanter のデータストアは、ジョブの再開やデータソースの変更とはかかわらない、トレーニング データ バッチへの決定論的なランダム アクセスを提供します。これは中間トレーニングなどの高度なトレーニング戦略をサポートするうえで非常に重要です。JAX の決定論と Levanter のデータストアにより、解釈可能性の研究者は、トレーニング中に特定のデータがモデルにどのような影響を与えるかを理解しやすくなります。

一貫性のあるフレームワークを作成

課題: JAX は強力なエンジンを提供していますが、既存の高レベル フレームワークはいずれも、読みやすさ、大規模なスケーラビリティ、ビット単位の決定論という私たちの厳格な要件をすべて満たすものではありませんでした。トレーニング プロセス全体をオーケストレートするために、包括的で明確な方針を持つシステムが必要でした。

解決策:

• 私たちは JAX ネイティブ フレームワークである Levanter をゼロから構築し、ビット単位の決定論性、高度な分散戦略によるスケーラビリティ、レジリエンスを備えた、私たちの求めるシステムにすることを目指しました。

• この目標を実現できたのは、JAX が単なるライブラリではなく、新しいツールを構築するための「メタフレームワーク」だからです。私たちは、JAX の成熟した高性能な TPU サポートと、高レベルの抽象化（jit）と低レベルの制御（Pallas）とのシームレスな統合を基盤として開発を進めました。

• このアプローチは JAX コミュニティでは一般的で、Flax、Equinox、Orbax、Optax などのライブラリが相互に連携して機能する活気に満ちたエコシステムを生み出しています。このおかげで、私たちのようなチームが強力なソリューションを構築できるようになっています。

仕組み: Marin-8B の構築

前述の原則、ツール、ライブラリは、Marin-8B のトレーニング実行中に実装され、実際に機能しました。モデル アーキテクチャは Llama スタイルの Transformer です。

Marin-8B-Base: モデル アーキテクチャの概要

Marin-8B のトレーニングは静的でモノリシックな実行ではなく、適応的なプロセスであり、社内では「Tootsie」プロセスと呼ばれていました。これは、実際の研究ワークフローを包み隠さず描写したもので、詳細が一般に公開されています。このプロセスは 12 兆を超えるトークンに及び、新しいデータ、手法、さらには大規模なマルチスライス TPU 構成（2x v5e-256 から 1x v4-2048 ポッド）への移行を途中で行うなど、異なるハードウェア構成に適応する複数のフェーズを経て実施されました。チームは継続的に混合データを改良し、より高品質なソースを組み込んで学習率やバッチサイズなどのハイパーパラメータの調整を行い、パフォーマンスを最適化しました。この「ごちゃごちゃした」現実のプロセスは強力な教材であり、JAX と Levanter スタックがこれらの重大な変化に対応しながらもビット単位の再現性を維持できたという能力は、その技術の堅牢性を強力に実証しています。

Marin コミュニティに参加

Marin プロジェクトは、基盤モデル開発の将来に参加し、JAX エコシステムに貢献することを広く呼びかける招待状です。Marin の歩みは「オープン性の次のステップとは何なのか？」という問いへの答えを表しています。この「オープンラボ」創出の取り組みは、JAX エコシステムの技術的能力によって可能になりました。そのパフォーマンス、ポータビリティ、再現性を重視した基礎設計は、研究の「完全な旅」を実現するための重要な要素となっています。

データ手法からトレーニング ログまで、あらゆる情報を共有することで、私たちは完全に再現可能なリソースを提供することを目指しています。このリソースにより、研究者は作業を詳細に検証し、そのうえでの構築し、信頼性を高めることができます。これは、より透明性の高い AI の未来に向けたコラボレーションの一歩であると考えています。ぜひこの「オープンラボ」に参加し、Marin を活用してください。研究に貢献することで、革新的かつ信頼性の高い次世代の基盤モデル構築にご協力ください。

プロジェクトの主なリソースは公式ウェブサイト marin.community です。そこから、Hugging Face でリリースされたモデルにアクセスする、GitHub の「オープンラボ」を探索する、Marin のドキュメントを閲覧する、Levanter のトレーニング フレームワークの詳細を学ぶ、といったことが可能です。シンプルな推論例を使用して、Colab で Marin を試すこともできます。

また、Discord チャンネルでは活発な議論が行われており、他のデベロッパーと直接交流できます。エコシステムを初めて利用する方のために、JAX の公式ドキュメントにはクイックスタート ガイドなどの優れたリソースも掲載されていますので、そちらもご活用ください。