こんにちは。Chief Research Oficerの西鳥羽です。今回はDeepSpeedを用いてHuggingface Transformersの複数ノードでの学習をする方法を紹介します。
Huggingface Transformersは事前学習済みモデルを簡単に扱うことができるフレームワークです。BERTなどの言語モデルをはじめとして最近はWhisperなどの音声モデル、DETRなどの画像モデルも扱えるようになっています。Huggingface Transformersでは数多くの事前学習済みモデルを用意しているため事前学習を行わなくても用いることは可能ですが、多くのモデルで事前学習にも対応しています。
Huggingface Transformerでは複数GPUが搭載されている単一のサーバーでの学習に対応していて、そちらは特に設定の変更などは無く学習の実行ができます。複数のGPUが搭載された複数のサーバーを用いた学習は、このブログでも紹介したことがある(紹介1 紹介2)DeepSpeedというライブラリを用いると行うことができます。
DeepSpeedでの並列化のしくみ
DeepSpeedでの複数サーバーでの学習を行う際、それぞれのサーバーでのプロセスの起動、終了を含めた管理はDeepSpeedが行ってくれます。従ってユーザーから行う操作は transformers/examples/pytorch/language-modeling/run_mlm.py などの学習プログラムを実行するノードからのみ行います。 学習プログラムを実行すると、DeepSpeedはまず各ノード各GPUを管理するプロセスを立ち上げるために全てのノードに搭載されているGPUの枚数分のSSH接続を行います。
それぞれのプロセスが生成された後、すべてのプロセスで学習データの読み込み、モデルの読み込みが行われます。そのあとにモデルの学習が行われます。モデルの学習はパラメータや勾配の共有などを行いながら全プロセス協調して行われます。始めのプロセスの立ち上げの通信がSSHで行われていたため、パラメータの共有などもSSHが使われて遅いのでは? と思われるかもしれませんが、学習時のモデルのパラメータや勾配の共有の通信はNCCLを通じた高速な通信を用いて行われます。そのため、Infinibandを用いるなどの高速な通信の性能を活かすことができます。
また、注意点としては学習データの読み込み時にディスクアクセスがボトルネックになって時間がかかってしまうことがあります。これは学習データの読み込み・前処理・キャッシュのためにHuggingface Datasetsを用いることが多いかと思われますが、Huggingface Datasetsの仕様上、学習を始める前にデータの大量の読み込み処理が数回発生します。この読み込み処理が全プロセスでほぼ同時に行われてしまうため、特に共有ディスクにそれらのファイルを配置してアクセスできるようにしている場合にディスクアクセスがボトルネックになってしまいます。このような事態が発生してしまう場合には学習データおよびキャッシュファイルをそれぞれのサーバーのローカルのSSDにコピーしておいたりBeeONDなどのSSDの性能を活かすことができる分散ファイルシステム上に移動しておくなどして同時アクセスへの対応を行っておくことによって学習実行までの時間を短縮することができます。
事前準備
Huggingface TransformersおよびPytorchのインストール手順を説明します。単一のサーバーで学習する時のインストール方法と変わらないので既にHuggingface Transformersを用いて学習を行う環境を構築している方はスキップしてもらって構いません。インストールの詳しい方法、或いは最新バージョンに対応した方法は Installation を参照してください。また、CUDAおよびNCCLをインストールしておく必要がありますが、ここでは割愛します。1/30現在ではPytorchおよびDeepSpeedがCUDA11系しかサポートしていないのでCUDA11系を事前にインストールしてください。
Pytorchのインストール
DeepSpeedはPytorchのみに対応しています。以下のコマンドを実行してPytorchをインストールしてください*1。
pip install torch
Huggingface Transformersのインストール
Pytorchのインストールが完了した後はHuggingface Transformersのインストールを行います。複数サーバーでの学習はリリースバージョンでも行うことは可能ですが、exampleの学習スクリプトを用いるためEditable Installを行います。
git clone https://github.com/huggingface/transformers.git cd transformers pip install -e .
複数サーバーでの学習のためのセットアップ
DeepSpeedを用いて複数サーバーで学習を行う際、単一のサーバーで学習するときに加えていくつかの設定が必要になります。
DeepSpeedのインストール
Huggingface Transformersをインストールした後にDeepSpeedのインストールを行います。 pip install deepspeed などのコマンドでインストールを行ってください*2。DeepSpeedは機能が豊富で全てをインストールしようとすると追加で依存ライブラリのインストールが必要になったりもするので、細かくインストールする機能を制御したい場合はDeepSpeedのInstallation Detailsを参照してください。
SSHの設定
DeepSpeedを用いて複数ノードでの学習を行う際、パスワードやパスフレーズの入力無しでSSHログインができる必要があります。そのため、ユーザー名やSSH接続を受け付けるポート番号、認証に用いる鍵の設定などがデフォルトではない場合 .ssh/config を修正して学習で用いる予定の全てのサーバーにパスワード、パスフレーズ無しのログインができるように設定します。
環境変数
複数サーバーで学習する際、学習で用いる各サーバーにSSHログインをするため、学習プログラム実行時に設定している環境変数は引き継がれません*3。
そのため、HF_HOME で学習データやモデルのキャッシュが置かれる場所を変更する場合など、SSHログインした後に環境変数が設定されるように Bourne ShellやBashで読み込まれる .profile に記述しておく必要があります*4。以下はモデルや学習データのキャッシュの保存場所を /tmp/cache に変更する設定の例です。.profile に記述します。
export HF_HOME=/tmp/cache
この他、PYTHONPATH や NCCL_IB_DISABLE などTransformers以外にも依存ライブラリが用いる設定のために指定している環境変数があれば同様に .profile に記述します。
学習に用いるサーバーおよびGPUの設定
学習プログラムを実行する時には他のサーバーへのログインやプロセスの生成などの処理はDeepSpeedが行ってくれるため、直接コマンドを実行するサーバー以外での操作必要は無いのですが、どのサーバーを使うかの情報を記述する必要があります。サーバーアドレス名、GPU数の順番に記述します。例えばworker-1, worker-2 というアドレスの2台のサーバーでそれぞれ8枚のGPUを用いて学習する時、以下のように記述します。
worker-1 slots=8 worker-2 slots=8
学習の実行
学習の実行時には deepspeed コマンドを経由して実行します。以下の例はrun_mlm.pyを用いて事前学習を行う時の例です。事前準備の学習に用いるサーバーおよびGPUの設定を記述したファイルを hostfile.txt という名前で保存しているとします。
また、DeepSpeedには学習の省メモリ化、高速化を行う機能もあり*5、その設定を ds.conf に記述し、--deepspeed ds.conf というオプションで指定しています。こちらに関しては単一ノードでの設定と同様なので割愛します。この設定に関しては過去のブログ記事や DeepSpeed Integration を参照してください。それ以外の学習のパラメータやハイパーパラメータは単一でのサーバーでの学習の時と同様にコマンドライン引数で与えることができます。
deepspeed --hostfile=hostfile.txt \
run_mlm.py \
--deepspeed ds.conf \
--model_name_or_path /path_to_model_dir \
--dataset_name /path/to/datasets \
--max_seq_length 512 \
--preprocessing_num_workers 128 \
--do_train \
--per_device_train_batch_size 16 \
--per_device_eval_batch_size 128 \
--learning_rate 2e-4 \
--weight_decay 0.01 \
--adam_beta1 0.9 \
--adam_beta2 0.999 \
--adam_epsilon 1e-6 \
--warmup_steps 10000 \
--save_steps 100 \
--logging_steps 100 \
--save_total_limit 2 \
--dataloader_num_workers 2 \
--output_dir /path/to/output_dir \
--max_steps 50000
まとめ
今回はDeepSpeedを用いてHuggingface Transformersの複数ノードでの学習を行う方法を紹介しました。以前の記事で紹介したZeROなどの省メモリ化や16bit浮動小数点を用いた高速化などの機能と組み合わせて用いることも可能なので、Huggingface Transformersのみの時よりも大きなパラメータサイズのモデルの学習を行うことが可能になります。
Transformerをベースとした言語モデルはモデルサイズが大きい方が性能が良くなる傾向にあるので、大きなモデルを学習する方法が提供されているのはありがたいものです。
*1:また、Pytorchが標準でサポートしているバージョンではないCUDAをインストールしている場合にはこちらの Install Pytorchを参照してください。
*2:環境変数無しでインストールした場合はJITコンパイルされます。この方法だと実行開始時にコンパイルを行うため試行錯誤しにくく、筆者は DS_BUILD_OPS=1 pip install deepspeed と指定してインストール時にコンパイルするようにしています。こうすることで実行を開始するまでの時間が短縮されます。
*3:筆者は確認してませんが、 ~/.deepspeed_env というファイルに記述することで環境変数が設定されるかもしれません。https://www.deepspeed.ai/getting-started/#multi-node-environment-variables
*4:ここでは .profile としましたが、必ずしも .profileである必要はありません。シェルでログインしたときに読み込まれるファイルは幾つかあり、そのいずれかのファイルでも構いません。
*5:むしろ学習の高速化、省メモリ化の方がメインの機能です。Training Overview and Features - DeepSpeed
