NVIDIA の Claude 3.5 Sonnet 超え（？）の Llama-3.1-Nemotron-70B-Instruct を試す

2024年10月17日 02:39

tl;dr

NVIDIA から Llama-3.1-Nemotron-70B-Instruct が公開
3 つのベンチマークで Claude 3.5 Sonnet や GPT-4o を超える性能を誇る
- Arena Hard で 85.0、AlpacaEval で 57.6、MT-Bench で 8.98
最大入力トークン数は 128k で最大出力トークン数は 4k
Llama-3.1-70B-Instruct を初期方策として、Llama-3.1-Nemotron-70B-Reward と HelpSteer2-Preference プロンプトを用いて RLHF
NVIDIA NIM、transformers、Ollama、MLX から動作を試した
MLX の変換済みモデルを MLX-Community にアップロードした

モデルページとベンチマーク

論文はこちら。解説は岡野原さんの投稿がわかりやすい。

Llama-3.1-Nemotoron-70B-Rewardは新しく作られたアライメント用データセット HelpSteer2-preferenceを使い、回答品質を大きく改善。例えばArena-Hardではフロンティアモデルに続くOSSで1位を達成しているhttps://t.co/TgSN1jyNrG…
— Daisuke Okanohara / 岡野原大輔 (@hillbig) October 16, 2024

実際に動かしてみないとわからない。ということで動かしてみましょう。

NVIDIA NIM から動かしてみる

NVIDIA NIM で使えるようです。NVIDIA NIM は、LLM などの AI モデルを本番環境へデプロイすることを最適化、簡素化するためのマイクロサービスソリューション。うーむ、よくわからん。要はモデルをデプロイして、API にして使えるよ、ということのよう。気になる方は下記公式ページをご覧ください。動かしたくてうずうずしているので先に進みます。

開いたらすぐ左にプロンプトを打ち込むチャットがありますので、そちらにプロンプトを入れるだけ。

もう少し聞いてみましょう。

Input: 核融合発電について教えてください。

Output:
核融合発電について詳しく説明します。

核融合発電とは？
核融合発電は、原子核が融合してエネルギーを発生させる原理を利用した発電方法です。太陽や他の星々では、水素原子核（プロトン）がヘリウム原子核に融合することでエネルギーを発生させています。人工的にこのプロセスを再現し、発生したエネルギーを電気に変換するのが核融合発電です。

核融合反応の基本
最も研究が進んでいる核融合反応は、以下の2つの反応です：
D-T (重水素-トリチウム) 反応
反応式: D (重水素) + T (トリチウム) → α (ヘリウム核) + n (中性子) + エネルギー
特徴: 発生する中性子を利用して熱を生成し、熱交換器で蒸気を生み、蒸気タービンで発電します。
D-D (重水素-重水素) 反応
反応式: D + D → ³He (ヘリウム同位体) + n + エネルギー (または) D + D → T + p (陽子) + エネルギー
特徴: この反応も中性子を発生させ、熱を生成しますが、反応率はD-T反応に比べて低いです。

核融合発電のメリット
無限に近い資源: 使用する燃料（例：リチウム、重水素）は豊富に存在します。
CO2排出ゼロ: 発電過程で二酸化炭素（CO2）を排出しません。
長期的安定供給: 天候や時間帯に依存しない基幹電源として機能できます。
廃棄物の低放射性: 核分裂型原子力発電に比べて、半減期の短い廃棄物しか生成しないと考えられています。

核融合発電の課題
高温・高圧環境の維持: 核融合反応を維持するためにプラズマを約1億度の高温で維持する必要があります。
反応を自持した維持: これまで、核融合反応をエネルギー出力が入力エネルギーよりも大きい（ブレーキ・イーブン）状態で長時間維持できていません。
コスト削減と商用化: 現在の技術では、発電コストが高いままです。商業的に競争力のある発電方法としての実現が課題です。

現在の状況と将来展望
ITER (国際熱核融合実験炉) プロジェクト: フランスに建設中のITERは、2025年からプラズマ実験を開始予定。世界中の国々が協力して取り組んでいる重要プロジェクトです。
商業化へのステップ: ITER以降、DEMO（実証炉）を経て、商業炉の実現を目指しています。DEMOの運転開始は2050年代半ばを予定しています。
民間企業の参入: ゴーグルやビル・ゲイツが出資する企業を含む、多くの民間企業が核融合研究に参入しています。新技術の導入やコスト削減に期待が寄せられています。

核融合発電は、将来の持続可能なエネルギー源として非常に有望ですが、技術的課題の解決と商業化への道のりはまだ長い状況です。

NVIDIA NIM

若干翻訳がこなれていない以外はそこそこまとも。ただ、2025 年に ITER のプラズマ実験（ファーストプラズマの点火）はちょっと厳しいかな（情報がかなり古い）。また、プロトンとヘリウムではさすがに核融合反応は起こらない。

総じてハルシネーションの低減さえできれば日本語モデルとしても十分使えるような感触を持ちました。

ではローカル環境で動かしてみます。

Mac Studio 2023
Apple M2 Ultra
Memory 192GB

ローカル環境（macOS）

transformers で動かしてみる

環境構築。自分は手に馴染むので uv を使っていますがそこはよしなに。

mkdir playground-Llama-3.1-Nemotron-70B-Instruct-HF
cd playground-Llama-3.1-Nemotron-70B-Instruct-HF

uv init --python 3.11
uv add torch transformers accelerate

touch main.py

まずは公式サンプルをそのまま動かしてみたいですね。

# main.py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_name = "nvidia/Llama-3.1-Nemotron-70B-Instruct-HF"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)

prompt = "How many r in strawberry?"
messages = [{"role": "user", "content": prompt}]

tokenized_message = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True)
response_token_ids = model.generate(tokenized_message['input_ids'].cuda(),attention_mask=tokenized_message['attention_mask'].cuda(),  max_new_tokens=4096, pad_token_id = tokenizer.eos_token_id)
generated_tokens =response_token_ids[:, len(tokenized_message['input_ids'][0]):]
generated_text = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
print(generated_text)

私の環境は macOS ですので、macOS で動かしたい方は下記のように `device_map="mps"` に変更し、.cuda() の部分をto("mps") に変更してください。また、macOS は bfloat16 に対応していないので、float16 に変更することも忘れずに。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_name = "nvidia/Llama-3.1-Nemotron-70B-Instruct-HF"
# model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="mps")
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="mps")
tokenizer = AutoTokenizer.from_pretrained(model_name)

prompt = "How many r in strawberry?"
messages = [{"role": "user", "content": prompt}]

tokenized_message = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True)
response_token_ids = model.generate(tokenized_message['input_ids'].to("mps"),attention_mask=tokenized_message['attention_mask'].to("mps"),  max_new_tokens=4096, pad_token_id = tokenizer.eos_token_id)
generated_tokens =response_token_ids[:, len(tokenized_message['input_ids'][0]):]
generated_text = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
print(generated_text)

では実行しましょう。

uv run main.py

量子化なしの 70B LLM がローカル環境で推論できることを夢見た時代が私にもありました。使用メモリが 150GB を超えており、Unified Memory だとしても 192GB では有効 GPU 利用率が 75% ゆえに 144GB までしか GPU として使えません。（正確には限界突破の方法があったりするのですが、自分の環境でなぜか `Result too large` と出力され、設定できません）

仕方ありません。max tokens を 128 まで減らしてリトライです。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_name = "nvidia/Llama-3.1-Nemotron-70B-Instruct-HF"
# model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="mps")
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="mps")
tokenizer = AutoTokenizer.from_pretrained(model_name)

prompt = "How many r in strawberry?"
messages = [{"role": "user", "content": prompt}]

tokenized_message = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True)
response_token_ids = model.generate(tokenized_message['input_ids'].to("mps"),attention_mask=tokenized_message['attention_mask'].to("mps"),  max_new_tokens=128, pad_token_id = tokenizer.eos_token_id)
generated_tokens =response_token_ids[:, len(tokenized_message['input_ids'][0]):]
generated_text = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
print(generated_text)

> uv run main.py
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [01:10<00:00, 2.35s/it]
Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)
A sweet question!

Let's count the "R"s in "strawberry":

1. S
2. T
3. R
4. A
5. W
6. B
7. E
8. R
9. R
10. Y

There are **3 "R"s** in the word "strawberry".

たったこれだけの出力ですが 7 分 4 秒かかりました。悲しいですね！

ただ、結果を見てみてください。`strawberry` に `r` は 3 つなのでちゃんと答えられていますね。ちなみに ChatGPT も Claude も答えられません。これは冗談抜きにすごい。

Ollama で動かしてみる

ollama run nemotron

A 70B instruct model customized by @nvidia from @AIatMeta llama 3.1.

😍 We can count letters now. Let's go team!

👇👇👇https://t.co/FsMxeckOmr pic.twitter.com/ALoDNYu2Uz
— ollama (@ollama) October 16, 2024

Ollama 公式も対応しましたので試してみます。Ollama が起動している状態でまずはモデルのダウンロードをしましょう。

ollama pull nemotron

デフォルトでは Q4_K_M がダウンロードされます。自分の環境ではダウンロードにおおよそ 34 分程度かかりました。気長に待ちましょう。

ollama run nemotron --verbose "こんにちは。"

簡単にオプションを解説。`--verbose` オプションは推論にかかった時間やトークン数などの数値を見ることができます。ollama run 実行時に単に文字列を与えてあげることでプロンプトを指定することも可能です。

LLM 以前によく存在した分岐型のチャットボット感が漂っていますね！

MLX で動かしてみる

ここからは macOS をお持ちの方のみ。mlx-lm を使って推論させてみましょう。

Llama-3.1-Nemotron-70B-Instruct の MLX 変換版を MLX Community にアップロードしました。

とりあえず一番軽い 4bit 版のものを動かしてみましょう。レスポンスはまあこれまでと同じなので書きませんが、よかったら試してみてください。macOS で推論する場合 4bit だと 64GB あれば動くかと思います。

uv を使用する場合

uv run --with mlx-lm mlx_lm.generate --model mlx-community/Llama-3.1-Nemotron-70B-Instruct-HF-4bit --prompt "こんにちは。" --max-tokens 128 --temp 0.0

pip を使用する場合

pip install mlx-lm
python -m mlx_lm.generate --model mlx-community/Llama-3.1-Nemotron-70B-Instruct-HF-4bit --prompt "こんにちは。" --max-tokens 128 --temp 0.0

もし MLX に興味のある方がいらしたら ChatMLX というオープンソースの macOS 用アプリがあるので試してみてください。

以上となります。Claude 3.5 Sonnet に汎用性能で超えるかというと少し懐疑的には考えていますが、一部のタスクでは確かに超えているものもあるといった印象です。モデルのパラメータサイズはそこそこ大きいですが、NVIDIA NIM であればすぐ動かすこともできるので遊んでみてください。

直近これらのイベント↓に出没します！もしご都合の合う方がいらしたら遊びに来てください！