Published on

transformersでdeepspeedを使ってみる(環境構築編)

Author
株式会社イエローバックの機械学習エンジニアです

DeepSpeed

マイクロソフトが公開しているディープラーニング向けの最適化ライブラリ「DeepSpeed」を HuggingFace transformers で使ってみます。 DeepSpeed を使うと巨大なモデルの学習を複数ノード、複数 GPU を使って効率的に行うことができますが、GPU1 個だけの場合でも少し大きめのモデルの学習ができるという恩恵を受けることができます。

本記事ではメモリが 11GB と少なめの RTX2080Ti GPU1 個を使って、どの程度のモデルで学習できるかを確認します。

環境構築

まずは、必要な環境の構築をおこないます。

HF transformers で DeepSpeed を使う場合、学習実行時に CuDA ライブラリを用いてビルドが行われます。 そのため CuDA 対応の pytorch に加えて、同一バージョンの CuDA ライブラリをあらかじめインストールしておく必要があります。

環境

以下の環境を利用します。

  • GPU: RTX2080Ti x 1
  • CPU: Intel Core i9-9900K
  • MEM: 64GB
  • OS: Ubuntu 20.04

Ubuntu パッケージインストール

まずは必要な Ubuntu のパッケージをインストールします。

$ sudo apt install -y sudo apt-get update\
  build-essential language-pack-ja wget \
  mecab mecab-ipadic-utf8 mecab-jumandic-utf8 libmecab-dev \
  python3.8-dev python3-venv software-properties-common jq \
  llvm-9 libaio1 libaio-dev cmake

CuDA のインストール

参照 https://developer.nvidia.com/cuda-downloads?target_os=Linux&target_arch=x86_64&target_distro=Ubuntu&target_version=2004&target_type=debnetwork

レポジトリを設定します。

$ wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
$ sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
$ sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/7fa2af80.pub
$ sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /"
$ sudo apt-get update

nVidia ドライバをインストールします。kernel モジュールもインストールされるためインストール後には一度リブートします。

$ sudo apt-get install -y --no-install-recommends nvidia-driver-465
$ sudo reboot

CuDA 11.1 をインストールします。

$ sudo apt-get install -y --no-install-recommends cuda-11-1
$ sudo reboot

CuDA のバイナリへのパスを設定します。CuDA は複数バージョンがインストール可能なので、ここでは明示的に CuDA 11.1 を指定しています。

$ PATH=/usr/local/cuda-11.1/bin:$PATH

必要に応じて~/.bashrc などへの設定も必要です。

python packages

python の仮想環境を使用します。

$ python3 -m venv ~/py38-ds-cuda11
$ source ~/py38-ds-cuda11/bin/activate

pytorch は cuda 11.1 対応のものをインストールします。

$ pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html
$ pip install transformers fugashi unidic-lite datasets wheel
$ pip install --no-cache-dir deepspeed

以前に deepspeed をインストールしていると、キャッシュに残っていてそのままインストールされてしまいます。 以前のものが cuda10.2 など別の cuda 対応の deepspeed だったりするとバージョンミスマッチとなってしまうので、 明示的に"--no-cache-dir"を指定してキャッシュを使用しないようにしています。

インストール状況確認

deepspeed をインストールすると ds_report コマンドがインストールされます。 このコマンドで deepspeed 関連のインストールされているバージョン情報等が確認できます。

$ ds_report
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
cpu_adam ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires CUDA version 10.1+, does not currently support >=11 or <10.1
sparse_attn ............ [NO] ....... [NO]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
async_io ............... [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/ono/py38-ds-cuda11/lib/python3.8/site-packages/torch']
torch version .................... 1.8.1+cu111
torch cuda version ............... 11.1
nvcc version ..................... 11.1
deepspeed install path ........... ['/home/ono/py38-ds-cuda11/lib/python3.8/site-packages/deepspeed']
deepspeed info ................... 0.3.16, unknown, unknown
deepspeed wheel compiled w. ...... torch 1.8, cuda 11.1

deepspeed で sparse_attn を有効にするには、現状では CUDA 10.1 か 10.2 が必要なようです。 ただ、HF transformers は現時点では DeepSpeed Sparse Attention に対応していないため、この WARNING は無視してしまいます。

参考 https://github.com/huggingface/transformers/issues/9606

おわりに

次回は、実際に deepspeed を使って学習をおこないます。

Discuss on Twitter