使用 TorchDistributor 的分散式訓練

本文說明如何使用 TorchDistributor 在 PyTorch ML 模型上執行分散式定型。

TorchDistributor 是 PySpark 中的開放原始碼模組，可協助使用者在 Spark 叢集上使用 PyTorch 進行分散式訓練，因此可讓您以 Spark 作業的形式啟動 PyTorch 訓練作業。 在背景下，它會初始化背景工作角色之間的環境和通道，並利用CLI命令 torch.distributed.run 跨背景工作節點執行分散式訓練。

TorchDistributor API 支援下表所示的方法。

方法和簽章	描述
init(self, num_processes, local_mode, use_gpu)	建立 TorchDistributor 的實例。
run(self, main, *args)	藉由叫用 main(**kwargs) main 是函式來執行分散式定型，並在 main 是檔案路徑時執行 CLI 命令 torchrun main *args 。

需求

• Spark 3.4
• Databricks Runtime 13.0 ML 或更新版本

筆記本的開發工作流程

如果模型建立和定型程式完全從本機計算機上的筆記本或 Databricks Notebook 發生，您只需要進行次要變更，即可讓您的程式代碼準備好進行分散式定型。

1. 準備單一節點程序代碼： 使用 PyTorch、PyTorch Lightning 或其他以 PyTorch/PyTorch Lightning 為基礎的架構來準備及測試單一節點程式代碼，例如 HuggingFace 定型器 API。

2. 準備標準分散式定型的程式代碼： 您需要 將單一程式定型轉換為分散式定型。 讓此分散式程式代碼全部包含在一個定型函式中，您可以搭配 TorchDistributor使用。

3. 在定型函式內移動匯入： 在定型函式中新增必要的匯入，例如 import torch。 這樣做可讓您避免常見的挑剔錯誤。 此外， 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. 建立新的 Notebook 初始化存放庫中的新 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)

疑難排解

筆記本工作流程的常見錯誤是在執行分散式定型時找不到或挑選物件。 當連結庫匯入語句未散發至其他執行程式時，就會發生這種情況。

若要避免此問題，請在定型方法中呼叫TorchDistributor(...).run(<func>)的定型函式頂端，以及包含在定型方法中呼叫的任何其他使用者定義函式內的所有匯入語句（例如import torch， ）。

Notebook 範例

下列筆記本範例示範如何使用 PyTorch 執行分散式定型。

Databricks 筆記本上的端對端分散式訓練

取得筆記本

分散式微調擁抱臉部模型筆記本

取得筆記本

PyTorch 檔案筆記本上的分散式訓練

取得筆記本

使用 PyTorch 閃電筆記本進行分散式訓練

取得筆記本

使用 Petastorm 筆記本載入分散式數據

取得筆記本