次の方法で共有

TorchDistributor を使用した分散トレーニング

  • [アーティクル]

この記事では、TorchDistributor を使用して PyTorch ML モデルによる分散トレーニングを行う方法について説明します。

TorchDistributor は、ユーザーが自分の Spark クラスターで PyTorch を使用して分散トレーニングを行うのに役立つ PySpark のオープンソース モジュールであるため、Spark ジョブとして PyTorch トレーニング ジョブを起動できます。 その内部では、ワーカー間の環境と通信チャネルが初期化され、CLI コマンド torch.distributed.run を使用してワーカー ノード間で分散トレーニングが実行されます。

TorchDistributor API は、次の表に示されているメソッドをサポートします。

メソッドとシグネチャ 説明
init(self, num_processes, local_mode, use_gpu) TorchDistributor のインスタンスを作成してください。
run(self, main, *args) メインが関数の場合は main(**kwargs) を呼び出して分散トレーニングを実行し、メインがファイル パスの場合は CLI コマンド torchrun main *args を実行してください。

要件

  • Spark 3.4
  • Databricks Runtime 13.0 ML 以降

ノートブックの開発ワークフロー

モデルの作成とトレーニングのプロセスが、完全にローカル コンピューター上のノートブックまたは Databricks Notebook から行われる場合は、コードを分散トレーニング用に準備するために、小さな変更を加えるだけで済みます。

  1. 単一ノード コードを準備する: PyTorch、PyTorch Lightning、または PyTorch/PyTorch Lightning に基づく他のフレームワーク (HuggingFace Trainer API など) を使用して、単一ノード コードを準備してテストしてください。

  2. 標準の分散トレーニング用にコードを準備する:単一プロセス トレーニングを分散トレーニングに変換する必要があります。 この分散されたコードをすべて、TorchDistributor で使用できる 1 つのトレーニング関数内に含めてください。

  3. トレーニング関数内でインポートを移動する: トレーニング関数内に、import torch などの必要なインポートを追加してください。 これにより、一般的な pickle エラーを回避できます。 さらに、モデルとデータが関連付けられている device_id は、次によって決定されます。

    device_id = int(os.environ["LOCAL_RANK"])

  4. 分散トレーニングを開始する: 望ましいパラメーターを使用して TorchDistributor をインスタンス化し、.run(*args) を呼び出してトレーニングを開始してください。

トレーニング コードの例を次に示します。

from pyspark.ml.torch.distributor import TorchDistributor

def train(learning_rate, use_gpu):
  import torch
  import torch.distributed as dist
  import torch.nn.parallel.DistributedDataParallel as DDP
  from torch.utils.data import DistributedSampler, DataLoader

  backend = "nccl" if use_gpu else "gloo"
  dist.init_process_group(backend)
  device = int(os.environ["LOCAL_RANK"]) if use_gpu  else "cpu"
  model = DDP(createModel(), **kwargs)
  sampler = DistributedSampler(dataset)
  loader = DataLoader(dataset, sampler=sampler)

  output = train(model, loader, learning_rate)
  dist.cleanup()
  return output

distributor = TorchDistributor(num_processes=2, local_mode=False, use_gpu=True)
distributor.run(train, 1e-3, True)

外部リポジトリからトレーニングを移行する

既存の分散トレーニング手順が外部リポジトリに保存されている場合は、次の手順を行うと Azure Databricks に簡単に移行できます。

  1. リポジトリをインポートする: 外部リポジトリを Databricks Git フォルダーとしてインポートしてください。
  2. 新しいノートブックを作成する リポジトリ内で新しい Azure Databricks Notebook を初期化してください。
  3. 分散トレーニングを開始する ノートブック セルで、次のように TorchDistributor を呼び出してください。
from pyspark.ml.torch.distributor import TorchDistributor

train_file = "/path/to/train.py"
args = ["--learning_rate=0.001", "--batch_size=16"]
distributor = TorchDistributor(num_processes=2, local_mode=False, use_gpu=True)
distributor.run(train_file, *args)

トラブルシューティング

ノートブック ワークフローの一般的なエラーは、分散トレーニングの実行時にオブジェクトが見つからないか、pickle できないことです。 これは、ライブラリの import ステートメントが他の実行プログラムに配布されていない場合に発生することがあります。

この問題を回避するには、TorchDistributor(...).run(<func>) で呼び出されるトレーニング関数の先頭と、トレーニング メソッド内で呼び出されるその他のユーザー定義関数の内部の "両方" に、すべての import ステートメント (例えば、import torch) を含めてください。

NCCL エラー: ncclInternalError: Internal check failed.

マルチノード トレーニング中にこのエラーが発生した場合、それは通常、GPU 間のネットワーク通信に問題があることを示します。 この問題は、NCCL (NVIDIA Collective Communications Library) が GPU 通信に特定のネットワーク インターフェイスを使用できない場合に発生します。

このエラーを解決するには、次のスニペットをトレーニング コードに追加して、プライマリ ネットワーク インターフェイスを使用します。

import os
os.environ["NCCL_SOCKET_IFNAME"] = "eth0"

サンプルの Notebook

次のノートブックの例は、PyTorch を使用して分散トレーニングを実行する方法を示しています。

Databricks ノートブックでのエンドツーエンドの分散トレーニング

ノートブックを入手

Hugging Face モデル ノートブックの分散微調整

ノートブックを入手

PyTorch ファイル ノートブックでの分散トレーニング

ノートブックを入手

PyTorch Lightning ノートブックを使用した分散トレーニング

ノートブックを入手