DatabricksにおけるTorchDistributorを用いた分散トレーニング

Distributed training with TorchDistributor | Databricks on AWS [2023/3/30時点]の翻訳です。

本書は抄訳であり内容の正確性を保証するものではありません。正確な内容に関しては原文を参照ください。

プレビュー
本機能はパブリックプレビューです。

本書では、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が関数の場合はmain(**kwargs)を呼び出すことで分散トレーニングを実行し、mainがファイルパスの場合はCLIコマンドtorchrun main *argsを実行します。

要件

• Spark 3.4
• Databricks Runtime 13.0 ML以降

ノートブックにおける開発ワークフロー

モデルの作成、トレーニングプロセスのすべてが、ご自身のローカルマシンのノートブック、あるいはDatabricksノートブックで行われる場合、分散トレーニングを行えるようにするには、少々のコード変更のみです。

1. シングルノードコードの準備: PyTorch、PyTorch Lightning、HuggingFace Trainer APIのようにPyTorch/PyTorch Lightningをベースとしたその他のフレームワークを用いたシングルノードのコードを準備、テストします。

2. 標準的な分散トレーニングのコードの準備: シングルプロセスのトレーニングを分散トレーニングに変換する必要があります。TorchDistributorで使用できるように、一つのトレーニング関数の中にすべての分散処理コードを含めるようにします。

3. import文をトレーニング関数内に移動: import torchのように必要なimport文をトレーニング関数内に追加します。こうすることで、よくあるpickleのエラーを回避することができます。さらに、モデルとデータが紐づけられるdevice_idは以下によって決定されます:

   Python

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

4. 分散トレーニングの起動: 必要なパラメーターを指定してTorchDistributorのインスタンスを作成し、トレーニングを起動するために.run(*args)を呼び出します。

以下にトレーニングコードのサンプルを示します:

Python

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に移行することができます:

1. リポジトリのインポート: Databricks Repoとして外部リポジトリをインポートします。
2. 新規ノートブックの作成: リポジトリ内で新規Databricksノートブックを作成します。
3. 分散トレーニングの起動: ノートブックのセルで、以下のようにTorchDistributorを呼び出します。

Python

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文が他のエグゼキューターに分散されない場合に生じることがあります。

この問題を避けるために、すべてのimport文(import torchなど)をTorchDistributor(...).run(<func>)で呼び出されるトレーニング関数の最初の部分と、トレーニングメソッドで呼び出されるその他のユーザー定義関数内の両方に含めます。

サンプルノートブック

Databricksノートブックにおけるエンドツーエンドの分散トレーニングノートブック

PyTorchファイルにおける分散トレーニング

PyTorch Lightningを用いた分散トレーニング

Databricksクイックスタートガイド

Databricksクイックスタートガイド

Databricks無料トライアル

Databricks無料トライアル