トレーニング速度が60%アップ！わずか5行のコードで、PyTorch 1.6は自動混合精度トレーニングをネイティブにサポートします。

[[333418]]

PyTorch 1.6 ナイトリーでは、自動混合精度トレーニングをサポートするサブモジュール amp が追加されます。期待する価値あり。パフォーマンスと、Nvidia Apex と比べてどのような利点があるのか​​を見てみましょう。

PyTorch 1.6 でリリースされる torch.cuda.amp 混合精度トレーニング モジュールは、わずか数行のコードを追加するだけで、大規模モデルのトレーニングを 50 ～ 60% 高速化するという約束を果たします。

PyTorch 1.6 で追加される予定の最もエキサイティングな機能の 1 つは、自動混合精度トレーニングのサポートです。

混合精度トレーニングは、PyTorch のデフォルトである単精度浮動小数点数 fp32 に代わる半精度浮動小数点数 fp16 で可能な限り多くの演算を実行することで、ニューラル ネットワークのトレーニング時間を大幅に短縮する手法です。最新世代の NVIDIA GPU には、高速 fp16 行列演算用に設計された専用のテンソル コアが搭載されています。

しかし、これまでこれらの Tensor Core は、精度を下げた演算をモデルに手動で記述する必要があったため、使いにくいままでした。ここで、自動化された混合精度トレーニングが役立ちます。今後リリースされる torc h.cuda.amp API を使用すると、わずか 5 行のコードでトレーニング スクリプトに混合精度トレーニングを実装できるようになります。

混合精度の仕組み

混合精度トレーニングの仕組みを理解する前に、まず浮動小数点数を確認する必要があります。

コンピュータ エンジニアリングでは、1.0151 や 566132.8 などの 10 進数は、伝統的に浮動小数点数として表されます。無限に正確な数値（π など）を作成できますが、それを保存するためのスペースは有限であるため、精度（数値を四捨五入する前に含めることができる小数点以下の桁数）とサイズ（数値を保存するために使用する桁数）の間で妥協する必要があります。

浮動小数点数の技術標準 IEEE 754 では、次の標準が設定されています。fp64 (倍精度または「double」とも呼ばれます) の最大丸め誤差は約 2^-52 です。fp32 (単精度または「single」とも呼ばれます) の最大丸め誤差は約 2^-23 です。fp16 (半精度または「half」とも呼ばれます) の最大丸め誤差は約 2^-10 です。

Python の float 型は fp64 ですが、メモリに敏感な PyTorch ではデフォルトの dtype として fp32 が使用されます。

混合精度トレーニングの基本的な考え方はシンプルです。精度を半分にし (fp32 → fp16)、トレーニング時間を半分にします。

最も難しいのは、これを安全に行うことです。

浮動小数点数が小さくなるほど、丸め誤差が大きくなることに注意してください。 「十分に小さい」浮動小数点数に対して実行される演算は、値をゼロに丸めます。これはアンダーフローと呼ばれ、バックプロパゲーション中の勾配更新の多くまたはほとんどが非常に小さくなるものの、ゼロにはならないため問題となります。丸め誤差はバックプロパゲーション中に蓄積され、これらの数値を 0 または NaN に変える可能性があります。これにより、勾配の更新が不正確になり、ネットワークが収束しなくなる可能性があります。

2018 年の ICLR 論文「Mixed Precision Training」では、あらゆる場所で fp16 を使用するだけで、2^-24 未満の勾配更新 (サンプル ネットワークのすべての勾配更新の約 5%) が「飲み込まれる」ことがわかりました。

混合精度トレーニングは、モデル トレーニングが発散することなく fp16 を使用できるようにする一連のテクニックです。 3つの異なる技術を組み合わせたものです。

まず、重み行列のコピーを 2 つ保持します。1 つは fp32 の「マスター コピー」、もう 1 つは fp16 の半精度コピーです。勾配の更新は fp16 行列を使用して計算されますが、fp32 行列で更新されます。これにより、グラデーション更新の適用がより安全になります。

2 番目に、異なるベクトル演算は異なる速度でエラーを蓄積するため、異なる方法で処理する必要があります。一部の操作は fp16 では常に安全ですが、他の操作は fp32 でのみ安全です。ニューラル ネットワーク全体を fp16 で実行するのではなく、一部の部分に半精度を使用し、他の部分に単精度を使用する方が適切です。この dtype の混合のため、この手法は「混合精度」と呼ばれます。

3 番目に、損失スケーリングを使用します。損失スケーリングとは、バックプロパゲーションを実行する前に、損失関数の出力を何らかのスカラー数 (論文では 8 から始めることを推奨) で乗算することを指します。乗法的に増加する損失値は乗法的に増加する勾配更新を生成し、多くの勾配更新を fp16 安全しきい値 2^-24 を超えて「ブースト」します。勾配更新を適用する前に必ずスケーリングを元に戻すようにしてください。また、スケーリングをあまり大きく選択しすぎて inf 重み更新 (オーバーフロー) を生成し、ネットワークが反対方向に発散する原因とならないようにしてください。

これら 3 つの手法を組み合わせることで、著者らはさまざまなネットワークを非常に短い時間で収束するようにトレーニングすることができます。ベンチマークに関しては、わずか 9 ページのこの論文を読むことをお勧めします。

テンソルコアはどのように機能しますか?

混合精度トレーニングではメモリが節約されますが (fp16 行列は fp32 行列の半分のサイズになります)、特別な GPU サポートがなければモデル トレーニングは高速化されません。チップ上には、半精度演算を高速化できる何かが必要です。最近の世代の NVIDIA GPU では、これは Tensor コアと呼ばれています。

Tensor コアは、2 つの 4 × 4 fp16 行列を乗算し、その結果を 3 つ目の 4 × 4 fp16 または fp32 行列に加算する (「融合乗算加算」) という非常に特殊な操作に最適化された新しいタイプの処理ユニットです。

この演算を基本的な構成要素として使用して、より大きな fp16 行列乗算演算を実装できます。ほとんどのバックプロパゲーションは行列乗算に簡略化できるため、Tensor コアはネットワーク内のほぼすべての計算集約型レイヤーに適用できます。

注意: 入力行列は fp16 である必要があります。 Tensor コアを搭載した GPU でトレーニングを行っており、混合精度トレーニングを使用していない場合は、グラフィック カードの性能を 100% 引き出せない可能性があります。 fp32 で定義された標準の PyTorch モデルは、チップに fp16 の計算を適用することはないため、非常に強力な Tensor コアはすべてアイドル状態になります。

Tensor コアは、2017 年後半に前世代の Volta アーキテクチャで導入され、現世代の Turing でいくつかの改善が見られ、今後の Ampere でもさらに改善される予定です。クラウドで一般的に利用できる 2 つの GPU は、V100 (5120 CUDA コア、600 テンソル コア) と T4 (2560 CUDA コア、320 テンソル コア) です。

覚えておく価値のあるパズルのもう 1 つのピースはファームウェアです。 Tensor Core 操作は CUDA 7.0 以降でサポートされていますが、初期の実装にはバグが多いと言われているため、CUDA 10.0 以降を使用することが重要です。

Pytorch 自動混合精度はどのように機能しますか?

この重要な背景を踏まえて、ようやく新しい PyTorch amp API を詳しく調べ始めることができます。

混合精度トレーニングは技術的にはずっと可能でした。つまり、fp16 でネットワークの一部を手動で実行し、損失スケーリングを自分で実装するのです。自動混合精度トレーニングの興味深い部分は、「自動」の部分です。学習する必要がある新しい API プリミティブは、torch.cuda.amp.GradScalar と torch.cuda.amp.autocast だけです。混合精度トレーニングを有効にするのは、トレーニング スクリプトの適切な場所にドロップするだけです。

例として、混合精度トレーニングを使用するネットワークのトレーニング ループの一部を以下に示します。 #NEW マーカーは、新しいコードが追加された場所を示します。

自己トレーニング()
 X = torch.tensor(X, dtype=torch.float32)
 y = torch.tensor(y, dtype=torch.float32) 
 
 
オプティマイザー = torch.optim.Adam(self.parameters(), lr=self.max_lr)
スケジューラ = torch.optim.lr_scheduler.OneCycleLR(
    オプティマイザー、self.max_lr、
    サイクルモメンタム=False、
    エポック = self.n_epochs、
    ステップごとのエポック = int (np.ceil(len(X) / self.batch_size))、
 ）
バッチ = torch.utils.data.DataLoader(
    torch.utils.data.TensorDataset(X, y)、
    batch_size=self.batch_size、shuffle=True の場合
） 
 
 
 ＃新しい
スケーラー = torch.cuda.amp.GradScaler() 
 
 
範囲内のエポックの場合(self.n_epochs):
 i、(X_batch、y_batch)をenumerate(batches) で指定します:
        X_batch = X_batch.cuda()
        y_batch = y_batch.cuda()
        オプティマイザ.zero_grad() 
 
 
        ＃新しい
        torch.cuda.amp.autocast() を使用する場合:
            y_pred = モデル(X_batch).squeeze()
            損失 = self.loss_fn(y_pred, y_batch) 
 
 
        ＃新しい
        スケーラー.スケール(損失).後方()
        lv = loss.detach().cpu().numpy()
 i % 100 == 0の場合:
            print(f "エポック {epoch + 1}/{self.n_epochs}; バッチ {i}; 損失 {lv}" ) 
 
 
        ＃新しい
        スケーラー.ステップ(オプティマイザー)
        スケーラー.更新() 
 
        スケジューラ

新しい PyTorch GradScaler オブジェクトは、損失スケーリングの PyTorch 実装です。 「混合精度の仕組み」のセクションで説明したように、トレーニング中に勾配がゼロに縮小するのを防ぐために、何らかの形のスケーリングが必要です。最適な損失乗数は、非常に小さな勾配を保持するのに十分高いですが、非常に大きな勾配が無限大に丸められ、逆の問題が発生するほど高くはありません。

PyTorch は指数バックオフを使用してこの問題を解決します。 Gradscalar は小さな損失乗数から始まり、そのたびに 2 倍になります。この段階的な倍増動作は、GradScalar が inf 値を含む勾配更新に遭遇するまで継続されます。 Gradscalar はこのデータ バッチを破棄し (つまり、勾配の更新をスキップし)、損失乗数を半分に削減し、倍増時間をリセットします。

このように損失乗数を上下させることにより、PyTorch は時間の経過とともに適切な損失乗数を近似することができます。 TCP 輻輳制御に精通している読者にとっては、ここでの核となる考え方は非常に馴染み深いものとなるはずです。アルゴリズムで使用される正確な数は設定可能で、ドキュメント文字列からデフォルト値を直接確認できます。

 torch.cuda.amp.GradScaler(
    init_scale= 65536.0 、growth_factor= 2.0 、backoff_factor= 0.5 、
    growth_interval = 2000 、有効 = True
 ）

Gradscalar では、その操作を実装するために、勾配更新計算 (オーバーフローのチェック) とオプティマイザー (破棄されたバッチを no-op に変換する) を制御する必要があります。このため、 loss.backwards() は scaler.scale(loss).backwards() に置き換えられ、 optimizer.step() は scaler.step(optimizer) に置き換えられます。

GradScalar はオーバーフローを検出して停止できますが (inf は常に不良であるため)、アンダーフローを検出して停止することはできません (0 は通常有効な値であるため)。初期値を低く設定しすぎ、成長間隔を長く設定しすぎると、GradScalar が介入する前にネットワークがアンダーフローして発散する可能性があります。このため、非常に大きな初期値を選択するのがよいでしょう。

最後に、GradScalar はステートフル オブジェクトであることに注意してください。この関数を使用してモデル チェックポイントを保存するには、モデルの重みをディスクに書き込んで読み取る必要があります。これは、state_dict および load_state_dict オブジェクト メソッド (PyTorch ドキュメントで説明) を使用して簡単に実行できます。

自動混合精度トレーニング パズルのもう半分は、torch.cuda.amp.autocast コンテキスト マネージャーです。 Autocast は fp32->fp16 変換を実装します。 「混合精度の仕組み」で説明したように、異なる演算ではエラーの蓄積率が異なるため、すべての演算を fp16 で安全に実行できるわけではありません。 amp モジュールのドキュメントからの次のスクリーンショットは、PyTorch で利用可能なさまざまな操作で autocast がどのように機能するかを示しています。

このリストは主に行列の乗算と畳み込み、および単純な線形関数で構成されています。

これらの操作は fp16 では安全ですが、入力が fp16 と fp32 の混合である場合でも動作が中断しないようにアップキャスト ルールが適用されます。このリストには、行列/ベクトルの内積とベクトルの外積という、他の 2 つの基本的な線形代数演算も含まれていることに注意してください。

対数、指数、三角関数、正規関数、離散関数、および (大きな) 合計は fp16 では安全ではないため、fp32 で実行する必要があります。

このリストを見ると、ほとんどのレイヤーは基本的な線形代数演算に内部的に依存しているため、自動キャストの恩恵を受けると思われますが、ほとんどのアクティベーション関数はそうではありません。畳み込み層が大きな勝者です。

自動キャストを有効にするのは非常に簡単です。必要なのは、モデルのフォワード パスを自動キャスト コンテキスト マネージャーでラップすることだけです。

 torch.cuda.amp.autocast() を使用する場合:
    y_pred = モデル(X_batch).squeeze()
    損失 = self.loss_fn(y_pred, y_batch)

このようにフォワード パスをラップすると、バックワード パスの自動キャスト (loss.backwards() など) が自動的にオンになるため、autocast を 2 回呼び出す必要がなくなります。

PyTorch のベスト プラクティス (インプレース操作の回避など) に従っている限り、自動キャストは基本的に「正常に機能します」。マルチ GPU DistributedDataParallel API を使用することも可能です (GPU ごとに 1 つのプロセスのみを使用するという推奨戦略に従う限り)。少し調整するだけで、マルチ GPU DataParallel API も使用できます。 Pytorch ドキュメントの自動混合精度の例ページの「複数の GPU の操作」セクションは、このトピックに関する便利なリファレンスです。私の意見では、覚えておくべき重要なポイントの 1 つは、「バイナリ クロス エントロピーではなく、ロジットによるバイナリ クロス エントロピーを優先する」ということです。

ベンチマーク

ここまでで、混合精度とは何か、Tensor Core とは何か、そして PyTorch API が自動混合精度をどのように実装するかについて学びました。残っているのは、実際のパフォーマンスベンチマークを確認することだけです。

私は、V100 (旧世代の Tensor コア) と T4 (現世代の Tensor コア) の Spell API を使用して、3 つの非常に異なるニューラル ネットワークを自動混合精度でトレーニングし、1 回は自動混合精度なしでトレーニングしました。最新の PyTorch 1.6 nightly と CUDA 10.0 を搭載した AWS EC2 インスタンス (それぞれ p3.2xlarge と g4dn.xlarge) を使用しました。すべてのモデルの収束は一貫しています。つまり、どのモデルでも、混合精度ネットワークと元のネットワーク間のトレーニング損失に違いは見られません。トレーニングされたネットワークは次のとおりです。

• Feedforward は、Rossman Store Samples Kaggle コンペティションのデータを使用してトレーニングされたフィードフォワード ニューラル ネットワークです。
• UNet、中規模のバニラUNet画像セグメンテーションネットワーク。Segmented Bob Ross Imagesデータセットでトレーニング済み。
• BERT、bert-base-uncasedバックボーン（huggingface経由）を使用した大規模なNLP変換モデル、およびKaggleコンペティションTwitter Sentiment Extractionのデータ

結果は次のとおりです。

フィードフォワード ネットワークは非常に小さいため、混合精度トレーニングを行ってもメリットはありません。

UNet は、7,703,497 個のパラメータを持つ中規模の畳み込みモデルであり、混合精度トレーニングから大きなメリットを得られます。興味深いことに、V100 と T4 はどちらも混合精度トレーニングの恩恵を受けますが、その恩恵は T4 の方がはるかに大きく、5% の時間節約に対して最大 30% の時間節約となっています。

BERT は大規模なモデルであり、時間を節約するためにここで混合精度トレーニングを使用すると、中規模モデルでは「あれば便利」から「必須」になります。 Volta または Turing GPU でトレーニングする場合、自動混合精度により、大規模モデルのトレーニング時間が 50% ～ 60% 短縮されます。

これは、特に、追加される複雑さが最小限であることを考慮すると、大きな利点です。モデルのトレーニング スクリプトを変更するコードは 4 行か 5 行だけです。私の意見では：

混合精度は、モデル トレーニング スクリプトに対して行う最初のパフォーマンス最適化の 1 つです。

メモリはどうですか?

「混合精度の仕組み」セクションで説明したように、fp16 行列はメモリ内で fp32 行列の半分のサイズなので、混合精度トレーニングのもう 1 つの利点はメモリ使用量です。 GPU メモリのボトルネックは GPU の計算能力よりもはるかに小さいですが、最適化の価値はまだたくさんあります。メモリを効率的に使用すればするほど、GPU で使用できるバッチ サイズが大きくなります。

PyTorch は、モデルのトレーニング プロセスの開始時に一定量の GPU メモリを予約し、トレーニング中はこのメモリを保持します。これにより、トレーニング中に他のプロセスが GPU メモリを過剰に取得し、PyTorch トレーニング スクリプトが OOM エラーでクラッシュすることが防止されます。

混合精度トレーニングを有効にすると、PyTorch のメモリ保持動作にどのような影響があるかを示します。

興味深いことに、両方の大規模モデルは混合精度への切り替えによるメリットを享受していますが、UNet は BERT よりも切り替えによるメリットをはるかに多く享受しています。 PyTorch のメモリ割り当ての動作は私にとって非常にわかりにくいので、なぜこのようなことが起こるのか全くわかりません。

要約する

自動混合精度トレーニングは、近日リリース予定の PyTorch 1.6 の使いやすく強力な新機能であり、最新の NVIDIA GPU で大規模モデルのトレーニングを 60% 高速化することが期待されています。

このテクノロジーはしばらく前から存在していましたが、これまでネイティブの PyTorch API がなかったため、一般ユーザーが簡単に利用できるものではありませんでした。

ソース コードから直接混合精度トレーニングについて詳しく知るには、PyTorch マスター ドキュメントの自動混合精度パッケージと自動混合精度の例のページを参照してください。

この機能を自分でテストしてみませんか?最新の PyTorch ナイトリーをインストールするのは非常に簡単です。インストール方法については、PyTorch ホームページの手順を参照してください。

これらのベンチマークを自分で再現してみませんか?すべてのモデル ソース コードは、GitHub の ResidentMario/spell-feedforward-rossman、ResidentMario/spell-unet-bob-ross、および ResidentMario/spell-tweet-sentiment-extraction リポジトリで入手できます。