24GBのコンシューマーGPUでRLHFを使用して20B LLMを微調整する
使用されたのは、24GBのコンシューマーGPUを利用して、RLHF（Reinforcement Learning with Human Feedback）を使用して20B LLM（Large Language Model）を微調整しました
私たちは、trlとpeftの統合を正式にリリースし、Reinforcement Learningを用いたLarge Language Model (LLM)のファインチューニングを誰でも簡単に利用できるようにしました！この投稿では、既存のファインチューニング手法と競合する代替手法である理由を説明します。
peftは一般的なツールであり、多くのMLユースケースに適用できますが、特にメモリを多く必要とするRLHFにとって興味深いです！
コードに直接深く入りたい場合は、TRLのドキュメンテーションページで直接例のスクリプトをチェックしてください。
イントロダクション
LLMとRLHF
言語モデルとRLHF（Reinforcement Learning with Human Feedback）を組み合わせることは、ChatGPTなどの非常に強力なAIシステムを構築するための次の手段として注目されています。
RLHFを用いた言語モデルのトレーニングは、通常以下の3つのステップを含みます：
1- 特定のドメインまたは命令のコーパスで事前学習されたLLMをファインチューニングする
2- 人間によって注釈付けされたデータセットを収集し、報酬モデルをトレーニングする
3- ステップ1で得られたLLMを報酬モデルとデータセットを用いてRL（例：PPO）でさらにファインチューニングする
ここで、ベースとなるLLMの選択は非常に重要です。現時点では、多くのタスクに直接使用できる「最も優れた」オープンソースのLLMは、命令にファインチューニングされたLLMです。有名なモデルとしては、BLOOMZ、Flan-T5、Flan-UL2、OPT-IMLなどがあります。これらのモデルの欠点は、そのサイズです。まともなモデルを得るには、少なくとも10B+スケールのモデルを使用する必要がありますが、モデルを単一のGPUデバイスに合わせるだけでも40GBのGPUメモリが必要です。
TRLとは何ですか？
trlライブラリは、カスタムデータセットとトレーニングセットアップを使用して、誰でも簡単に自分のLMをRLでファインチューニングできるようにすることを目指しています。他の多くのアプリケーションの中で、このアルゴリズムを使用して、ポジティブな映画のレビューを生成するモデルをファインチューニングしたり、制御された生成を行ったり、モデルをより毒性のないものにしたりすることができます。
trlを使用すると、最も人気のあるDeep RLアルゴリズムの1つであるPPOを分散的に実行するか、単一のデバイスで実行することができます！これを可能にするために、私たちはHugging Faceエコシステムのaccelerateを活用していますので、任意のユーザーは実験を興味深いスケールまでスケールアップできます。
RLで言語モデルをファインチューニングするには、以下のプロトコルに従います。これには元のモデルの2つのコピーが必要です。アクティブモデルが元の動作/分布からあまりにも大きく逸脱しないようにするために、最適化ステップごとに参照モデルのログットを計算する必要があります。これにより、最適化プロセスには最低でも2つのモデルのコピーが常に1つのGPUデバイスに必要とされるため、セットアップを単一のGPUに合わせるのはますます困難になります。
trlでは、参照モデルとアクティブモデルの間で共有レイヤーを使用することもできます。この機能の具体的な例は、解毒化の例で紹介されています。
大規模トレーニング
大規模なトレーニングは課題があります。最初の課題は、モデルとそのオプティマイザの状態を利用可能なGPUデバイスに合わせることです。1つのパラメータが使用するGPUメモリの量は、「精度」（具体的にはdtype）に依存します。最も一般的なdtypeはfloat32（32ビット）、float16、およびbfloat16（16ビット）です。最近では、トレーニングと推論についても「エキゾチックな」精度がサポートされていますが（一定の条件と制約があります）、例えばint8（8ビット）などです。要するに、モデルをGPUデバイスにロードするために、10億のパラメータごとにfloat32精度では4GB、float16では2GB、int8では1GBのメモリが必要です。このトピックについて詳しく学びたい場合は、このブログ記事を参照してください：https://huggingface.co/blog/hf-bitsandbytes-integration。
AdamWオプティマイザを使用する場合、各パラメータには8バイトが必要です（例：モデルが1Bのパラメータを持つ場合、モデルの完全なAdamWオプティマイザには8GBのGPUメモリが必要です-参照）。
これらの課題に対処するために、多くの技術が大規模なスケールで採用されています。最も一般的なパラダイムは、パイプライン並列処理、テンソル並列処理、およびデータ並列処理です。
データ並列処理では、同じモデルが複数のマシン上で並列にホストされ、各インスタンスに異なるデータバッチが供給されます。これは最も直接的な並列化戦略であり、単一のGPUの場合と同様の戦略です。また、すでにtrlによってサポートされています。パイプライン並列処理とテンソル並列処理では、モデル自体が複数のマシンに分散されます。パイプライン並列処理では、モデルはレイヤごとに分割されます。一方、テンソル並列処理では、テンソル演算（例：行列の乗算）が複数のGPUに分割されます。これらのモデル並列処理戦略では、多くのデバイスにモデルの重みをシャードする必要があり、プロセス間でのアクティベーションと勾配の通信プロトコルを定義する必要があります。これは実装が容易ではなく、Megatron-DeepSpeedやNemoなどのフレームワークの採用が必要になる場合があります。また、スケーリングLLMトレーニングに必要な他のツールとして、Adaptive activation checkpointingとfused kernelsを強調することも重要です。並列処理のパラダイムに関する詳細な情報は、こちらをご覧ください。
したがって、私たちは次の質問を自問しました：データ並列処理だけでどこまで行けるのでしょうか？既存のツールを使用して、（アクティブモデル、参照モデル、およびオプティマイザの状態を含む）超大規模なトレーニングプロセスを単一のデバイスに収めることはできるでしょうか？答えは「はい」と思われます。主な要素は、アダプタと8ビット行列の乗算です！次のセクションでこれらのトピックを説明しましょう：
8ビット行列の乗算
効率的な8ビット行列の乗算は、LLM.int8()という論文で最初に紹介された方法であり、大規模なモデルの量子化時のパフォーマンスの低下問題を解決することを目的としています。提案された方法は、リニアレイヤで裏側で適用される行列の乗算を2つの段階に分解します。アウトライアーな隠れ状態の部分はfloat16で行われ、非アウトライアーな部分はint8で行われます。
要するに、8ビット行列の乗算を使用すると、完全精度のモデルのサイズを4倍（したがって、ハーフ精度モデルの場合は2倍）縮小することができます。
低ランク適応とPEFT
2021年に発表された「LoRA: Low-Rank Adaption of Large Language Models」という論文では、大規模な言語モデルの微調整を行うために、事前学習済みの重みを凍結し、クエリと値の層のアテンション行列の低ランクバージョンを作成する方法が示されました。これらの低ランク行列は、元のモデルよりもはるかに少ないパラメータを持ち、より少ないGPUメモリでの微調整が可能になります。著者らは、低ランクアダプタの微調整によって、完全な事前学習モデルの微調整と同等の結果を得たことを示しています。
この技術により、LLMの微調整に必要なメモリ要件を少なくすることができます。ただし、アダプタレイヤでの追加の行列乗算により、前進および逆伝搬はおおよそ2倍遅くなるというデメリットもあります。
PEFTとは何ですか？
Parameter-Efficient Fine-Tuning（PEFT）は、LLM上でアダプタレイヤの作成と微調整をサポートするために作成されたHugging Faceのライブラリです。peftは、DeepSpeedとBig Model Inferenceを活用した大規模モデルのための🤗 Accelerateとシームレスに統合されています。
このライブラリは、最先端のモデルを多数サポートしており、以下のような幅広い例を提供しています：
- 因果言語モデリング
- 条件付き生成
- 画像分類
- 8ビットint8トレーニング
- Dreamboothモデルの低ランク適応
- セマンティックセグメンテーション
- シーケンス分類
- トークン分類
このライブラリは現在も活発に開発中であり、今後数ヶ月にわたって発表される予定の多くの機能があります。
20Bパラメータモデルの低ランクアダプタを使用した微調整
前提条件が整ったので、上記で言及されたツールを使用して、単一の24GB GPU上で20BパラメータLLMをRLを使用して微調整する方法をステップバイステップで説明しましょう。
ステップ1：8ビット精度でアクティブモデルを読み込む
transformersを使用してLLMの「無料のメモリ削減」を行うには、LLM.int8で説明されている方法を使用してモデルを8ビット精度で読み込むことです。これは、from_pretrainedメソッドを呼び出す際にload_in_8bit=Trueフラグを追加することで実行できます（詳細はこちらをご覧ください）。
前のセクションで述べたように、モデルを読み込むために必要なGPUメモリの量を計算するための「ハック」は、「パラメータの数十億」という考え方です。1バイトは8ビットですので、完全精度のモデル（32ビット=4バイト）では、10億パラメータごとに4GBが必要です。半精度のモデルでは10億パラメータごとに2GB、int8モデルでは10億パラメータごとに1GBが必要です。
まずはじめに、8ビットでアクティブなモデルを読み込みましょう。次のステップでは何をする必要があるかを見てみましょう！
ステップ2： peftを使用して追加のトレーニング可能なアダプターを追加する
ステップ2では、モデル内にアダプターを読み込み、これらのアダプターをトレーニング可能にします。これにより、アクティブなモデルに必要なトレーニング可能な重みの数が劇的に減少します。このステップでは、peftライブラリを活用し、わずかなコードで実行することができます。アダプターをトレーニングした後は、簡単にハブにプッシュして後で使用することができます。
ステップ3：同じモデルを使用して参照およびアクティブなロジットを取得する
アダプターを非アクティブにすることができるため、PPOのために参照とアクティブなロジットを取得するために同じモデルを使用することができます。これには、peftライブラリのdisable_adaptersコンテキストマネージャーという機能を活用しています。
トレーニングスクリプトの概要：
次に、transformers、peft、trlを使用して20Bパラメータのgpt-neoxモデルをトレーニングする方法について説明します。この例の最終目標は、メモリ制約のある設定でポジティブな映画のレビューを生成するためにLLMを微調整することでした。同様のステップは、対話モデルなどの他のタスクにも適用することができます。
全体的には、3つの主なステップとトレーニングスクリプトがあります：
- スクリプト – imdbデータセットのテキスト生成のために8ビットモデルで凍結された低ランクアダプターの微調整。
- スクリプト – アダプターレイヤーをベースモデルの重みにマージし、これらをハブに保存します。
- スクリプト – 低ランクアダプターの感情微調整によるポジティブなレビューの作成。
これらのステップを24GBのNVIDIA 4090 GPUでテストしました。24GBのGPUで全体のトレーニングを実行することは可能ですが、完全なトレーニングランは単一のA100で🤗リサーチクラスター上で行われました。
トレーニングプロセスの最初のステップは、事前学習済みモデルの微調整でした。通常、これには高性能な80GB A100 GPUが数台必要ですが、私たちは低ランクアダプターをトレーニングすることを選びました。これは因果関係言語モデリングの設定として扱い、imdbデータセットの例を1エポックトレーニングしました。imdbデータセットには、映画のレビューとそのポジティブまたはネガティブな感情を示すラベルがあります。
適応されたモデルを取り出し、RLでさらなる微調整を行うためには、まず適応された重みを結合する必要がありました。これは、事前学習済みモデルとアダプターを16ビット浮動小数点数で読み込み、適切なスケーリングが適用された重み行列で概要を作成することで実現されました。
最後に、凍結されたimdbで微調整したモデルの上に、別の低ランクアダプターを微調整することができました。RLアルゴリズムに報酬を提供するために、imdb感情分類器を使用しました。
この実験の完全なWeights and Biasesレポートは、こちらでご覧いただけます。さらにプロットやテキストの生成をチェックしたい場合は、ぜひご覧ください。
結論
peftとbitsandbytesライブラリを活用して、trlに新機能を実装し、RLHFを使用して大規模な言語モデルを合理的なコストで微調整することができるようになりました。私たちは、24GBのコンシューマーGPU上で40GBの<bfloat16でのgpt-neo-xの微調整が可能であることを示しました。この統合がコミュニティに広く使用され、RLHFを活用してより大きなモデルを微調整し、素晴らしい成果物を共有するために使用されることを期待しています。
この統合の限界を押し上げるための次のステップについて興味深い方向性を特定しました。
- この統合は、複数のGPU環境でどのようにスケーリングするのでしょうか？私たちは主に、GPUの数に対してこの統合がどのようにスケーリングするか、Data Parallelismをそのまま適用できるか、または関連するライブラリのいずれかで新しい機能の採用が必要かを探求します。
- トレーニング速度を向上させるためにどのようなツールを活用できるのでしょうか？この統合の主な欠点は、全体的なトレーニング速度です。将来的には、トレーニングをより速くするための可能な方法を探求することに興味があります。
参考文献
- 並列性のパラダイム：https://huggingface.co/docs/transformers/v4.17.0/en/parallelism
transformersにおける8ビット統合：https://huggingface.co/blog/hf-bitsandbytes-integration- LLM.int8論文：https://arxiv.org/abs/2208.07339
- 勾配チェックポイントの説明：https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-activation-checkpointing.html
We will continue to update VoAGI; if you have any questions or suggestions, please contact us!
Was this article helpful?
93 out of 132 found this helpful
Related articles
