xAI、3140億パラメータの「Grok-1」ベースモデルをオープンソース公開

xAIは、大規模言語モデル「Grok-1」のベース重みとネットワークアーキテクチャをApache 2.0ライセンスでオープンソース公開しました。

今回公開されたのは、2023年10月に事前学習を完了した時点の生のチェックポイントです。対話タスク向けなどの微調整は一切施されておらず、スクラッチから学習された純粋なベースモデルとなっています。

アーキテクチャと主要スペック

Grok-1は、総パラメータ数3140億を誇る混合エキスパート（MoE）モデルです。合計8基のエキスパートのうち、トークンごとに2基のエキスパートを動的に選択して処理を行います。

ネットワーク構成は64層で、埋め込み次元数は6144です。アテンション機構には48個のクエリヘッドと8個のキー・バリューヘッドを採用し、最大コンテキスト長は8192トークンに対応しています。

トークナイザーには語彙数131072のSentencePieceを採用しているほか、Rotary Embedding（RoPE）、アクティベーションシャーディング、8ビット量子化のサポートを備えています。

動作環境と動かし方

GitHubリポジトリには、JAXをベースにした重みの読み込みと推論テスト用のサンプルコードが用意されています。

3140億パラメータという極めて巨大なサイズのため、付属のコードを検証するだけでも大容量GPUメモリを備えたマシンが必須となります。なお、リポジトリに含まれるMoE層の実装は、カスタムカーネルを使わずに正確性を担保することを優先したため、推論効率は最適化されていないと案内されています。

重みファイルは公式のマグネットリンクを使用したBitTorrent経由、またはHugging Face Hubから直接ダウンロード可能です。

手元の環境で依存関係をインストールし、取得したチェックポイントを配置した上で run.py を実行すると、テスト入力に対する推論テストが行えます。

なぜ重要か

率直に言って、現時点でこのモデルを直接動かせるのは、十分なGPUクラスタを保有する研究機関やインフラ開発者に限られます。ローカルPCで手軽に試したい一般のエンジニアにはハードルが高いのが実情です。

それでも、3140億パラメータ規模の最先端MoEベースモデルが制限の少ないApache 2.0で公開された意義は極めて大きいと私たちは考えます。今後オープンソースコミュニティ主導で、より効率的な推論カーネルの開発や対話向けファインチューニングが急速に進むはずです。