PyTorch 2.0 の動作: HuggingFace および TIMM モデルの高速化!

PyTorch 2.0 では、torch.compile() の簡単な行でモデルのトレーニング速度を 30% ～ 200% 向上させることができます。このチュートリアルでは、この高速化を実際に再現する方法を示します。

torch.compile()  できるさまざまなコンパイラ バックエンドを簡単に試すことができます。これにより、PyTorch コードの実行が高速化されます。それは torch.jit.script()  ソース コードを変更せずに nn.Module 上で直接実行できるドロップイン置換です。

前回の記事では、torch.compile が任意の PyTorch コード、制御フロー、ミューテーションをサポートし、動的シェイプをある程度サポートしていることを紹介しました。

163 のオープンソース モデルをテストした結果、torch.compile() が 30% ～ 200% の高速化を実現できることがわかりました。

opt_module = torch.compile(module)

詳細なテスト結果については、以下を参照してください。

https://github.com/pytorch/torchdynamo/issues/681

このチュートリアルでは、その利用方法を説明します。 torch.compile()  モデルのトレーニングを高速化します。

要件と設定

GPU の場合 (GPU が新しいほど、パフォーマンスの向上がより顕著になります):

pip3 install numpy --pre torch[dynamo] --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117

CPUの場合:

pip3 install --pre torch --extra-index-url https://download.pytorch.org/whl/nightly/cpu

オプション: インストールの確認

git clone https://github.com/pytorch/pytorch
cd tools/dynamo
python verify_dynamo.py

オプション: Docker のインストール

必要な依存関係はすべて PyTorch の Nightly Binaries ファイルで提供されており、次の方法でダウンロードできます。

docker pull ghcr.io/pytorch/pytorch-nightly

アドホックな実験の場合、コンテナがすべての GPU にアクセスできることを確認してください。

docker run --gpus all -it ghcr.io/pytorch/pytorch-nightly:latest /bin/bash

始める

  簡単な例

まず簡単な例を見てみましょう。GPU が新しいほど、速度の向上がより明らかになることに注意してください。

import torch
   def fn(x, y):
       a = torch.sin(x).cuda()
       b = torch.sin(y).cuda()
       return a + b
   new_fn = torch.compile(fn, backend="inductor")
   input_tensor = torch.randn(10000).to(device="cuda:0")
   a = new_fn()

この例では実際に速度が向上するわけではありませんが、出発点としては役立ちます。

この例では、トーチ.cos()  そして トーチ.sin()  は点単位の演算の例です。より有名な点単位の演算は要素ごとに実行できます。 トーチ.relu()。

各演算子がメモリからテンソルを読み取り、いくつかの変更を加えてから、それらの変更を書き戻す必要があるため、eager モードでのポイント単位の操作は最適ではありません。

PyTorch 2.0 の最も重要な最適化の 1 つは融合です。

したがって、この例では、2 つの読み取りと 2 つの書き込みを 1 つの読み取りと 1 つの書き込みに変えることができます。これは、これらの GPU のボトルネックはメモリ帯域幅 (どの程度のメモリ帯域幅を使用できるか) であるため、これは重要です。計算 (GPU が浮動小数点演算をどれだけ速く実行できるか) ではなく、GPU が GPU にデータを送信します。

PyTorch 2.0 の 2 番目に重要な最適化は CUDA グラフです。

CUDA グラフは、Python プログラムから個々のコアを起動するオーバーヘッドを排除するのに役立ちます。

torch.compile() はさまざまなバックエンドをサポートしていますが、その中で最も注目に値するのは、Triton カーネルを生成できる Inductor です。

https://github.com/openai/triton

これらのカーネルは Python で書かれており、ただし、ほとんどの手書きの CUDA カーネルよりも優れています。上記の例が trig.py という名前であるとすると、triton カーネルを生成するコードを実行して実際に検査できます。

TORCHINDUCTOR_TRACE=1 python trig.py

@pointwise(size_hints=[16384], filename=__file__, meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': 0, 'constants': {}, 'configs': [instance_descriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]})
   @triton.jit
   def kernel(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
       xnumel = 10000
       xoffset = tl.program_id(0) * XBLOCK
       xindex = xoffset + tl.reshape(tl.arange(0, XBLOCK), [XBLOCK])
       xmask = xindex < xnumel
       x0 = xindex
       tmp0 = tl.load(in_ptr0 + (x0), xmask)
       tmp1 = tl.sin(tmp0)
       tmp2 = tl.sin(tmp1)
       tl.store(out_ptr0 + (x0 + tl.zeros([XBLOCK], tl.int32)), tmp2, xmask)

上記のコードからわかります: 罪  両方が存在するため、融合が発生しました 罪  オペレーターは Triton カーネルで実行され、一時変数はレジスターに保存され、非常に迅速にアクセスされます。

実際のモデルの例

例として、PyTorch Hub の resnet50 を取り上げます。

import torch
   model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
   opt_model = torch.compile(model, backend="inductor")
   model(torch.randn(1,3,64,64))

実際の操作では、モデルがコンパイルされているため、最初の実行が非常に遅いことがわかります。以降の実行は高速化されます。したがって、ベンチマーク テストを開始する前に、モデルをウォームアップするのが一般的です。

ご覧のとおり、ここではコンパイラ名を表すために「inductor」を使用していますが、利用可能なバックエンドはこれだけではなく、REPL で実行できます。 torch._dynamo.list_backends()  利用可能なバックエンドの完全なリストを表示します。

試してみることもできます aot_cudagraphs  または nvfuser  。

ハグ顔モデル例

PyTorch コミュニティでは、トランスフォーマーまたは TIMM 事前トレーニング済みモデルがよく使用されます。

https://github.com/huggingface/transformers

https://github.com/rwightman/pytorch-image-models

PyTorch 2.0 の設計目標の 1 つは、コンパイル スタックが実際に実行されるほとんどのモデルでそのまま動作する必要があることです。

ここでは、事前トレーニング済みモデルを HuggingFace ハブから直接ダウンロードして最適化します。

import torch
   from transformers import BertTokenizer, BertModel
   # Copy pasted from here https://huggingface.co/bert-base-uncased
   tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
   model = BertModel.from_pretrained("bert-base-uncased").to(device="cuda:0")
   model = torch.compile(model) # This is the only line of code that we changed
   text = "Replace me by any text you'd like."
   encoded_input = tokenizer(text, return_tensors='pt').to(device="cuda:0")
   output = model(**encoded_input)

モデルから削除された場合 to(device=”cuda:0″)  そして エンコードされた入力 , PyTorch 2.0 は、CPU 上で実行するように最適化された C++ カーネルとして生成されます。

BERT の Triton または C++ カーネルを調べることができます。これらは、上記の三角関数の例よりも明らかに複雑です。ただし、PyTorch を知っている場合はスキップできます。

同じコードを次のように使用すると、さらに良い結果が得られます。

* https://github.com/huggingface/accelerate

*DDP

同様に、TIMM の例を試してください。

import timm
   import torch
   model = timm.create_model('resnext101_32x8d', pretrained=True, num_classes=2)
   opt_model = torch.compile(model, backend="inductor")
   opt_model(torch.randn(64,3,7,7))

PyTorch の目標は、より多くのモデルに適応し、ほとんどのオープン ソース モデルの動作を高速化できるコンパイラーを構築することです。今すぐHuggingFace Hubにアクセスしてください、PyTorch 2.0 で TIMM モデルを高速化しましょう!

https://huggingface.co/timm