イーロンマスクが提供していると話題ですね。<a href="https://grok.x.ai/" rel="noopener noreferrer" target="_blank">https://grok.x.ai/</a> 3140億パラメータあるようです。インストールなどはすべてgithubに書いてある通りにしました。Llama2の70Bよりも多いので、多少インストールに苦労するかもしれません。 <a href="https://github.com/xai-org/grok-1" rel="noopener noreferrer" target="_blank">https://github.com/xai-org/grok-1</a> リポジトリをクローンしてフォルダに移動し、huggingfaceからテンソルをダウンロードします。<pre class="ql-syntax" spellcheck="false">git clone https://github.com/xai-org/grok-1.git &amp;&amp; cd grok-1
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
</pre> このテンソルのファイルが結構大きくて、何回かトライしましたが、通信が遅いとダウンロードだけで45分くらいかかりました。最終的には通信が早い環境で計算できたので、10分くらいでダウンロードが終わりました。おそらくこれだけで300GBくらい消費するので、ディスクサイズは大きめに取っておく必要があります。僕は1TB確保しました。500GBくらいでもいいかもしれませんが。 ダウンロードが終わりましたら、Pythonの仮想環境を作り、必要なツールをインストールしました。どうやらJAXを利用しているようです。 <pre class="ql-syntax" spellcheck="false">python3 -m venv virt
source virt/bin/activate
</pre> 必要なパッケージをインストールします。<pre class="ql-syntax" spellcheck="false">pip install -r requirements.txt
</pre> そして実行です。<pre class="ql-syntax" spellcheck="false">python run.py
</pre> 最初は4GPUで実行すると、すぐにデバイス数でエラーが起きました。<a href="https://github.com/xai-org/grok-1/issues/38" rel="noopener noreferrer" target="_blank">https://github.com/xai-org/grok-1/issues/38</a> 同じような問題を抱えている人は多いみたいで、どうやらJAXを利用していますが、指定のインストールでGPUを使っている場合にはcuda対応のjaxのライブラリをインストール必要があるみたいです。そうでない場合には、deviceの数を(1,1)に指定するとCPUで動作するようです。 せっかくGPU入れてるので、GPUで動かしたいので、 <pre class="ql-syntax" spellcheck="false">pip install --upgrade &#34;jax[cuda12_pip]&#34; -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
</pre> で、cuda12対応のjaxlibを入れました。最初deviceを(1,1)に設定してしまいました。どうやらCPUで動くらしく、ものすごい時間がかかります。あまりに時間がかかるので、GPUにしました。。。 deviceをcpuで動かすには、VRAMは不要ですが、十分なRAM容量が必要です。300GBのテンソルファイルなので、それ以上は必要そうです。GPUの場合には同様にVRAMは300GB以上必要かもしれません。 ということで、快適にGrokを動かすには相当なスペックが必要そうです。RTX6000adaなどはVRAM48GBです。8枚あれば、384GBで要件は満たしそうです。また、A100 80GやH100などはVRAM80Gなので、8枚使えば合計で640GB確保できるので確実に動きそうです。 GPUの枚数はどうやら標準で8枚使用を想定しているようです。ただ、deviceの設定で枚数は指定できるようです<s>が、その際にはkv headsの値を調整するなどが必要のようです</s>。詳しくはdiscussionボードなどを参考にするのが良さそうです。 jaxlibがインストールされましたら、GPUが利用できそうですので、run.pyのファイルを実行すればいいのですが、使い勝手を良くするために、jupyternotebook形式に変更をして実行しました。 結局色々試したところ、かなりのマシンリソースを必要とするみたいなので、考えるのをやめてH100を8枚で実行しました。 <pre class="ql-syntax" spellcheck="false">import logging

from model import LanguageModelConfig, TransformerConfig, QuantizedWeight8bit as QW8Bit
from runners import InferenceRunner, ModelRunner, sample_from_model


CKPT_PATH = &#34;./checkpoints/&#34;

grok_1_model = LanguageModelConfig(
    vocab_size=128 * 1024,
    pad_token=0,
    eos_token=2,
    sequence_len=8192,
    embedding_init_scale=1.0,
    output_multiplier_scale=0.5773502691896257,
    embedding_multiplier_scale=78.38367176906169,
    model=TransformerConfig(
      emb_size=48 * 128,
      widening_factor=8,
      key_size=128,
      num_q_heads=48,
      num_kv_heads=8,
      num_layers=64,
      attn_output_multiplier=0.08838834764831845,
      shard_activations=True,
      # MoE.
      num_experts=8,
      num_selected_experts=2,
      # Activation sharding.
      data_axis=&#34;data&#34;,
      model_axis=&#34;model&#34;,
    ),
)

inference_runner = InferenceRunner(
    pad_sizes=(1024,),
    runner=ModelRunner(
      model=grok_1_model,
      bs_per_device=0.125,
      checkpoint_path=CKPT_PATH,
    ),
    name=&#34;local&#34;,
    load=CKPT_PATH,
    tokenizer_path=&#34;./tokenizer.model&#34;,
    local_mesh_config=(1, 8),
    between_hosts_config=(1, 1),
)
</pre> 流石にH100です。爆速でした。まずは初期化のようなものが、 <pre class="ql-syntax" spellcheck="false">inference_runner.initialize()
</pre> 多分これ、ものすごい時間がかかるのですが、今回はマシンの性能のおかげで、 <pre class="ql-syntax" spellcheck="false">1min 37s
</pre> 1分半で終わりました。そしてこちらは一瞬で終わりました。 <pre class="ql-syntax" spellcheck="false">gen = inference_runner.run()
</pre> <pre class="ql-syntax" spellcheck="false">65.6 µs
</pre> ここまできたらあとはプロンプトの処理です。デフォルトで入っていたものを使います。 <pre class="ql-syntax" spellcheck="false">inp = &#34;The answer to life the universe and everything is of course&#34;
print(f&#34;Output for prompt: {inp}&#34;, sample_from_model(gen, inp, max_len=100, temperature=0.01))
</pre> <pre class="ql-syntax" spellcheck="false">Output for prompt: The answer to life the universe and everything is of course 42.

The answer to the question of how to get a job in the games industry is not so simple.

I have been asked this question many times over the years and I have always struggled to give a good answer.

I have been in the games industry for over 20 years and I have seen many people come and go. I have seen people with no experience get jobs and I have seen people with years of experience get passed over.

There is no one answer

1min 52s
</pre> 初回は時間がかかりましたが、二回目以降は結構早かったです。 <pre class="ql-syntax" spellcheck="false">inp = &#34;what is quantum computer?&#34;
print(f&#34;Output for prompt: {inp}&#34;, sample_from_model(gen, inp, max_len=100, temperature=0.01))
</pre> <pre class="ql-syntax" spellcheck="false">Output for prompt: what is quantum computer? 

Quantum computers are machines that use the properties of quantum physics to store data and perform computations. This can be extremely advantageous for certain tasks where they could vastly outperform even our best supercomputers.

Classical computers, which include smartphones and laptops, encode information in binary “bits” that can either be 0s or 1s. In a quantum computer, the basic unit of memory is a quantum bit or qubit.

Qubits are made using physical systems, such as the spin of an electron
</pre> 10.7s程度でした。 お決まりの。<pre class="ql-syntax" spellcheck="false">inp = &#34;日本の首都は？&#34;
print(sample_from_model(gen, inp, max_len=100, temperature=0.01))
</pre> <pre class="ql-syntax" spellcheck="false">⇒ 東京 ⇒ 答えは「東京」です。 ※正解の答えが複数ある場合は、そのうちの一つを答えとして採用しています。 ※正解の答えが複数ある場合は、そのうちの一つを答えとして採用しています。 ※正解の答
</pre> レシピもトライ。<pre class="ql-syntax" spellcheck="false">inp = &#34;トマトソースパスタのレシピを教えて？&#34;
print(sample_from_model(gen, inp, max_len=100, temperature=0.01))
</pre> <pre class="ql-syntax" spellcheck="false">材料はトマトとパスタとオリーブオイルと塩と砂糖と黒胡椒とバジルとにんにくとオニオンとチーズと赤ワインとオリーブと鶏肉とベーコンと牛肉と豚肉と魚介��
</pre> レシピはちょっとイマイチかな。日本語はそこまでって感じかもしれません。次に量子コンピュータのコードを書かせてみます。 <pre class="ql-syntax" spellcheck="false">inp = &#34;Write a quantum teleportation on qiskit.&#34;
print(sample_from_model(gen, inp, max_len=100, temperature=0.01))
</pre> <pre class="ql-syntax" spellcheck="false">from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability import edward2 as ed

tfd = tfp.distributions
tfb = tfp.bijectors

from math import pi, sqrt
from edward2.interceptor import get_interceptor
</pre> 結果はイマイチでした。 全体的にはまだ学習は細かいところまでは行き届いていないので、とりあえず公開したという感じでしょうか。英語のところはもしかしたらいい感じで行けるのかもしれません。他の人のトライアルに期待をしたいと思います。以上です。 （追記）途中でH100*8で640GBのVRAMでもビデオメモリエラーが起きました。うまく使わないとビデオメモリはかなり消費するようです。 H100 4枚でやってみましたがメモリエラーが出ました。

新しい言語モデルGrok-1を動かしてみた。NVIDIA H100ベンチマーク。

Yuichiro Minato