DatabricksにおけるDeepSpeed Distributorのご紹介

Last updated at 2024-08-01Posted at 2024-07-28

Introducing the DeepSpeed Distributor on Databricksの翻訳です。

本書は著者が手動で翻訳したものであり内容の正確性を保証するものではありません。正確な内容に関しては原文を参照ください。

GPU不足は現実のものとなり、大規模言語モデルのトレーニングのスケールアップと最適化を可能にすることで、AIプロジェクトのデリバリーを加速する助けとなります。DeepSpeedは、GPUのメモリー要件を最大4倍削減できるフレームワークです。DeepSpeedジョブの設定と起動は複雑なものですが、Databricksが手助けをします。だからこそ、TorchDistributorとspark-tensorflow-distributorの両方と連携するSparkの完全にオープンソースな拡張である新たなDeepSpeed Distributorを発表できることを嬉しく思っています。シンプルなPySparkの関数でDeepSpeedを用いたトレーニングジョブを起動することができます - クラスターの設定、通信やモニタリングなどの面倒なことはすべてあなたのために対応します！本書では、大規模言語モデル(LLM)やコンピュータービジョン(CV)アプリケーションの構築における、あなたのAIジャーニーを進めるために必要となるDeepSpeedのキーとなる機能を確認します。

GPUボトルネックの打破

現在のディープラーニングトレーニングにおける主なボトルネックは、GPUごとのRAMです - 単純に大きなものにはなっていません。

一般的なGPU - インスタンスタイプと利用可能なRAM

この問題に対応するために主に二つのパラダイムが出現しました: データ並列化とモデル並列化のトレーニングです。

データ並列化のトレーニングには、複数GPUへのデータの分割が含まれます。このシナリオでは、それぞれのGPUには、モデルの完全なコピーとトレーニングのアーティファクト(オプティマイザの状態と勾配)を格納する必要があり、トレーニングデータで使用するRAMを大幅に減少させることになります。

データ並列化トレーニング - 我々のVRAMの大部分を消費するすべてのトレーニングアーティファクトと完全なモデルを必要とします

一方、モデル並列化トレーニングでは、モデルを複数のGPUに分割するので、それぞれのマシンのRAMを開放します。このシナリオでは新たな問題が生じます - 分散GPU間で利用できるネットワーク帯域です。

モデル並列化トレーニング - モデルを分割できますが、完全なフォワードパス、バックワードパスを得るために、3つすべてのGPUを通じてデータを引き渡す必要があることを意味します

モデル並列化のパラダイムでは、とてつもないコストを引き起こすことなしに、より大きなモデルをトレーニングできるようにする、DeepSpeedのメモリー効率性を活用することができます。

DeepSpeedとDeepSpeed Distributorの理解

DeepSpeedライブラリは、3つのキーモジュールから構成されます: ディープラーニングモデルのトレーニング、推論、圧縮です。本書では、モデルのトレーニングにフォーカスします。

DeepSpeedが何に適しており、どのように設定するのかを理解するために、トレーニングプロセスを紐解きましょう。大規模言語モデルをトレーニングやその従兄弟であるコンピュータービジョンのトレーニングにおいては、2つの制約が存在します: トレーニングデータをロードする際の入力 / 出力のスピードと、GPUカードごとに利用できるメモリーの総量です。データロードの最適化に関しては、Hugging FaceとのコラボレーションやTorchdelta拡張をご覧ください。

モデルトレーニングにおいては、主にメモリーを消費する部分が3つあります: モデルとその重み、オプティマイザの状態と勾配、トレーニングデータのサイズです。トレーニング時間を短縮するために、GPUにロードするデータの量を最大化したいと考えます。これには、モデル、オプティマイザと勾配のRAMのフットプリントを制限するためにDeepSpeedを設定する必要があります。これらの設定をウォークスルーしていきましょう。

モデル精度の設定

LLMアプリケーションのエンドユーザーに高品質のチャット体験を確実に提供するために、可能な限り大きなモデルを使いたいと考えます。デフォルトでは、モデルの重みは32ビットの浮動小数点として格納されますので、70億パラメーターモデルであっても小規模なGPUのRAMを食い潰しますし、大規模なものであっても非常に大きなチャンクを占有することになります。DeepSpeedには代わりに16ビットの浮動小数点でのトレーニングを可能とする最適化が含まれており、モデルサイズを大幅に削減します。 16ビット浮動小数点精度を用いる際の経験則として、平均的なモデルが必要とするGPU RAM(GB)は、デフォルトのfp32におけるパラメーター数(十億)の4倍ではなく、パラメーター数の約2倍を必要とするというものがあります。さらに低い精度を用いることも可能ですが、これらのテクニックは依然として実験的なものとなっています。

選択できる16ビットには2つのタイプがあります: 通常のfp16とbfloat16です。新たなNvidiaのGPU(A10/A100/H100)を扱う際には、以下のようにJSON設定ファイルで設定できるbfloat16を使うことをお勧めします。

"bf16": { 
   "enabled": true 
}

古いGPUではbfloat16をサポートしていないので、以下のように古くて効率の悪いfp16にフォールバックします。

"fp16": { 
        "enabled": "auto", 
        "loss_scale": 0, 
        "loss_scale_window": 1000, 
        "initial_scale_power": 16, 
        "hysteresis": 2, 
        "min_loss_scale": 1 
    }

可能な場合にはbfloatと互換性のあるGPUを使うことをお勧めします。 fp16のトレーニングは不安定で、学習率を変更したり、NaNのロスを回避するために、トレーニングの際に手動による介入を必要する場合があると言われ続けており、結果として、fp16には数多くの設定が必要となる一因にもなっています。詳細に関してはdocumentationやmixed precision trainingに関するNvidiaのガイドをご覧ください。

オプティマイザと勾配の設定

メモリー節約に関して次に見ていく領域は、オプティマイザと勾配の計算です。これは、Zero Redundancy Optimizer (ZeRO)の設定の活用を通じて行うことができます。ZeROはDeepSpeedのコアであり、研究論文で初めて導入されました。著者は、トレーニングの過程で大量の重みや状態、勾配を格納することは冗長であり、GPU、さらにはCPUやディスクにさえパーティショニングできると述べています。これは、データを引き回す必要性を生じさせ、ネットワークトラフィックを増加させますが、GPU RAMの節約自体に価値があります。

ZeROには3つの設定があります:

ステージ1: オプティマイザとパーティショニング
ステージ2: オプティマイザと勾配のパーティショニング
ステージ3: オプティマイザ、勾配とパラメータのパーティショニング

ステージ3の先では、CPUや接続されているストレージに計算処理をオフロードすることもできますが、これは莫大なパフォーマンスの制約を受けることになります。可能な領域節約を理解するために、オリジナルのZeROの論文からの図を見てみましょう。ここでは、ステージ1=Pos、ステージ2 = Pos+g、ステージ3 = Pos+g+pとなっています。

ステージ1からステージ3に行くに従って、より多くのVRAM削減を見て取れます。VRAM消費量の計算式は若干複雑ですが、GPUごとのVRAM要件見積もりのビジュアライゼーションを以下に示します。

4GPUを前提としてどれだけ削減できるのか計算

ZeRoの3以降では、オプティマイザの状態、勾配、パラメータをCPUのRAMやアタッチされたNVMeドライブにオフロードすることができます。オフロードによってパフォーマンスに影響を与えることになりますが、より大きなモデルを小さな環境でトレーニングすることが可能になります。

オフロードはJSONで設定できます:

# cpu offload
"offload_optimizer": {"device": "cpu"},
"offload_param": {"device": "cpu"},

# gpu offload
"offload_optimizer": {"device": "nvme", 
    "nvme_path": "/local_disk0/optimizer"},


"offload_param": {"device": "nvme",
"nvme_path": "/local_disk0/param"}

Databricksクラスターでは、一時的なキャッシュとして使用できる直接アタッチされたストレージのマウントパスは /local_disk0となります。また、オフロードにはinitスクリプト(AWS/Azure)を通じてのインストールが必要となるOSレベルのライブラリであるlibaio-devのインストールが必要となります。オフロードを用いる際、十分なCPU RAMやNVMeストレージがあることを確認することが重要です。NVMeストレージを使う際、クラスターの作成時にローカルストレージのオートスケールを設定することをお勧めします。

DeepSpeedトレーニングジョブの起動方法

キーとなるDeepSpeedの設定の動作を理解したいので、我々のパイプラインをどのようにコーディングし、起動するのかを見ていきましょう。DeepSpeed Distributorの設計の美しいところは、あなたのワークフローで活用する際に大きなコードの修正が不要であるという点です。TorchDistributorのように、DeepSpeedフレーバーでは、2つの方法で活用が可能です:

既存のノートブックのコンテキストで記述された関数を起動
外部のPythonスクリプトを起動

ノートブックから関数を起動

この例では、trainというトレーニング関数があるものとし、引数training_argumentsとdatasetを受け付けるものとします。

シングルノードでの実行:

from pyspark.ml.deepspeed.deepspeed_distributor import DeepspeedTorchDistributor

distributor = DeepspeedTorchDistributor(numGpus=1, 
nnodes=1, 
localMode=True,
useGpu=True, 
deepspeedConfig = deepspeed_dict)

completed_trainer = distributor.run(train, training_arguments, dataset)

分散処理での実行:

from pyspark.ml.deepspeed.deepspeed_distributor import DeepspeedTorchDistributor

distributor = DeepspeedTorchDistributor(numGpus=4, 
nnodes=2, 
localMode=True, 
useGpu=True, 
deepspeedConfig = deepspeed_dict)

completed_trainer = distributor.run(train, training_arguments, dataset)

DeepSpeedTorchDistributorで利用できるパラメータを説明します:

numGpus: ノードあたりのGPUの数
nnodes: トレーニングを行うノードの数
localMode: Trueの場合、(シングルノードクラスターの)ドライバーノードでのみトレーニングを実行し、Falseの場合ワーカーでトレーニングを実行
useGpu: GPUを使ってトレーニングするかどうか
deepspeedConfig: すべての設定パラメータを含むJSONファイル

useGPUは使えますが、DeepSpeedはGPUを用いて利用するのが理想的であるということを伝えさせていただきます。

Pythonスクリプトの起動

ノートブックで定義されたtrain関数の実行に加え、Pythonファイルへのパスと、引数を指定するコマンドラインを指定することで、トレーニングスクリプトをロードして実行することもできます。

シングルノードでの実行

from pyspark.ml.deepspeed.deepspeed_distributor import DeepspeedTorchDistributor

distributor = DeepspeedTorchDistributor(numGpus=4, 
nnodes=1, 
localMode=False, 
useGpu=True, 
deepspeedConfig = deepspeed_dict)

completed_trainer = distributor.run('<path_to_file>/train.py', "--var1 value1", "--var2 value2")

分散処理での実行

from pyspark.ml.deepspeed.deepspeed_distributor import DeepspeedTorchDistributor

distributor = DeepspeedTorchDistributor(numGpus=4, 
nnodes=2, 
localMode=False, 
useGpu=True, 
deepspeedConfig = deepspeed_dict)

completed_trainer = distributor.run('<path_to_file>/train.py', "--var1 value1", "--var2 value2")

生産的になるためのトリックとティップス

キーとなるDeepSpeedの設定と、トレーニングジョブの起動方法を理解したので、最高にスムーズな体験を確実にするために注意すべきいくつかの検討事項をカバーします。

メモリー要件の見積もり

特定のモデルをファイにチューニングする際の要件の管理やり回は大変です。いかなるGPUリソースも消費する前に、トレーニングジョブの重要度を最初に評価するための有用な経験則がいくつか存在します。

ここまでで議論したように、モデルのパラメーター数は、VRAMの容量や必要とするトレーニングの状態を表現します。fp16やbfloat16でトレーニングする際、シンプルな数式を使うことができます:

2 x パラメーター数(十億単位) = 必要なVRAM(GB)

つまり、llama_v2_7bのような70億パラメータモデルは、bfloat16フォーマットでロードするためには約14GBのVRAMが必要となります。

次に検討すべきことは、どのようなタイプのファインチューニングを行なっているのかということです。モデルの重み、重みの勾配、アダプターの重みやオプティマイザの状態のメモリーフットプリントは、LoRA、QLoRA、完全なファンチューニング(アクティベーションや入力勾配を除く)のいずれかを用いているのかに依存します。パラメータあたりの合計メモリー使用量は以下の表にまとめています。

*注意現時点ではQLoRAとDeepSpeedのZeRo3には互換性がありません

モデルサイズのメモリ要件に対する数式と、ファインチューニングの追加のメモリー要件を組み合わせることで、トレーニング実行に必要な合計RAMを見積もることができます:

利用可能なリソースを用いて、我々のジョブを実行できるようにサイズアップしたので、残りのトレーニング設定を見ていきましょう。

Hugging Faceキャッシュの設定

Hugging Faceライブラリを使う際は、どのようにキャッシュが行われるのかを理解することが重要です。通常、キャッシュは個々のアイテムレベル(データセットやモデルに対する固有の呼び出し)で設定され、ライブラリレベル(すべてのデータセットやトランスフォーマーオブジェクト)とグローバルレベル(すべてのHugging Faceライブラリ向け)があります。

また、アイテムのインスタンスを作成する際のキャッシュの設定がグローバルの設定を上書きする階層構造が存在します。すべてのHugging Faceライブラリのデフォルトは、DatabricksノードのルートOSフォルダである~/.cacheになっています。これはすぐにルートOSフォルダを満杯にし、クラスターをクラッシュさせます。ルートへのプレッシャーを低減するために、DBFSあるいは/local_disk0のパスにキャッシュディレクトリを設定することをお勧めします。 これは、あなたがモデルやデータセットクラスのインスタンスを作成する前に、環境変数を設定することで行うことができます。

環境変数の設定:

import os

# Setting Hugging Face cache to /local_disk0
os.environ['HF_HOME'] = '/local_disk0/hf_home'
os.environ['HF_DATASETS_CACHE'] = '/local_disk0/hf_home'
os.environ['TRANSFORMERS_CACHE'] = '/local_disk0/hf_home'

DBFSへのキャッシュは、クラスターがシャットダウンした後の永続性を保証しますが、local_disk0へのキャッシュはHugging Faceデータセットのように頻繁にアクセスされるデータに対して優れたパフォーマンスを提供します。モデルは一度しかロードされないのでDBFSにトランスフォーマーモデルをキャッシュし、データセットはlocal_disk0にキャッシュすることが経験則となります。

DeepSpeedとMLflowの設定

MLflowはtransformersを強力にサポートしており、PythonセッションはMLflowへのログイントークンを用いて初期化されます。しかし、DeepSpeedを実行する際、個々のGPUはそれらの認証情報を継承しない個別のPythonプロセスを持つことになります。進めるためには、dbutilsを用いてこれらのパラメータをPython変数に保存し、DeepspeedTorchDistributorが分散する関数内で環境変数に割り当てます。

Databricksホストとトークンの取得と設定:

# logging host and token
import os

browser_host = dbutils.notebook.entry_point.getDbutils().notebook().getContext().browserHostName().get()
db_host = f"https://{browser_host}"

db_token = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()

os.environ['DATABRICKS_HOST'] = db_host
os.environ['DATABRICKS_TOKEN'] = db_token

また、主のノートブックプロセスしか自動で新規エクスペリメントを作成できないので、MLflowエクスペリメントを作成、設定する必要があります。最後に、MLflowのロギングを行う際、トレーニングプロセスのヘッドノードのみから実行するようにする必要があります。ヘッドノードは、Nvidia用語でいうところのglobal_rankが0のものとなります。

PyTorchによるGPUランクの取得:

global_rank = torch.distributed.get_rank()

Tensorboardインテグレーション

MLflowはトレーニング実行の可視化のためのロギングとモニタリングの機能を提供しますが、Tensorboardを実行のモニタリングに使うことは依然として有用であり、以下のコマンドを実行することでDatabricks内で起動することができます。

DatabricksノートブックでのTensorboardの起動

%load_ext tensorboard
# This sets up our tensorboard settings
experiment_log_dir = <log-directory>

%tensorboard --logdir $experiment_log_dir
# This starts Tensorboard

トレーニング実行中にTensorboardを使うには、自身のノートブックでこのコードを実行します。DBFSにログを書き込む際には、実行が完了するまでTensorboardで結果を参照できないことがあることに注意してください。

勾配の集約

GPUのVRAMの制約がない場合、トレーニングデータをより高速に処理するためにより大きなバッチを設定することができます。VRAMの制約のため、我々は通常バッチサイズを1桁にすることになります。勾配の集約は、十分なデータレコードを処理するまで重みの更新を避けるようにアルゴリズムに指示します。これによって、パフォーマンスに対して少ないコストでVRAMの要件を削減することができます。

DeepSpeedのJSON設定における勾配集約の設定:

"gradient_accumulation_steps": 4

勾配のチェックポイント

上で議論したように、トレーニング実行における勾配の計算処理は、VRAM消費の主な原因となっています。デフォルトでは、フォワードパスでのネットワークのアクティベーションの全ての値はバックワードパスで使用するので格納されています。勾配のチェックポイント戦略では、アクティベーションのあるポーションのみを格納します。これは、いくつかの部分では再計算が必要となり、トレーニングループを遅くさせるということを意味します(詳細はこちらをご覧ください)。

勾配チェックポイントの有効化:

model.gradient_checkpointing_enable()

我々の実験では、4つのA10GでZeRo 3の勾配チェックポイントのオフロードを用いた際、GPUあたり約6GBのVRAMを削減しました。

最適な設定の特定

DeepSpeedのトレーニンググループを設定するゴールは、可能な限り効率的に実行させるということです。最もフラストレーションを感じることの一つは、あなたのトレーニング実行の奥深くでアウトオブメモリーが発生するということです。

GPUにおける典型的なアウトオブメモリー(OOM)例外

このようなことが発生することを制限するために、最適な設定を確立するために以下のステップに従うことをお勧めします:

上述の表を用いてメモリー要件を見積もる
モデルの精度にはbfloat16を使う
オプティマイザと勾配の設定ではZeRo 3を使う

トレーニングの実行を試し、どうなるのかを確認します。アウトオブメモリー問題に遭遇するのであれば、徐々にオフロードや勾配集約ステップ、勾配チェックポイントを行います。トレーニングプロセスが遅くなる場合がありますが、完了の可能性は高まることでしょう。

モニタリングとデバッグ

DeepSpeedトレーニングの実行の監視、デバッグするベストな方法は、DatabricksのWeb UIのMetricsとDriver logsのタブを参照することです:

クラスターUIにおけるハードウェア使用率のメトリクスとシステムログ

メトリクスタブには、GPUやCPU利用率を含む個別のハードウェア設定のドロップダウンがあります。

DatabricksにおけるGPU利用率メトリクス

DatabricksにおけるロードとCPU利用率のメトリクス

ドライバーログには、ノートブックでは十分に表示されないような情報を含むターミナルセッションに表示されるのと同じデバッグの詳細が表示されます。

ドライバーログの追加詳細情報

まとめ

DatabricksにおけるDeepSpeed Distributorは、GPUのメモリー要件を大幅に削減することで、大規模言語モデルの効率的なトレーニングにおける重要な前進を示しています。リソース設定や管理をシンプルにし、スケーラブルで最適化されたAIプロジェクトの開発を実現します。こちらのコードを用いてすぐにDatabricksで活用してみましょう。

はじめてのDatabricks

Databricks無料トライアル

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up