12/02にPyTorch 2.0のアナウンスがありました。まだnightly版(α版)で正式リリースされるのは2023年3月頃のようですが、機能自体は試すことができるので早速使ってみました。
12/05現在、絶賛検証中なので結論のようなものは書けませんが、全体の傾向としては概ね公称通りに高速化の効果が認められました。
- 精度が低下することはない
- 小さなモデルに対して、学習は速くならず、コンパイルオーバヘッドのためepochsが少ない場合は全体として遅くなる、GPU使用率はAMPだと僅かに低くなる傾向
- 大きなモデルに対して、学習は速くなり(約5 ~ 30%高速化)、デフォルト設定ではVRAM使用率は少し低くなる(5 ~ 10%弱程度)
- GPUだけでなくCPUも効率良く使えるケースだと特に高い効果が期待できる
- コンパイルオプションはいくつかあるけどデフォルトで使うのが一番良さそう
あくまで後述する僕の実験環境での傾向ですので参考程度に見てください。学習に数日かかるような大きなモデルかつマルチGPU環境(A100以降)であれば、おそらく定量的にも公称通りの結果が期待できるかと思われます。誰か試してー
PyTorch 2.0の機能と使い方
PyTorch2.0では新しいコンパイラが搭載されて速くなります。ユーザー体験的には極めてシンプルです。
|
1 2 |
# model: torch.nn.Module model = torch.compile(model) |
これだけで学習が速くなるようです。公式ドキュメントはものすごく重厚長大で同じ事を何度も書いていましたが使い方はシンプル。一応理論的な部分も読みながら検証していますが、内容が難しいので自信のある方は公式ドキュメントから技術詳細を参照してください。
- PyTorch 2.0 developervendor-experience | PyTorch
- PEP 523 – Adding a frame evaluation API to CPython | peps.python.org
nightly版ではデフォルトで有効にしておらず、意図的にcompileメソッドを通す作法になるようです。
また、以下のような著名なモデルパッケージでたくさん効果検証しているらしく、当然ながら研究用途だけでなくプロダクション用途での活用も期待できます。
- 46 models from HuggingFace Transformers
- 61 models from TIMM: a collection of state-of-the-art PyTorch image models by Ross Wightman
- 56 models from TorchBench: a curated set of popular code-bases from across github
環境
環境に依って結果が大きく変わりそうなので丁寧に書いておきます。
- OS: Ubuntu 20.04 (x86_64)
- CPU: Intel(R) Xeon(R) CPU E5-2686 v4 @ 2.30GHz 8コア
- RAM: 16GB
- GPU: Tesla V100-SXM2-16GB
- CUDA: Driver Version: 510.73.08, CUDA Version: 11.6
- Python: 3.9.13
A100とCUDA 11.7の環境が一番効果があるようですが、個人でA100はさすがにコスト高いので断念。VRAM 16GBしかないのでそれほど大きなモデルは試せていません。
|
1 2 3 |
# インストール # zshの人は torch[dynamo] の部分はダブルクォートで囲っておく pip install numpy --pre torch[dynamo] torchvision torchaudio --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu116 |
Pytorch関連のモジュールは以下のようになりました。上記コマンドだとtorchvisionはcpu版がインストールされたので過去のnightly版cu116バージョンに強制的に入れ替えてみたのですが上手く動かなかったので素直に最新版を入れました。PyTorch 2.0と言っているのに1.14になっていますがnightlyではよくあることです。正式リリースではたぶん2.0になるのでしょう。
|
1 2 3 4 5 6 7 |
# Pythonモジュール pytorch-lightning 1.7.0 torch 1.14.0.dev20221204+cu116 torchaudio 0.14.0.dev20221204+cu116 torchmetrics 0.9.3 torchtriton 2.0.0+0d7e753227 torchvision 0.15.0.dev20221204+cpu |
インストール確認は公式のチェックツールが使えます。
|
1 2 3 4 5 6 7 |
$ git clone https://github.com/pytorch/pytorch $ cd tools/dynamo $ python verify_dynamo.py Python version: 3.9.13 `torch` version: 1.14.0.dev20221204+cu116 CUDA version: 11.6 All required checks passed |
CUDAは11.6未満だとダメっぽいです。ちなみにコードレベルでは _dynamo と _inductor というサブパッケージが使えるようになっていて、それぞれ新しいコンパイラを構成するコンポーネント(TorchDynamo/TorchInductor)となっています。
|
1 2 |
import torch._dynamo as dynamo import torch._inductor.utils as utils |
TorchDynamoはPythonのJITコンパイラ(FX Graphsを生成)でPyTorch 2.0の肝となる技術です。TorchInductorはコンパイラのバックエンドで、ハードウェア環境に適したオペレーションを実行します(CUDA環境ならTriton、CPU環境ならC++/OpenMP)。なので、仮にGPU環境が用意できなくてCPU環境で使ったとしてもTorchInductorのおかげで恩恵にあずかれるはずです。
検証
ここでは画像分類モデルで検証します。モデルはお手軽にtorchvisionから、データセットはKaggleから適当にお借りしました。
- モデル: Swin Transformer (swin_t/swin_b) / ConvNeXt (Large)
- データセット: Dogs vs. Cats | Kaggle (25,000 images)
- 実験環境: Native AMP(not apex) or FP32 / 10 epochs / ハードウェア環境は前述の通り
Swin Transformerのtiny(params: 28.3M)とbase(params: 87.8M)とConvNeXT large(params: 197.8M)の3種類で学習させました。その他の細かいハイパーパラメータ等は全て同じです。オプティマイザにはRAdamを使いましたがSGDとか他のものを使っても結果にはほぼ影響ありませんでした。
実装はPyTorch Lightningで書かれた既存の学習プログラムに model = torch.compile(model) の1行を加えただけです。
結果
まず前提条件として精度面ではDynamoを有効にしても低下しないこと、その上で学習速度やGPU使用率などの非機能要件面を比較します。ログ取得/可視化にはPyTorch LightningのDeviceStatsMonitorとClearMLを使いました。
Swin Transformer tiny (params: 28.3M)
まずはSwin Transformer tinyから。ClearMLで出力したグラフを貼ります。後述しますが、MLOpsツール併用によるロギング処理自体のコストを抑えるため粗めのプロットで描いています。
前提条件として精度面では低下していないことを確認できました(誤差レベルで多少変動アリ)。Precision/Recalも同様でこれ以降の検証モデルも精度面での低下は特に見られなかったので省略します。
* 学習時間
| Dynamo無効 AMP | Dynamo有効 AMP | Dynamo無効 FP32 | Dynamo有効 FP32 |
| 11:46m | 12:17m | 13:28m | 13:34m |
Dynamo無効 AMPが一番速い結果となりました。結果としては想定通りで、小さなモデルには効果が無く、むしろ最初にグラフコンパイルのコストがあるので少し遅くなるという結果でした。次に非機能要件を比較してみます。
GPU使用率はDynamo有効 AMPが一番低く、Dynamo無効と比べると平均5 ~ 10%程度節約できているようです。一方でFP32ではほぼ効果がありませんでした。
VRAM使用率ではAMP環境でのDynamoの効果が無く、逆にFP32では少し効果がありました。
念のためGPU温度を見てみるとDynamo有効 AMPが一番GPUに優しいことがわかります。ただし前述の通りVRAM使用率は変わらないのでバッチサイズを増やす余地はあまりなさそうです。
Swin Transformer base (params: 87.8M)
次にSwin Transformer baseを試します。モデルサイズによる判断基準の境界値として適切かどうか確認します。
* 学習時間
| Dynamo無効 AMP | Dynamo有効 AMP | Dynamo無効 FP32 | Dynamo有効 FP32 |
| 19.33m | 13:49m | 22:25m | 21:56m |
学習時間ではDynamo有効 AMPが一番効果があり、Dynamo無効と比べると約30%高速化しています。一方でFP32の方は効果はほぼ無いようでした。
GPU使用率ではDynamoの有無による変化はほぼ見られません、FP32でも同様です。
一方でVRAM使用率ではDynamoの有無による変化が見られました。AMPで約10%、FP32で約6%低下しています。ここで、Swin-Tでは特に違いがなかったCPU利用率も確認したところ、こちらは興味深い結果になっていたので載せます。
GPU利用率は変わらないのにCPU利用率はDynamo有効 AMPだと10%ほど増えています。仕組み的にはTorchInductorバックエンドだとCPU環境でも効果があるらしいので、それのおかげでしょうか？実際に学習処理はかなり高速化しているので、Dynamo/Inductorが効率よく働いていることは間違いなさそうです。
ConvNeXT Large (params: 197.8M)
最後に一番大きなモデルとなるConvNext Largeで試します。Dynamo効果の期待値としてはこれが一番高いです。
* 学習時間
| Dynamo無効 AMP | Dynamo有効 AMP | Dynamo無効 FP32 | Dynamo有効 FP32 |
| 16:57m | 15:56m | 33:40m | 32:17m |
結果は想定とは異なりDynamo有効 AMP環境でも6%ほどしか高速化しませんでした。。学習に数日かかるようなケースだと6%の改善でも十分な効果と言えそうですが公称値にはほど遠いですね。。コンパイラバックエンド周りのオプションがいくつかあるのでそれらを試してみると結果が変わりそうですが、それはもっと検証を進めてからまた記事に起こしたいと思います。
GPU使用率はDynamoの有無でほぼ変わりませんでした。大きいモデルなのでFP32だとがっつりGPUを使っています。
VRAM使用率はAMP環境においてはDynamoの有無でほぼ変わらなかったですが、FP32でなぜか使用率が下がっています。実行時に大量の警告がでていたので、上手くコンパイルできていなかった可能性が高いです。何度試しても同様だったのでFP32での検証は失敗として扱います。
コンパイルオプション
torch.compileの重要なオプションの一つにmodeがありますが、デフォルト値を含めて以下の3つが指定できます。
|
1 2 3 4 5 6 7 8 9 10 11 12 |
# API NOT FINAL # default: optimizes for large models, low compile-time # and no extra memory usage torch.compile(model) # reduce-overhead: optimizes to reduce the framework overhead # and uses some extra memory. Helps speed up small models torch.compile(model, mode="reduce-overhead") # max-autotune: optimizes to produce the fastest model, # but takes a very long time to compile torch.compile(model, mode="max-autotune") |
max-autotuneモードは残念ながらコンパイルが一度も通らず諦めています。。小さいモデル(Swin-T)でreduce-overheadの効果を期待して試してみたのですが、VRAM使用率が爆増しただけで学習時間はほぼ変わりませんでした。コンパイルでVRAMを逼迫させるくらいならコンパイルせずにバッチサイズを増やした方が速くなるので、小さなモデルなら無理にDynamoを使うのは避けた方が良いのかもしれません。これも引き続き調査します。
その他 いろいろ
まずはどうでもいいレベルの話ですけど、PyTorch Lightningでの学習時にDynamoを有効にしているとコンソールログに OptimizeModule と表記されるようになるので地味にありがたいです。
|
1 2 3 4 5 6 7 8 9 10 |
| Name | Type | Params ----------------------------------------------- 0 | net | OptimizedModule | 196 M <= Dynamo無効の場合はモデル名(例: ConvNext, Swin Transformer)が表記される 1 | criterion | CrossEntropyLoss | 0 2 | metrics | MetricCollection | 0 ----------------------------------------------- 196 M Trainable params 0 Non-trainable params 196 M Total params 392.467 Total estimated model params size (MB) |
コンパイルしても速くならない場合
Swin-Bはまぁまぁ効果があったのにConvNeXTだと効果が薄かったので深掘りしたいのですが、公式にはGraph breakとやらを避けろとあります。なんだそれ難しい。。
他にもプロセスがクラッシュしたときに対応方法などFAQを見ればある程度は方針がわかりますが、コンパイラの内部構造を先に勉強した方が近道かもしれません。
コンパイル時間の計測
手動でコードを差し込まないといけませんが、以下のコードでコンパイル時間を計測できるようです。学習処理が全部終わった後で実行します。
|
1 2 3 4 5 6 7 |
torch._dynamo.utils.compile_times() # 以下のようなCSV形式?のテキストが返ってくる convert_frame_assert.<locals>._convert_frame_assert, 0.6044, 0.0753, ... create_aot_dispatcher_function, 0.3882, 0.4755, 0.3467, ... compile_fx.<locals>.fw_compiler, 0.3782, 0.4505, 0.3293, ... ... 省略 |
想定外の出力でしたが、Dynamoの内部関数の実行単位でかかった時間が返ってくるようで、たぶん数字を全部足してやればコンパイルにかかった総時間になるのだと思います。Excelでそのまま読み込んでSwin-B AMPのケースで集計してみたところ129秒ほどかかっているようでした。だいたい2分ほどですね。それっぽい妥当な数字に見えますが正しい集計方法かはわかりません。
警告・エラー
nightly版なので警告とかエラーがたくさんでて滅入るのですがメモ代わりに書いておきます。W&BとかClearMLのようなMLOps系のツールを使っていろんなログを出すような環境でモデル学習しているとこういうエラーが増える気がします。
|
1 2 3 4 5 6 7 |
Exception ignored in: <function _after_at_fork_child_reinit_locks at 0x7f3d951420d0> Traceback (most recent call last): File "/home/ryo/.pyenv/versions/3.9.13/lib/python3.9/logging/__init__.py", line 255, in _after_at_fork_child_reinit_locks handler._at_fork_reinit() File "/home/ryo/.pyenv/versions/3.9.13/lib/python3.9/logging/__init__.py", line 894, in _at_fork_reinit self.lock._at_fork_reinit() AttributeError: 'NoneType' object has no attribute '_at_fork_reinit' |
Dynamoを無効にすると全く出ないエラーだったのでなんらかの影響はあるのでしょうが、ロギング処理時にロック失敗してるけど無視してもまぁ大丈夫みたいなエラーです。モデル学習処理自体には問題なさそうだったのですが、大量のエラーがでるのでコンソールログがだいぶ汚れてしまいますね。かと言って今時MLOpsツールを使わないのはツライので仕方ないです。
あとはCUDA環境だと↓のような警告も大量に出てしまいますが、Inductorのコア部分っぽいのでさらに難しいです。。
|
1 |
[2022-12-05 09:17:11,729] torch._inductor.lowering: [WARNING] using triton random, expect difference from eager |
おわりに
今回はPyTorch 2.0のTorchDynamo/TorchInductorによる高速化を試してみました。今回試したモデルだとSwin-Bで大きな効果は認めつつも、期待値が一番高かったConvNeXT Largeでの効果が薄くて謎を残す結果となりましたが、精度を落とすことなく学習が速くなることは確認できました。
コンパイルオプションがいろいろあって難しいのですが、デフォルト設定であれば model = torch.compile(model) を1行追加するだけ、ユーザー体験的には素晴らしいと思います。以前紹介したfunctorch(JAXライクなfunctorchで機械学習を速くする – part 1)を使う場合は根本から実装を書き換える必要がありますが、PyTorch 2.0を使えば実装レベルの移行コストはほぼ無いに等しいです。
今回は失敗した検証もありましたので、もう少しドキュメントやフォーラム、実装等を読み込みながら検証を続けたいと思います。予定では来年3月頃に正式リリースらしいので楽しみにしておきましょう。