gpt-neox-20b を 3090 x 2 で動かすメモ(3090 x 1 でも動く!)

ChatGPT(GPT-3.5?) しゅごいよねぇ...
ローカルで(GPU で)似たようなの動かしたい...

追記: 2023/06 時点ですと 8bit 量子化が成熟してきていてで 3090 x 1 で問題なく動きました. また日本語ですと rinna 3.6B がいい感じでした. https://zenn.dev/syoyo/articles/946c17666e10fb

LangChain + GPT-NEOX-Japanese-2.7b で日本語 LLM やりとり整備するメモ
https://qiita.com/syoyo/items/d0fb68d5fe1127276e2a

GPT-NEOX-Japanese-2.7b 試したけどなんか微妙だったネ...

とりま LLM では大きさは正義と思いますので, gpt-neox-20b 試します.

https://www.infoq.com/jp/news/2022/04/eleutherai-gpt-neox/
https://arxiv.org/abs/2204.06745v1

gpt-neox-20b は元モデルサイズが 43 GB くらい(fp16 でのウェイト)ありますが,
bitsandbytes(LLM.int8)で INT8 化 + 3090 x 2 GPUs に分割で 30 GB くらいのメモリで動くようになるやりかたがありました!

ありがとうございます.

一応 3090 でも動くのを確認しました. ただ長い文章あたえるとダメかもしれません.
(一応 gpt-neox-20b の最大トークン 2048 tokens くらいはいけそうであった)

環境

• 3090 x 2
• CPU mem 128 GB
• CUDA 12.0

とりあえず text-generation

from transformers import logging
logging.set_verbosity(logging.ERROR)

import torch
from accelerate import init_empty_weights, infer_auto_device_map, load_checkpoint_and_dispatch
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, logging
torch.cuda.empty_cache()
import os,shutil,json,random,truecase,datetime,wget,tarfile,pathlib
from fuzzywuzzy import fuzz


model_name = "EleutherAI/gpt-neox-20b"


config = AutoConfig.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

#
# weight 読み込まず, device_map 構築用にモデルを構築
#
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)


#
# 明示的に GPU の利用可能メモリを設定
# dtype=torch.int8 で 8bit 精度
# (内部的には bitsandbytes を呼んでいる模様?)
#
# GPU 1 枚で一部を CPU で計算でよければ,
# infer_auto_device_map で設定ではなく,
# `from_pretrain` のところで `device_map="auto", load_in_8bit=True` でいいかもしれません
# 
#
device_map = infer_auto_device_map(
                                  model, #using the empty model reference here saves time and RAM
                                  #change max_memory to suit your setup
                                  max_memory={0: "19GiB",1: "19GiB", "cpu": "90GiB"},
                                  #no_split_module_classes tells accelerate what sections of the model shouldn't be split during optimization
                                  no_split_module_classes=["GPTNeoXLayer"],
                                  #specify the doubly-shrunk weights' data type
                                  dtype=torch.int8
                              )

model = AutoModelForCausalLM.from_pretrained(
                                  model_name,
                                  device_map=device_map,
                                  config=config,
                                  load_in_8bit=True,
                              )


query_in = "..."

input_ids = tokenizer(query_in, return_tensors="pt").input_ids
input_token_length=input_ids.shape[1]
input_ids = input_ids.to('cuda')
output_tokens = model.generate(
        input_ids,
        do_sample=True,
        temperature=0.9,
        max_length=input_token_length+256,
    )
output_text = tokenizer.batch_decode(output_tokens)[0]
print(output_text)

日本語も AutoTokenizer でそのままいけます!
(Japanese 用 Tokenizer 使う必要がない)

日本語GPT（abeja/gpt-neox-japanese-2.7b）での文章生成を見て大規模コーパスの中身を想像する
https://note.com/oshizo/n/nd9af0d311810

ありがとうございます.

財務相 prompt 試します.

財務相は２６日、投機的な動きを背景に急激な為替変動が続けば「非常に不安定」な状況になるとして「非常に慎重」と言明した。

一方、米商務省は「市場過熱は続いている」との見通しを示している。

黒田総裁は２６日の閣議後会見で、「大規模な携帯電話やＩＴの開発が進み、これが市場の動きを加速させている」と指摘。「金融システムの調整が求められている」と述べ、企業に買い物のタイミングを見直して買ったり、売ったりするよう勧めた。<|endoftext|>

Super cooo! EOF token がでちゃいましたけどまあ後処理で suppress できますネ.

ずんだもん応答を試します!

prompt = """ずんだもんは東北に住む活発でかわいい女の子です。

    あなた「きみの名前は？」
    ずんだもん「ボクの名前はずんだもんなのだ。」

    あなた「今日の天気は？」
    ずんだもん「ボクは今日の天気は晴れだと思うのだ」

    あなた「今日は何時に起きたの？」
    ずんだもん「ボクは7時に起きたのだ」

    あなた「年齢は？」
    ずんだもん「ボクの年齢は秘密なのだ」

    あなた「朝ごはんは何食べたい？」
    ずんだもん「ボクはおにぎり食べたいのだ」

    あなた「昼ごはんは何食べたい？」
    ずんだもん「"""

ずんだもんは東北に住む活発でかわいい女の子です。

    あなた「きみの名前は？」
    ずんだもん「ボクの名前はずんだもんなのだ。」

    あなた「今日の天気は？」
    ずんだもん「ボクは今日の天気は晴れだと思うのだ」

    あなた「今日は何時に起きたの？」
    ずんだもん「ボクは7時に起きたのだ」

    あなた「年齢は？」
    ずんだもん「ボクの年齢は秘密なのだ」

    あなた「朝ごはんは何食べたい？」
    ずんだもん「ボクはおにぎり食べたいのだ」

    あなた「昼ごはんは何食べたい？」
    ずんだもん「ボクはお餅食べたいのだ。」

    あなた「夜ごはんは何食べる？」
    ずんだもん「ボクはパスタ食べたいのだ。」

    あなた「あなたは何歳？」
    ずんだもん「ボクは24歳です。ボクは日本人なのだ。」

    あなた「あなたの家はあなたのお父さんのもんだと言ってくれた。」
    ずんだもん「そう、いいんですか？でも、あなたのお父さんはいないので、お父さんのもんのお母さんです。」

    あなた「あなたのお母さんは今、どこにいるの？」
    ずんだもん「お母さんは昼ごはん作って、パソコンを使ってるの。」

    あなた「ボクはどこにいるの？

Ultra super cooooool! (最後切れているのは max_length 超えでしょう)
(ずんだもんさん, AI により年齢が 24 才というのもばれて(?)しまう(?) )

GPT-NEOX-Japanese-2.7b ではうまく埋めてくれませんでしたが, gpt-neox-20b はいい感じになりました!
(パラメータの設定次第というのもあるかもであるが)

メモリ消費

ふぇぇ... gpt-neox-20b, bitsandbytes(8bit int) で 3090 x 2 で動いたよぉ... 🥺🥺🥺🥺✌️✌️✌️✌️✌️🌸🌸🌸🌸🌸💖💖💖💖💖 pic.twitter.com/JV2zGYUs1Q

— syoyo.eth 🌸 レイトラ ® 🐯 4 周年 🎉 (@syoyo) January 8, 2023

メモリ消費は 21 GB + 4 GB という感じでした. ギリ 3090 x 1 でも動作しました(CUDA_VISIBLE_DEVICES=0 で一個に限定)が,
長い tokens は生成できないと思われるのでやっぱり 28 GB くらいは必要で GPU 二台構成(もう一方はいくらかスペック低くてもよいと思われる. 3090 + 3070 とか)にするのがよさそです!

4080 16 GB x 2 でも動くと思われます.
INT8 サポートは Turing 以降なので, GPU マイニングとかで 3070 とか余っている方は 3070 x 4(8 GB x 4 = 32 GB)でも動くかと思われます.

bitsandbytes が ROCm 対応すれば, RX6800(16GB) x 2 でも動くと思われます.

あとは huggingface accelerate(ややこしい名前...) で GPU に入らない部分を CPU にオフロードさせる手もあるでしょう.

日本語?

The Pile データセットのレポートを見ると,

そこそこのウェイトを占める Pile CC データセットから jusText では日本語や中国語は扱えないので, それらの言語は除いたとあります.

それでも GPT-NEOX-Japanese-2.7b とかよりもいいかんじにふるまう感はあります.

「人とAIが協調するためには、」で text-generation したら以下のような結果が得られました!

人とAIが協調するためには、AI開発者と人間の間に人間界からプラットフォームを開発する必要がある。

「プラットフォームは、コンピュータプログラムの組み立て時に行われる汎用的な仕組みの開発の手助けを行なうんだ」と
Zooxが説明する。

Zooxは2017年の10月、自律走行車のプラットフォームに「Zoox共同創業者のチャン・ファン・フーンが開発した、高度なAI
技術を活用し、自律走行車を操作する方法を、大学の研究者やデータサイエンティストの仕事を手伝った。だが、これがAI
や自律走行車などと協調して動作するプラットフォームなのだ」とツイートした。

Zooxは、自律走行車を支援するプラットフォームを次に紹介する。Zooxによると

テキスト生成とかは gpt-neox-20b に任せて,
日本語特化の処理は, 前処理や後処理で Ginza や Sudachi などでもいいのかもしれませんネ.

チャットとかしたい!

LangChain とかと組み合わせてチャットとかできるようにしたいネ.

のあとのほうを参考にしてプロンプトエンジニアリングをがんばればいけるっぽいが...

あと, 最近(?) LangChain では LLM backend で gpt-neox-20b 選べるようになっています.

また, chat データで訓練した gpt-next-20b が出てました: https://huggingface.co/togethercomputer/GPT-NeoXT-Chat-Base-20B

ただ試したところ日本語は微妙でした.

<人間>: こんにちは.
<ボット>: Hi. How may I help you today?
<人間>: 何か買い物に行きたいんですが、どこがおすすめですか
<ボット>: どこで買い物をするべきですか？
<人間>: すみません、先日おばあちゃんに会っていただいたばあちゃんとかわりに
<ボット>: すみません、先日おばあちゃんに会っていただいたばあちゃんとかわりにどうぞ
<人間>: こちらです
<ボット>: こちらの方に向かいます
<人間>: こちらの方に向かいます
<ボット>: どうぞ
<人間>: こちらがお

日本語データセットでﾌｧｲﾝﾁｭｰﾝすれば改善するかもです.

Fine tune は...?

一応 A100(40 GB) x 8 でいけるっぽいが...

268 GB 必要なのかしらん

fine tune は 推論 * 1.5 ~ 2.0 のメモリ(i.e. int8 なら 30 ~ 50 GB)で行けると思われます.
ある程度時間かかってもいいのであれば, メモリ積んだ CPU で fine tune という手もあるかもです.
CPU だと 16 cores CPU(e.g. Ryzen 3950x) で GPU より 10 倍遅いくらいと思われるので, fine tune が GPU で 1 ~ 2 時間なのであれば, CPU で一日でできる計算になる.

あとは DeepSpeed で nvme 積んだりで頑張るか, petals で分散処理も考えられます. https://github.com/bigscience-workshop/petals

QLoRA ﾌｧｲﾝﾁｭｰﾝで 3090 x 1 でいけそうです.

TODO

• メモリ管理あたりもう少し明示的に設定できるようにしてみる
  • huggingface transformers も accelerate も, ぺろっと試すにはいいが, いろいろ手を加えたり微調整したいとなるといろいろ抽象化していてやりずらい...
• ChatGPT みたいなチャットボットをつくる
• fine-tune ためす
• きたる GPT-4 レベル(1T とか?)のモデルを動かせるように 4090 65,536 枚くらい調達しておく