Google
Cloud Platform
Amazon Web Services
Microsoft Azureトーチディストリビューター による分散トレーニング
この記事では、 TorchDistributor を使用して PyTorch 機械学習モデルで分散トレーニングを実行する方法について説明します。
TorchDistributor は PySpark のオープンソース モジュールであり、ユーザーが Spark クラスターで PyTorch を使用して分散トレーニングを行うのに役立つため、PyTorch トレーニング ジョブを Spark ジョブとして起動できます。 内部的には、ワーカー間の環境と通信チャネルを初期化し、CLIコマンド torch.distributed.run を利用してワーカーノード間で分散トレーニングを実行します。
トーチディストリビューター API は、次の表に示すメソッドをサポートしています。
メソッドと署名 |
説明 |
|---|---|
|
トーチディストリビューターのインスタンスを作成します。 |
|
main が関数の場合は |
ノートブック の開発ワークフロー
モデルの作成とトレーニング プロセスが完全にローカル コンピューター上のノートブックまたは Databricks ノートブックから行われる場合は、小さな変更を加えるだけで、コードを分散トレーニング用に準備できます。
単一ノード コードを準備します。 PyTorch、PyTorch Lightning、または HuggingFace Trainer API のような PyTorch/PyTorch Lightning に基づくその他のフレームワークを使用して、単一ノード コードを準備してテストします。
標準の分散トレーニング用のコードを準備します。 単一プロセスのトレーニングを分散トレーニングに変換する必要があります。この分散コードをすべて、
TorchDistributorで使用できる 1 つのトレーニング関数に含めます。トレーニング関数内でインポートを移動する: トレーニング関数内で
import torchなどの必要なインポートを追加します。 そうすることで、一般的な酸洗いエラーを回避できます。 さらに、モデルとデータが関連付けられるdevice_idは、次の要素によって決まります。device_id = int(os.environ["LOCAL_RANK"])
分散トレーニングの開始: 目的のパラメーターを使用して
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)
外部リポジトリ からのトレーニングの移行
外部リポジトリに格納されている既存の分散トレーニング手順がある場合は、次の手順を実行して Databricks に簡単に移行できます。
リポジトリをインポートする:外部リポジトリをDatabricks Git フォルダーとしてインポートします。
新しいノートブックを作成する リポジトリ内の新しい Databricks ノートブックを初期化します。
分散トレーニングの開始 ノートブックのセルで、次のように
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)
