「イーロン・マスクの生成AI「Grok-1」、オープンソース化」ということで、3月11日にGPT-3を超える、3140億パラメーターのMoEがオープンソースとしてリリースされました。
GitHubやHuggingfaceからダウンロードできるようになっているので、自分の環境で動作するかどうか、確認してみました。
私の使っているGPUは、RTX 4060Ti 16GB で、ミッドレンジのものになると思います。
これで動作するのでしょうか。
Grok-1のスペックについて
Parameters: 314B
Architecture: Mixture of 8 Experts (MoE)
Experts Utilization: 2 experts used per token
Layers: 64
Attention Heads: 48 for queries, 8 for keys/values
Embedding Size: 6,144
Tokenization: SentencePiece tokenizer with 131,072 tokens
Additional Features:
Rotary embeddings (RoPE)
Supports activation sharding and 8-bit quantization
Maximum Sequence Length (context): 8,192 tokens
Parameters: 314B 3140億のパラメーターは、オープンソースの中では最大になると思います。
Architecture: Mixture of 8 Experts (MoE)
８つのモデルが協調して動作するということです。
モデルの構造が分からないので、LayerからEmbedding Sizeの数字が大きいのか小さいのかよく分かりません。QKVは小さいように思えますが、どうなのでしょうか。
環境構築
実行環境
WSL2を使う。
最新のCUDA Toolkitが必要
cuDNNも必要
リポジトリのコピー
git clone https://github.com/xai-org/grok-1.git
cd grok-1
仮想環境の設定と初期化
python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
JAXが正しく入らないので、以下を実行
pip install -U "jax[cuda12_pip]==0.4.25" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
モデルのダウンロード
pip install huggingface_hub[hf_transfer]
huggingface-cli download xai-org/grok-1 --repo-type model --include ckpt-0/* --local-dir checkpoints --local-dir-use-symlinks False
15:13 開始
16:22 終了
約70分
モデルファイルは、トータルで310GBとなっています。ファイル数が770です。
実行
python run.py
結果
全く動きません。
まず、デフォルトでは８つの演算デバイスを使用します。要するに、GPUが8つ必要です。
設定を変えて、１つのGPUで動かそうとしました。
local_mesh_config=(1,1) に変更
モデルファイルのロードで、メモリが足りないです。（OSError: [Errno 28] No space left on device: './checkpoints/ckpt-0/tensor00000_000' -> '/dev/shm/tmp_d8pupbr'）
モデルをロードする前に、すでに15GB程度の消費がありましたが、何で使ったのかが分かりません。
モデルファイルはトータルで310GBなので、8枚にロードするとすると、一枚あたり39GBということになります。一般向けのGPUじゃ無理ですね。A100あたりを8枚使うと、余裕で動きそうです。
ということで、クラウドサービス等で動かす場合は、課金に注意というところでしょうか。