DeepSpeedとAccelerateを使用した非常に高速なBLOOM推論

High-speed BLOOM inference using DeepSpeed and Accelerate.

この記事では、176BパラメータのBLOOMモデルを使用してトークンごとのスループットを非常に高速に取得する方法を紹介します。

モデルは352GBのbf16（bfloat16）ウェイト（176*2）を必要とするため、最も効率的なセットアップは8x80GBのA100 GPUです。また、2x8x40GBのA100または2x8x48GBのA6000も使用できます。これらのGPUを使用する主な理由は、この執筆時点ではこれらのGPUが最大のGPUメモリを提供しているためですが、他のGPUも使用できます。たとえば、24x32GBのV100を使用することもできます。

単一のノードを使用すると、通常、最速のスループットが得られます。なぜなら、ほとんどの場合、ノード内のGPUリンクハードウェアの方がノード間のものよりも速いためですが、常にそうとは限りません。

もしハードウェアがそれほど多くない場合でも、CPUやNVMeのオフロードを使用してBLOOM推論を実行することは可能ですが、もちろん、生成時間は遅くなります。

• 🤗 Accelerateは、PyTorchのおかげで非常に大規模なモデルを実行する方法です
• AutoTrainによる画像分類
• 日本語安定拡散

また、GPUメモリの半分の容量を必要とする8ビット量子化ソリューションについても説明します。これにはBitsAndBytesとDeepspeed-Inferenceライブラリが必要です。

ベンチマーク

さらなる遅延なしでいくつかの数値を示しましょう。

一貫性を保つために、この記事のベンチマークはすべて同じ8x80GBのA100ノードで実行され、512GBのCPUメモリを持つJean Zay HPCで行われました。JeanZay HPCのユーザーは、約3GB/sの読み取り速度（GPFS）で非常に高速なIOを利用しています。これはチェックポイントの読み込み時間に重要です。遅いディスクは読み込み時間が遅くなります。特に複数のプロセスでIOを同時に行っている場合はさらに重要です。

すべてのベンチマークは、100トークンの出力を貪欲に生成しています：

Generate args {'max_length': 100, 'do_sample': False}

入力プロンプトはわずかなトークンで構成されています。以前のトークンのキャッシュもオンになっています。常にそれらを再計算すると非常に遅くなるためです。

まず、生成の準備が完了するまでにかかった時間（つまり、モデルの読み込みと準備にかかった時間）を見てみましょう：

Deepspeed-Inferenceには、事前にシャードされたウェイトリポジトリが付属しており、読み込みに約1分かかります。Accelerateの読み込み時間も優れており、わずか2分です。他のソリューションはここでははるかに遅いです。

読み込み時間は重要であるかどうかは、一度読み込んだら追加の読み込みオーバーヘッドなしに繰り返しトークンを生成できるため、場合によります。

次に、トークン生成の最も重要なベンチマークです。ここでのスループット指標は単純であり、100個の新しいトークンを生成するのにかかった時間を100で割り、バッチサイズで割ったものです。

ここでは8x80GBのGPU上でのスループットがmsec単位で表示されています：

ここで、OOMはGPUメモリに収まりきらない大きすぎるバッチサイズによるメモリ不足（Out of Memory）の状態を表します。

Deepspeed-Inferenceのテンソル並列性（TP）とカスタム統合CUDAカーネルを使用して、1ミリ秒未満のスループットを実現することができます！これは本当に驚くべきことです！ただし、他のモデルに対してこのソリューションを使用する場合、それを動作させるために開発者の時間が必要になる場合があります。

Accelerateも非常に高速です。これは単純なパイプライン並列性（PP）のアプローチを使用しており、非常にシンプルなため、どのモデルでもそのまま動作するはずです。

Deepspeed-ZeROは複数の生成ストリームを並列に処理できるため、スループットは8または16でさらに分割できます。8または16のGPUがgenerate呼び出し中に使用された場合、バッチサイズ64を処理できるため、スループットは約4ミリ秒となり、3つのソリューションは非常に近いです。

これらの数値がどのように計算されたかをもう一度見直してみましょう。バッチサイズ128の場合、fp16モードでDeepspeed-Inferenceを使用して100個の新しいトークンを生成するために実時間で8832ミリ秒かかりました。したがって、スループットはwalltime/(batch_size*new_tokens)または8832/(128*100) = 0.69となります。

今度は、Deepspeed-InferenceとBitsAndBytesによって提供される量子化されたint8ベースのモデルのパワーを見てみましょう。これにより、bfloat16またはfloat16の推論に比べて、元のGPUメモリの半分しか必要ありません。

スループット（4x80GB A100）：ミリ秒単位

ベンチマークの結果を再現するには、以下で説明する3つのスクリプトのいずれかに--benchmarkを追加してください。

ソリューション

まず、デモリポジトリをチェックアウトしてください：

git clone https://github.com/huggingface/transformers-bloom-inference
cd transformers-bloom-inference

この記事では、bloom-inference-scripts/配下にある3つのスクリプトを使用します。

フレームワーク固有のソリューションはアルファベット順に表示されます：

HuggingFace Accelerate

Accelerate

Accelerateは、次の方法で推論のための大きなモデルを処理します：

1. 空の重みでモデルをインスタンス化します。
2. 各レイヤーのサイズとデバイス（GPU、CPU）上の利用可能なスペースを分析し、各レイヤーを配置するデバイスを決定します。
3. モデルのチェックポイントをビット単位でロードし、各重みをそのデバイスに配置します。

それは、入力と出力を正しいデバイスに転送し、モデルの重みがCPU（またはディスク）にオフロードされる際には、転送のフックを使用してモデルが正常に実行されることを保証します。また、順伝播が完了すると、重みは再度オフロードされます。

複数のGPUにモデル全体を格納する十分なスペースがある場合、すべてのレイヤーが実行されるまで、制御を1つのGPUから次のGPUに切り替えます。一度に1つのGPUのみが動作するため、GPUのアイドリングにもかかわらず、合理的なスループットを実現します。

また、同じコードが任意のセットアップで実行できるため、非常に柔軟です。Accelerateは最初に使用可能なすべてのGPUを使用し、次にRAMがいっぱいになるまでCPUにオフロードし、最後にディスクにオフロードします。CPUやディスクへのオフロードは処理を遅くします。例えば、ユーザーはBLOOMをコードの変更なしで単に2つのA100で実行し、スループットは8×80 A100では10ミリ秒に対して1トークンあたり15秒になったと報告しています。

このソリューションの詳細については、Accelerateのドキュメントを参照してください。

セットアップ

pip install transformers>=4.21.3 accelerate>=0.12.0

実行

シンプルな実行方法は次の通りです：

python bloom-inference-scripts/bloom-accelerate-inference.py --name bigscience/bloom --batch_size 1 --benchmark

BitsAndBytesから8ビットの量子化ソリューションをアクティベートするには、まずbitsandbytesをインストールしてください：

pip install bitsandbytes

そして、前のコマンドラインに--dtype int8を追加してください：

python bloom-inference-scripts/bloom-accelerate-inference.py --name bigscience/bloom --dtype int8 --batch_size 1 --benchmark

4つ以上のGPUがある場合、次のようにして4つだけを使用するように指定できます：

CUDA_VISIBLE_DEVICES=0,1,2,3 python bloom-inference-scripts/bloom-accelerate-inference.py --name bigscience/bloom --dtype int8 --batch_size 1 --benchmark

この場合、OOMが発生しない最大のバッチサイズは40です。スクリプト内を見ると、最初のGPUをアクティベーションと以前のトークンのキャッシュのみを処理するようにメモリ割り当てマップを調整する必要がありました。

DeepSpeed-Inference

DeepSpeed-Inferenceは、テンソルパラレリズムと効率的な融合CUDAカーネルを使用して、大規模なバッチサイズ（128）で1トークンあたりの推論時間が<1ミリ秒という超高速な推論を実現します。

セットアップ

pip install deepspeed>=0.7.3

実行

1. 最速の方法は、TP-pre-sharded（TP = Tensor Parallel）チェックポイントを使用することです。これは非プリシャードのbloomチェックポイントに比べて約1分しかかからず、非常に高速です：

deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py --name microsoft/bloom-deepspeed-inference-fp16

1a. もし元のブルームチェックポイントを実行したい場合、読み込みに10-20分かかりますが、以前のソリューションと同じスループットで実行されます。

deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py --name bigscience/bloom

2a. 8ビット量子化バージョンでは、通常の半精度バージョンのGPUメモリの半分しか必要ありません。

deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py --name microsoft/bloom-deepspeed-inference-int8 --dtype int8

ここでは microsoft/bloom-deepspeed-inference-int8 を使用し、スクリプトを int8 で実行するように指定しました。

そしてもちろん、4つの80GB A100 GPUが十分です:

deepspeed --num_gpus 4 bloom-inference-scripts/bloom-ds-inference.py --name microsoft/bloom-deepspeed-inference-int8 --dtype int8

この場合、OOMなしで実行できる最大バッチサイズは128です。

ここで、より良いパフォーマンスにつながる2つの要素が見られます。

1. ここでは、Accelerateのパイプライン並列処理（PP）ではなく、Tensor Parallelism（TP）を使用してスループットが向上しました。Accelerateは非常に汎用的であるため、GPUの使用率を最大化するのは困難です。すべての計算はまずGPU 0で行われ、次にGPU 1、などと進みます。つまり、7つのGPUは常にアイドル状態です。一方、DeepSpeed-InferenceはTPを使用しており、テンソルをすべてのGPUに送信し、各GPUで生成の一部を計算し、その結果をすべてのGPU同士で通信し、次の層に移ります。これにより、すべてのGPUが同時にアクティブになりますが、通信量が多くなります。

2. DeepSpeed-Inferenceはまた、メモリの割り当てを抑え、GPU間でのテンソルのコピーを行わないためのカスタムCUDAカーネルを使用しています。これにより、メモリ要件が少なくなり、カーネルの起動回数が減少し、スループットが向上し、より大きなバッチサイズが可能になり、全体的なスループットが向上します。

さらに例に興味がある場合は、GPU上のDeepSpeed-Inferenceを使用したAccelerate GPT-J推論またはAccelerate BERT推論をご覧ください。

Deepspeed ZeRO-Inference

Deepspeed ZeROは、ほとんどのモデルを数個または数百のGPUにスケールし、トレーニングまたは推論を行うことができる魔法のシャーディング手法を使用します。

Setup

pip install deepspeed

Run

現在のスクリプトはすべてのGPUで同じ入力を実行しますが、各GPUで異なるストリームを実行し、n_gpu 倍のスループットを得ることもできます。これはDeepspeed-Inferenceではできません。

deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-zero-inference.py --name bigscience/bloom --batch_size 1 --benchmark

ZeROでは、ユーザーは同時に複数のユニークなストリームを生成することができるため、全体的なパフォーマンスは、スループット（秒/トークン）を参加するGPUの数で割ったものです。したがって、8つまたは16個のGPUを使用する場合、8倍から16倍の速さで実行できます。

また、8つの大型GPUがない場合は、1つの小型GPUでオフロードソリューションを試すこともできますが、実行には長い時間がかかります。

CPUオフロード（1x GPUs）:

deepspeed --num_gpus 1 bloom-inference-scripts/bloom-ds-zero-inference.py --name bigscience/bloom --batch_size 8 --cpu_offload --benchmark

NVMeオフロード（1x GPUs）:

deepspeed --num_gpus 1 bloom-inference-scripts/bloom-ds-zero-inference.py --name bigscience/bloom --batch_size 8 --nvme_offload_path=/path/to/nvme_offload --benchmark

以下のHTMLコードの日本語訳です：

/path/to/nvme_offload を、高速なNVMeドライブ上で約400GBの空きメモリがある場所に調整してください。

その他のクライアントとサーバーのソリューション

transformers-bloom-inferenceには、サーバーソリューションを含む、より効率的なソリューションがさらにあります。

以下はいくつかのプレビューです。

サーバーソリューション：

• Mayank Mishraは、このブログ記事で説明されているすべてのデモスクリプトをウェブサーバーパッケージに変換しました。ダウンロードはこちらからできます。

• Nicolas Patryは、スーパーエフィシェントなRustベースのウェブサーバーソリューションを開発しました。

その他のクライアント側のソリューション：

• Thomas Wangは、非常に高速なカスタムCUDAカーネルBLOOMモデルを開発しています。

• The JAX team @HuggingFaceは、JAXベースのソリューションを開発しました。

このブログ記事は、公開後数か月経ってから読んでも古くなる可能性があるため、最新のソリューションを見つけるためにはtransformers-bloom-inferenceをご利用ください。

ブログのクレジット

以下の方々に、良い質問をして記事の読みやすさを向上させるのにご協力いただきました。Olatunji RuwaseさんとPhilipp Schmidさん、本当にありがとうございます。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!