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

[Distributed training with TorchDistributor \| Databricks on AWS](https://docs.databricks.com/machine-learning/train-model/distributed-training/spark-pytorch-distributor.html) [2023/3/30時点]の翻訳です。

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

:::note info
**プレビュー**
本機能は[パブリックプレビュー](https://docs.databricks.com/release-notes/release-types.html)です。
:::

本書では、[TorchDistributor](https://github.com/apache/spark/blob/master/python/pyspark/ml/torch/distributor.py)を用いて、どのようにPyTorch MLモデルの分散トレーニングを実行するのかを説明します。

TorchDistributorは、SparkクラスターでPyTorchを用いた分散トレーニングを支援するPySparkのオープンソースモジュールであり、SparkジョブとしてPyTorchトレーニングを起動することができます。内部では、環境とワーカー間のコミュニケーションチャンネルを初期化し、ワーカーノードにおける分散トレーニングを実行するためにCLIコマンド`torch.distributed.run`を活用しています。

[TorchDistributor API](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.ml.torch.distributor.TorchDistributor.html)は以下の表にあるメソッドをサポートしています。

| メソッドとシグネチャ | 説明 |
|:--|:--|
| `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ノートブック](https://qiita.com/taka_yayoi/items/24a897cf40bba6d9e305)で行われる場合、分散トレーニングを行えるようにするには、少々のコード変更のみです。

1. **シングルノードコードの準備:** PyTorch、PyTorch Lightning、HuggingFace Trainer APIのようにPyTorch/PyTorch Lightningをベースとしたその他のフレームワークを用いたシングルノードのコードを準備、テストします。
1. **標準的な分散トレーニングのコードの準備:** [シングルプロセスのトレーニングを分散トレーニングに変換](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html)する必要があります。`TorchDistributor`で使用できるように、一つのトレーニング関数の中にすべての分散処理コードを含めるようにします。
1. **import文をトレーニング関数内に移動:** `import torch`のように必要なimport文をトレーニング関数内に追加します。こうすることで、よくあるpickleのエラーを回避することができます。さらに、モデルとデータが紐づけられる`device_id`は以下によって決定されます:

    ```py:Python
    device_id = int(os.environ["LOCAL_RANK"])
    ```

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

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

```py: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](https://qiita.com/taka_yayoi/items/b89f199ff0d3a4c16140)として外部リポジトリをインポートします。
1. **新規ノートブックの作成:** リポジトリ内で新規Databricksノートブックを作成します。
1. **分散トレーニングの起動:** ノートブックのセルで、以下のように`TorchDistributor`を呼び出します。

```py: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ノートブックにおけるエンドツーエンドの分散トレーニングノートブック

https://docs.databricks.com/_extras/notebooks/source/deep-learning/torch-distributor-notebook.html

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

https://docs.databricks.com/_extras/notebooks/source/deep-learning/torch-distributor-file.html

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

https://docs.databricks.com/_extras/notebooks/source/deep-learning/torch-distributor-lightning.html

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

[Databricksクイックスタートガイド](https://www.amazon.co.jp/dp/B09V1YXFVQ/)


### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)