PyTorch Windows における分散データ並列処理サポートの概要
Training models is one of the issues that frustrates machine learning developers the most, and that won’t be the case for the foreseeable future. In particular, parts that take a very long time cannot be changed by anything. However, if you have access to a (for now) fortunate environment of multiple GPUs, Distributed Data Parallel (DDP) should help you train your model. Model training for DDP is performed transparently across multiple GPUs. Multiple GPUs can be located on a single machine or isolated on multiple machines. DDP leverages all GPUs to maximize compute power, significantly reducing training time.
Although for quite some time DDP was only available on Linux. Changed in PyTorch 1.7. DDP support on Windows was introduced by Microsoft in PyTorch 1.7 and has been improving ever since. In this article, we’ll discuss how DDP can help you with your training experience on Windows.
tutorial
Set up two reference machines of the same specification, one on Windows and one with Linux, on Azure and model training on the same code and datasets.
We’ll use a very useful Azure resource called Data Science Virtual Machine (DSVM). This is an easy-to-handle VM image with many machine learning tools pre-installed. At the time of writing, PyTorch 1.8.1 (Anaconda), which is included in the DSVM image, will also be used in this demo.
This article uses a “Standard NC24s_v3” size that gives you four NVIDIA Tesla V100 GPUs at your disposal.
To better understand how DDP works, let’s learn some basic concepts.
One important concept to understand is process groups. This is the basic underlying tool for DDP. A process group, as the name implies, is a group of processes. Each process is responsible for training workloads on one dedicated GPU. You also need a way to coordinate groups of processes (and more importantly, the GPUs on which they underpin them) to enable them to communicate with each other. This is called a “backend” in PyTorch (–dist-backend in script parameters). The NCCL and MPI backends are not currently available on Windows, so PyTorch 1.8 uses Gloo as the backend. For more information about the backend, see the PyTorch documentation. Finally, you need an area for the backend to exchange information. This is called a “store” in PyTorch (–dist-url in the script parameter). For more information about stores, see the PyTorch documentation.
多少混乱を招きそうなその他の概念に、「ワールド サイズ」と「ランク」があります。ワールド サイズは、基本的にはトレーニング ジョブに参加するプロセスの数を指します。前述のとおり、各プロセスは 1 つの専用 GPU を担当します。このため、ワールド サイズは使用される GPU の合計数に等しくなると考えると、わかりやすくなると思います。次の「ランク」は、各プロセスのインデックス番号として表すことができ、1 つの特定のプロセスを識別するのに使用することができます。また、すべてのプロセスを調整する「コントローラー」の役割を果たす、ランク 0 のプロセスが常に必要になる点に注意してください。ランク 0 のプロセスが存在しないと、トレーニング全体を正常に実行できません。
必要な基礎知識について学んだら、実際のトレーニングを開始しましょう。データセットには ImageNet 2012 の小規模なサブセットを使用します。データセットをダウンロードして、ファイルシステムのどこかに配置したと想定してください。このデモでは、「D:\imagenet-small」としましょう。
当然ながらトレーニング スクリプトも必要です。ここでは PyTorch Examples リポジトリの ImageNet トレーニング スクリプトを使用し、ターゲット モデルには ResNet50 を使用します。このトレーニング スクリプトは、通常のトレーニング スクリプトに、「torch.distributed」や「torch.multiprocessing」など DDP の性能により提供されるパッケージを加えたものと考えることができます。スクリプトにはそれほど多くのロジックは含まれていないため、これに基づいて独自のスクリプトを簡単に設定できます。詳しくは、こちらの 入門チュートリアルも参考にしてください。
マシン 1 台の場合、FileStore を使用するだけでさらに容易な設定が可能です。コマンド全体は以下のようになります。
「world-size 1」と「rank 0」を使用していることにお気付きかと思います。これは、必要なワールド サイズとランクを、利用可能な GPU に基づいてスクリプトが計算するためです。実際に使用されるワールド サイズは、利用可能な GPU の数と同じ 4 です。また、各プロセスのランクには、0 から始まる適切な数字が自動で割り当てられます。
コマンドライン引数に慣れていない場合は、環境変数を使用して DDP の引数を初期化することもできます。これは、展開を自動化する必要がある場合に便利です。詳しくは、PyTorch ドキュメントの the PyTorch documentationを参照してください。
ここまですべて問題なければ、トレーニング ジョブがすぐに開始されます。
トラブルシューティング
うまくいかない部分がある場合は、以下に示すトラブルシューティングのヒントを参考にしてみてください。
- Windows で FileStore を使用している場合、使用しているファイルがその他のプロセスによってロックされていないことを確認してください。これは、トレーニング プロセスを強制終了した場合に発生する可能性があります。この状況になると、スクリプトが FileStore を初期化できず DDP トレーニング プロセスがフリーズする可能性があります。回避策として、前回のトレーニング プロセスを手動で終了し、ファイルを削除した後に、次のトレーニングを実行してください。
- TcpStore を使用している場合、ネットワークにアクセス可能でポートが利用可能であることを確認してください。利用できない場合、スクリプトが TcpStore を初期化できずトレーニングがフリーズする可能性があります。ランク 0 のプロセスは、指定されたポートにバインドしてリッスンし、その他のプロセスは、そのポートへの接続を試みます。TCP 接続の問題のデバッグには、netstat などのネットワーク監視ツールを使用すると便利です。
- トレーニングの実行中は、nvidia-smi などのツールを使用して GPU 負荷を監視できます。GPU が最大活用され使用率 100% で実行されるのが理想ですが、GPU 負荷が低い場合は、バッチ サイズの拡張や DataLoader ワーカー数の増加が効果的です。
- DDP で使用される GPU の数は、有効なバッチ サイズにも影響する点に注意してください。たとえば 1 基の GPU のバッチ サイズに 128 を使用し、その後 GPU が 2 基の DDP に切り替えたとします。この場合、a) バッチを分割して各 GPU でバッチ サイズ 64 を使用する、b) 各 GPU のバッチ サイズに 128 を使用して、有効バッチ サイズを 256 にする、という 2 つの選択肢があります。GPU メモリの制約を考えなければ、この選択はほぼ自らの裁量で決定できます。いずれかの方法に合わせてスクリプトに調整を加えてください。選択肢 b) を選んだ場合、同様のトレーニング結果を得るためには最初の学習率の調整が必要になる点にもご注意ください。
ベンチマーク
時間 | GPU 1 基 (DDP なし) | GPU 2 基 | GPU 4 基 |
---|---|---|---|
Linux | 56 分 58 秒 | 31 分 7 秒 | 17 分 20 秒 |
Windows | 58 分 55 秒 | 31 分 55 秒 | 19 分 3 秒 |
データからわかるように、GPU を追加することによって、全体的に期待どおりの高速化が実現されました。GPU を 2 基使用することでトレーニング時間はほぼ半分になり、GPU 4 基だと約 1/4 になりました。
損失曲線からわかるのは、トレーニング時間の短縮が不十分なトレーニング結果につながるわけではないということです。時間の経過と共にモデルが徐々に訓練されていくことが期待できます。
今回は、Windows 上の DDP において、Linux 上の DDP に匹敵するパフォーマンス向上を、精度を犠牲にせずに実現できることを、簡単なデモを通して説明いたしました。マイクロソフトは現在も PyTorch チームと緊密に協力しながら、Windows における PyTorch エクスペリエンスの改善に取り組んでいます。Windows 上の DDP がサポートされたことで、トレーニングのパフォーマンスは今後さらに飛躍的に向上するでしょう。ぜひ皆さんにお試しいただき、フィードバックをお寄せいただければと思います。