コードを1行変更するだけで、PyTorchのトレーニングを3倍高速化できます。これらの「高度なテクニック」が鍵となります。

最近、ディープラーニング分野の著名な研究者であり、Lightning AI のチーフ人工知能教育者である Sebastian Raschka 氏が、CVPR 2023 で「最小限のコード変更による PyTorch モデルトレーニングのスケーリング」と題した基調講演を行いました。

セバスチャン・ラシュカ氏は、自身の研究成果をより多くの人々と共有するために、このスピーチを記事にまとめました。この記事では、最小限のコード変更で PyTorch モデルのトレーニングを拡張する方法を説明し、低レベルのマシン最適化ではなく、混合精度法とマルチ GPU トレーニング モードを活用することに重点を置いていることを示します。

この記事では、Visual Transformer (ViT) をベースモデルとして使用しています。ViT モデルは、基本データセットからゼロから開始し、約 60 分のトレーニング後にテスト セットで 62% の精度を達成します。

GitHub アドレス: https://github.com/rasbt/cvpr2023

以下は元の記事です。

ベンチマークの構築

次のセクションでは、Sebastian が、大規模なコード リファクタリングを行わずにトレーニング時間と精度を向上させる方法について説明します。

ここでの焦点はモデルとデータセットの詳細ではないことに注意することが重要です (読者が多くの依存関係をダウンロードしてインストールすることなく、自分のマシンで再現できるように、できるだけシンプルにすることを意図しています)。ここで共有されているすべての例は GitHub で公開されており、読者は完全なコードを調べて再利用できます。

スクリプト 00_pytorch-vit-random-init.py の出力。

トレーニングをゼロから始めないでください

現在、テキストや画像に基づいてディープラーニング モデルをゼロからトレーニングすることは、多くの場合非効率的です。通常、事前トレーニング済みのモデルを使用し、モデルを微調整して、時間とコンピューティング リソースを節約しながら、より優れたモデリング結果を実現します。

上記で使用したのと同じ ViT アーキテクチャを別のデータセット (ImageNet) で事前トレーニングし、微調整すると、はるかに短い時間でより優れた予測パフォーマンスを実現できます。20 分 (3 トレーニング エポック) で 95% のテスト精度が得られます。

00_pytorch-vit-random-init.py と 01_pytorch-vit.py の比較。

コンピューティングパフォーマンスの向上

最初からトレーニングする場合と比較して、微調整によってモデルのパフォーマンスが大幅に向上することがわかります。下の棒グラフはこれを要約したものです。

00_pytorch-vit-random-init.py と 01_pytorch-vit.py の比較棒グラフ。

もちろん、モデルのパフォーマンスはデータセットやタスクによって異なる場合があります。しかし、多くのテキストおよび画像タスクでは、一般的な公開データセットで事前トレーニングされたモデルから始める価値があります。

次のセクションでは、予測精度を犠牲にすることなくトレーニング時間を短縮するためのさまざまな手法について説明します。

オープンソースライブラリ Fabric

最小限のコード変更で PyTorch のトレーニングを効率的に拡張する 1 つの方法は、オープンソースの Fabric ライブラリを使用することです。これは、PyTorch の軽量ラッパー ライブラリ/インターフェースと考えることができます。 pip 経由でインストールします。

 pip install lightning

以下で説明するすべてのテクニックは、純粋な PyTorch でも実装できます。 Fabric は、このプロセスをより便利にすることを目的としています。

コードを高速化する高度なテクニックを検討する前に、まず Fabric を PyTorch コードに統合するために必要な小さな変更について説明しましょう。これらの変更を行った後は、コードを 1 行変更するだけで、高度な PyTorch 機能を簡単に使用できるようになります。

PyTorch コードと Fabric を使用するように変更されたコードの違いは最小限であり、次のコードに示すように、いくつかの小さな変更のみが含まれます。

通常の PyTorch コード (左) と Fabric を使用した PyTorch コード

上の図をまとめると、通常の PyTorch コードを PyTorch + Fabric に変換するには次の 3 つの手順が必要です。

• Fabric をインポートし、Fabric オブジェクトをインスタンス化します。
• Fabric を使用して、モデル、オプティマイザー、データ ローダーを設定します。
• 損失関数は loss.backward() の代わりに fabric.backward() を使用します。

これらの小さな変更により、既存のコードをさらにリファクタリングすることなく、PyTorch の高度な機能を活用できるようになります。

以下の「高度な機能」に進む前に、モデルのトレーニング実行時間と予測パフォーマンスが以前と同じであることを確認してください。

01_pytorch-vit.py と 03_fabric-vit.py の比較結果。

前の棒グラフからわかるように、トレーニングの実行時間と精度は予想どおり以前とまったく同じです。あらゆる変動はランダム性に起因する可能性があります。

前のセクションでは、Fabric を使用して PyTorch コードを変更しました。なぜそんなに面倒なことをするのでしょうか?次に、混合精度や分散トレーニングなどの高度な手法を試します。コードを1行変更して、次のコードを置き換えます。

 fabric = Fabric(accelerator="cuda")

変更する

fabric = Fabric(accelerator="cuda", precisinotallow="bf16-mixed")

04_fabric-vit-mixed-precision.py スクリプトの比較結果。スクリプトアドレス: https://github.com/rasbt/cvpr2023/blob/main/04_fabric-vit-mixed-precision.py

混合精度トレーニングを使用することで、同じ予測パフォーマンスを維持しながら、トレーニング時間を約 18 分から 6 分に短縮しました。このトレーニング時間の短縮は、Fabric オブジェクトをインスタンス化するときにパラメータ「precisinotallow="bf16-mixed"」を追加するだけで実現できます。

混合精度の理解

混合精度トレーニングでは、基本的に 16 ビットと 32 ビットの精度を使用して、精度が失われないようにします。 16 ビット表現で勾配を計算すると、32 ビット形式よりもはるかに高速になり、メモリも大幅に節約されます。この戦略は、メモリや計算に制限がある状況で非常に有益です。

すべてのパラメータと演算が 16 ビット浮動小数点数に変換されるわけではないため、「低」精度トレーニングではなく「ハイブリッド」精度トレーニングと呼ばれます。代わりに、トレーニング中に 32 ビットと 16 ビットの操作を切り替えるため、「混合」精度と呼ばれます。

下の図に示すように、混合精度トレーニングには次の手順が含まれます。

• 計算を高速化するために重みを低い精度 (FP16) に変換します。
• 勾配を計算します。
• 数値の安定性を維持するために、勾配をより高い精度 (FP32) に戻します。
• スケールされたグラデーションを使用して元の重みを更新します。

このアプローチにより、ニューラル ネットワークの精度と安定性を維持しながら効率的なトレーニングが実現します。

より詳細な手順は次のとおりです。

• 重みを FP16 に変換: このステップでは、最初は FP32 形式で表されるニューラル ネットワークの重み (またはパラメーター) が、精度の低い FP16 形式に変換されます。これによりメモリ使用量が削減され、FP16 操作に必要なメモリが少なくなるため、ハードウェアによる処理が高速化されます。
• 勾配を計算する: ニューラル ネットワークの順方向および逆方向の伝播に、精度の低い FP16 重みを使用します。このステップでは、ネットワークの重みに対する損失関数の勾配 (偏導関数) を計算します。これらの勾配は、最適化プロセス中に重みを更新するために使用されます。
• 勾配を FP32 に戻す: FP16 形式で勾配を計算した後、より高精度の FP32 形式に戻します。この変換は数値の安定性を維持し、精度の低い演算を使用するときに発生する可能性のある勾配の消失や爆発などの問題を回避します。
• 学習率を乗算して重みを更新します。FP32 形式で表現された勾配に学習率を乗算して、重み (最適化中のステップ サイズを決定するスカラー値) を更新します。

ステップ 4 の積は、元の FP32 ニューラル ネットワークの重みを更新するために使用されます。学習率は最適化プロセスの収束を制御するのに役立ち、優れたパフォーマンスを実現するために非常に重要です。

ブレインフロート16

先ほど、「float 16 ビット」精度のトレーニングについて説明しました。前のコードでは、precisinotallow="16-mixed" ではなく precisinotallow="bf16-mixed" が指定されていることに注意してください。これらは両方とも有効なオプションです。

ここで、「bf16-mixed」の「bf16」は、Brain Floating Point (bfloat16) を意味します。 Google は、特に Tensor Processing Unit (TPU) における機械学習およびディープラーニング アプリケーションで使用するためにこの形式を開発しました。 Bfloat16 は従来の float16 形式に比べてダイナミック レンジを拡張しますが、ある程度の精度が犠牲になります。

拡張されたダイナミック レンジにより、bfloat16 は非常に大きな数値と非常に小さな数値を表現できるようになり、ディープラーニング アプリケーションで発生する可能性のある値の範囲に適したものになります。ただし、精度が低いと、一部の計算の精度に影響したり、場合によっては丸め誤差が発生したりする可能性があります。しかし、ほとんどのディープラーニング アプリケーションでは、この精度の低下はモデリングのパフォーマンスにほとんど影響を与えません。

bfloat16 はもともと TPU 用に開発されましたが、NVIDIA Ampere アーキテクチャに基づく A100 Tensor Core GPU をはじめ、いくつかの NVIDIA GPU が bfloat16 のサポートを開始しています。

次のコードを使用して、GPU が bfloat16 をサポートしているかどうかを確認できます。

 >>> torch.cuda.is_bf16_supported() True

GPU が bfloat16 をサポートしていない場合は、precisinotallow="bf16-mixed" を precisinotallow="16-mixed" に変更できます。

マルチGPUトレーニングと完全にシャード化されたデータ並列処理

次に、マルチ GPU トレーニングを変更してみます。これは、複数の GPU を利用できる場合にモデルのトレーニングを大幅に高速化できるため便利です。

ここでは、データ並列処理とテンソル並列処理の両方を活用する、より高度な手法である Fully Sharded Data Parallelism (FSDP) を紹介します。

Fabric では、FSDP を使用して、次の方法でデバイスの数とマルチ GPU トレーニング戦略を増やすことができます。

 fabric = Fabric( accelerator="cuda", precisinotallow="bf16-mixed", devices=4, strategy="FSDP" # new! )

06_fabric-vit-mixed-fsdp.py スクリプトの出力。

現在、4 つの GPU を使用することで、コードの実行時間は約 2 分となり、混合精度トレーニングのみを使用した場合と比べて 3 倍近く高速化しました。

データ並列性とテンソル並列性を理解する

データ並列処理では、データのミニバッチが分割され、各 GPU にモデルのコピーが存在します。このプロセスは、複数の GPU で並行して動作することでモデルのトレーニングを高速化します。

データ並列処理の仕組みの概要は次のとおりです。

1. 同じモデルがすべての GPU に複製されます。
2. 各 GPU は、入力データの異なるサブセット (異なるデータのミニバッチ) を受け取ります。
3. すべての GPU は、モデルに対して順方向および逆方向の伝播を独立して実行し、独自のローカル勾配を計算します。
4. 勾配はすべての GPU にわたって収集され、平均化されます。
5. 平均勾配はモデルのパラメータを更新するために使用されます。

各 GPU は異なるデータ サブセットを並列に処理し、勾配を平均化してパラメータを更新することでモデル全体のトレーニング プロセスが高速化されます。

このアプローチの主な利点は速度です。各 GPU が異なるデータのミニバッチを同時に処理するため、モデルはより短時間でより多くのデータを処理できます。これにより、特に大規模なデータセットを扱う場合に、モデルのトレーニングに必要な時間を大幅に短縮できます。

ただし、データの並列処理にもいくつかの制限があります。最も重要なのは、各 GPU にモデルとパラメータの完全なコピーがなければならないことです。これにより、モデルが単一の GPU のメモリに収まる必要があるため、トレーニングできるモデルのサイズが制限されます。これは、現代の ViT や LLM では実現できません。

データ並列処理とは異なり、テンソル並列処理ではモデル自体が複数の GPU に分割されます。また、データ並列処理では、各 GPU がモデル全体に​​対応する必要があり、大規模なモデルをトレーニングするときに制限となる可能性があります。テンソル並列処理により、モデルを分割し、複数のデバイスにトレーニングを分散することで、単一の GPU では大きすぎる可能性のあるモデルのトレーニングが可能になります。

テンソル並列処理はどのように機能しますか?行列の乗算を想像してください。計算を分散する方法は、行ごとと列ごとの 2 つがあります。簡単にするために、計算を列ごとに分散することを検討してください。たとえば、下の図に示すように、大規模な行列乗算演算を複数の独立した計算に分解し、それぞれを異なる GPU で実行することができます。次に、結果を連結して結果を取得し、計算負荷を効果的に軽減します。