こんにちは、CCCMKホールディングス TECH LABの三浦です。
この頃はとても暑い日が続いています。年を取るにつれ、だんだん暑いのが苦手になってきたように思います。日差しが強い中外を歩くと、明らかに身体にダメージを受けているような気がします。出来るだけ日陰を歩くようにしたり、帽子をかぶるようにして、暑さを少しでも抑えるようにしないとですね。
先日参加したDatabricksのDATA+AI SUMMITでは大規模言語モデル(LLM)に関するトピックがたくさんあり、特にオープンソースのLLMについてはこれからどんどん活用していきたい、と考えています。なのですが、私がこれまで使ってきたNVIDIA V100-16GBの環境では推論処理だけでも動かすのが難しいサイズのオープンソースのLLMもあります。
これまで色々試してきた感じですと、大体10B(100億)くらいのパラメータ数のLLMだとGPUメモリで処理が出来て、それよりも大きいLLMになるとGPUメモリに乗り切らずにOut Of Memory(OOM)エラーが出てしまうような印象があります。
どんな出力が得られるのかちょっと確認してみたいな、という時にOOMエラーが出るとそれ以上手が打てなかったのですが、Microsoftが開発した深層学習の処理を最適化するライブラリDeepSpeedを使うことで、ある程度今の環境でも大きなLLMの推論処理を走らせることが出来ました。
今回はその内容について、ご紹介したいと思います。
DeepSpeedのZeRO-Inference
DeepSpeedはMicrosoftが開発した深層学習処理の最適化ライブラリです。
特に巨大なサイズの深層学習モデルの最適化を目的とした様々な機能が含まれています。例えばGPUメモリに乗り切らないモデルのパラメータをCPUメモリに一度格納しておき、必要な部分だけGPUメモリにコピーして高速に処理させるといったCPU-Offloadや、マルチGPUでモデルのパラメータや勾配などの情報を単純にコピーして持たせるのではなく、分割してそれぞれのGPUメモリに持たせるZeRO(Zero Redundancy Optimizer)などを使用することが出来ます。
それらによって、大規模な深層学習モデルの学習と推論を少ない計算リソースで効率的に行うことが可能です。推論においては今回は"ZeRO-Inference"という機能を使用しましたが、それ以外にも"DeepSpeed-Inference"という機能もあります。
DeepSpeed-Inferenceは複数のGPUでモデルを分割してメモリに保持することで、1つのGPUメモリでは乗り切らないようなモデルを扱うことを可能にします。一方ZeRO-Inferenceの方はCPUメモリを併用することでGPUメモリに乗り切らないモデルを扱うことを可能にしています。
チュートリアルなどを参考に使ってみる
ひとまずDeepSpeedを使ってHuggingfaceに公開されているMosaicMLの"mosaicml/mpt-30b"というモデルを動かしてみました。
今回参考にしたドキュメントをご紹介します。まずZeRO-InferenceについてはDeepSpeedのドキュメントを参考にしました。
HuggingfaceのモデルをZeRO-Inferenceで動かす方法についてはHuggingface Transformersのドキュメントを参考にしました。
実行環境
DatabricksのRuntime 13.1 for Machine Learningで、V100-16GBを2基積んだ環境を使用しました。
DeepSpeedのセットアップ
必要なcudaのライブラリと、libaio-devをインストールします。必要なcudaライブラリについては、Databricksのdollyのレポジトリにあるtrain_dolly.pyを参考にしました。
!wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/libcusparse-dev-11-7_11.7.3.50-1_amd64.deb -O /tmp/libcusparse-dev-11-7_11.7.3.50-1_amd64.deb && \ wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/libcublas-dev-11-7_11.10.1.25-1_amd64.deb -O /tmp/libcublas-dev-11-7_11.10.1.25-1_amd64.deb && \ wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/libcusolver-dev-11-7_11.4.0.1-1_amd64.deb -O /tmp/libcusolver-dev-11-7_11.4.0.1-1_amd64.deb && \ wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/libcurand-dev-11-7_10.2.10.91-1_amd64.deb -O /tmp/libcurand-dev-11-7_10.2.10.91-1_amd64.deb && \ dpkg -i /tmp/libcusparse-dev-11-7_11.7.3.50-1_amd64.deb && \ dpkg -i /tmp/libcublas-dev-11-7_11.10.1.25-1_amd64.deb && \ dpkg -i /tmp/libcusolver-dev-11-7_11.4.0.1-1_amd64.deb && \ dpkg -i /tmp/libcurand-dev-11-7_10.2.10.91-1_amd64.deb && \ apt install libaio-dev -y
続いてpipを使用してdeepspeedをインストールします。また、"mosaicml/mpt-30b"を動かすために必要になる"einops"というライブラリもインストールします。
pip install deepspeed einops
ZeRO-Inferenceの設定ファイル
ZeRO-Inferenceの設定ファイルを用意します。ZeROには3つの最適化に関するstageがありますが、モデルのパラメータを最適化するstage3の設定ファイルをベースにします。
%%writefile zero_infer.json
{
"fp16": {
"enabled": "auto",
"loss_scale": 0,
"loss_scale_window": 1000,
"initial_scale_power": 16,
"hysteresis": 2,
"min_loss_scale": 1
},
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"offload_param": {
"device": "cpu",
"pin_memory": true
},
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto",
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_gather_16bit_weights_on_model_save": true
},
"steps_per_print": 2000,
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": false
}
Pythonスクリプト
次にモデルをダウンロードして推論(テキスト生成)を行うPythonスクリプトです。
%%writefile deepspeed_zeroinfer.py import torch import deepspeed from transformers.deepspeed import HfDeepSpeedConfig from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig import json import time import argparse import os parser = argparse.ArgumentParser(description="Using DeepSpeed ZeRO Inference") # Deepspeedが使用するargument parser.add_argument("--local_rank", type=int, default=-1, help="local rank passed from distributed launcher") # ここからこのスクリプト専用のargument parser.add_argument("model_name", type=str, help="set model name for Huggingface Models") parser.add_argument("input_text", type=str, help="set text to input model") args = parser.parse_args() # argumentから変数へ model_name = args.model_name # ロードするhuggingface model input_text = args.input_text # モデルに入力する文字列 # multi-GPU関連の設定 os.environ["TOKENIZERS_PARALLELISM"] = "false" # To avoid warnings about parallelism in tokenizers local_rank = int(os.getenv("LOCAL_RANK",0)) world_size = int(os.getenv("WORLD_SIZE",1)) torch.cuda.set_device(local_rank) deepspeed.init_distributed() # ベースとなるZeRO3 configの読み込み ds_config_file = "zero_infer.json" with open(ds_config_file) as f: ds_config = json.load(f) # 推論用に修正 model_config = AutoConfig.from_pretrained(model_name) hidden_size = model_config.hidden_size ds_config["train_batch_size"] = 1 * world_size ds_config["train_micro_batch_size_per_gpu"] = 1 ds_config["reduce_bucket_size"] = hidden_size*hidden_size ds_config["stage3_prefetch_bucket_size"] = 0.9 * hidden_size * hidden_size ds_config["stage3_param_persistence_threshold"] = 10 * hidden_size dschf = HfDeepSpeedConfig(ds_config) #zero3を使用するために必要(モデルロード前に実行する必要がある) model = AutoModelForCausalLM.from_pretrained(model_name) tokenizer = AutoTokenizer.from_pretrained(model_name) print("【deepspeed initialize start】") start_time = time.time() ds_engine = deepspeed.initialize(model=model, config_params=ds_config)[0] ds_model = ds_engine.module.eval() init_time = time.time() - start_time print("【deepspeed initialize end】") print(f"initializing took {init_time}sec") input_token = tokenizer( input_text, add_special_tokens=False, return_tensors="pt").to(model.device) with torch.no_grad(): start_time = time.time() outputs = ds_model.generate( **input_token, max_new_tokens=64, do_sample=True, temperature=0.7, top_p=0.9, repetition_penalty=1.05, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id) generate_time = time.time() - start_time gen_text = tokenizer.decode(outputs[0],skip_special_tokens=True) result = { "local_rank,":local_rank, "generate_time(s)":format(generate_time,".2f"), "gen_text":gen_text } print(result)
DeepSpeedを稼働させる
上記で作成したPythonスクリプトファイルを指定して、DeepSpeedを動かします。--num_gpusでは使用するGPUの数を指定することが出来ます。
!deepspeed --num_gpus 2 deepspeed_zeroinfer_for_mpt30b.py "The highest mountain in Japan is"
出力結果は以下にようになりました。
{'local_rank,': 0, 'generate_time(s)': '858.48', 'gen_text': "The highest mountain in Japan is Mount Fuji, with a height of 3776 meters. The second-highest peak is Kita-dake (3084 m) on the border between Nagano and Yamanashi Prefectures; it forms part of the Northern Japanese Alps range that includes most of Honshu's major peaks above 3000"}
{'local_rank,': 1, 'generate_time(s)': '858.71', 'gen_text': "The highest mountain in Japan is Mount Fuji.\n\n**futur** **e** _adj_ futur, à venir \n **What time do you have to be at the airport?** — **Quand est-ce que vous devez être à l'aéroport?**\n\n**future** _n"}
ご覧の通り、2つの結果が出力されています。これは2台のGPUによって並列で処理が行われたためです。ZeRO-Inferenceのメリットとして複数の入力データをそれぞれのGPUに分散して処理出来る点があります。今回は入力が1件だけなのでそのメリットを感じられませんが、多数のデータをバッチで処理する際などには大きな効果があると思います。
1つの入力に対して大体14分くらいの時間がかかりました。設定ファイルを調整すると、もう少し改善できるのかもしれません。
実はHuggingfaceのAccelerateでも同じことが出来る・・・？
DeepSpeedを使わない場合はどれくらいの処理時間がかかるのかを見ておこうと思い、DeepSpeedを使用しないスクリプトも用意して実行してみました。しかしDeepSpeedを使っていないにも関わらず、GPUメモリが溢れることなく処理が完了してしましまいました・・・(しかもちょっと処理時間も速い)。
{'generate_time(s)': '732.21', 'gen_text': 'The highest mountain in Japan is Mount Fuji, at 3776 metres (12388 feet).\n\n## **Where should I stay?**\n\nThere are lots of different types of accommodation available. You can choose to be right on the beach or high up with a view over town; you may want an apartment so that your kids have their'}
この原因が分からずしばらく悩んでいたのですが、HuggingfaceのAccelerateというライブラリによって実現されていることが分かりました。
DatabricksのRuntime 13.1 for Machine LearningにはAccelerateが予めインストールされており、しかもHuggingfaceのTransformersで事前学習済みモデルをロードする時、
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map="auto")
のようにdevice_map="auto"と指定すればAccelerateによってモデルのパラメータの配置を最適に行ってくれるようです。AccelerateとDeepSpeedは互換性もあるようです。
まとめ
ということで、今回は巨大なサイズのLLMを出来るだけ少ない計算リソース上で動かす、ということにチャレンジしてみた話をご紹介させて頂きました。LLMに関する技術は今まさにどんどん生み出されている時期で、なかなかキャッチアップが大変だな・・・と心から感じました。HuggingfaceのAccelerateについてもちゃんと調べておいた方が良さそうだと思います。また、もう少し処理速度を改善できる方法についても調べていこうと思います。
