共用方式為

使用 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 是函式) 來執行分散式訓練，並執行 CLI 命令 torchrun main *args (如果 main 是檔案路徑)。

需求

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

筆記本的開發工作流程

如果模型建立和訓練程序完全從本機電腦上的筆記本或 Databricks 筆記本發生，您只需要進行細微變更，即可讓您的程式碼準備好進行分散式訓練。

  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. 建立新筆記本：初始化存放庫中的新 Azure Databricks 筆記本。
  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)。

NCCL 失敗：ncclInternalError: Internal check failed.

在多節點訓練期間遇到此錯誤時，通常表示 GPU 之間的網路通訊出現了問題。 當 NCCL (NVIDIA 集體通訊程式庫) 無法使用特定網路介面進行 GPU 通訊時，就會發生此問題。

若要解決此錯誤，請在訓練程式碼中新增下列片段，以使用主要網路介面。

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

筆記本範例

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

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

取得筆記本

分散式微調 Hugging Face 模型筆記本

取得筆記本

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

取得筆記本

使用 PyTorch Lightning 的分散式訓練筆記本

取得筆記本