Freigeben über

Verteiltes Training mit TorchDistributor

  • Artikel

In diesem Artikel wird beschrieben, wie Sie mit TorchDistributor verteiltes Training für PyTorch-ML-Modelle durchführen.

TorchDistributor ist ein Open-Source-Modul in PySpark, mit dem Benutzer verteilte Schulungen mit PyTorch auf ihren Spark-Clustern durchführen können, sodass Sie PyTorch-Trainingsaufträge als Spark-Aufträge starten können. Im Hintergrund initialisiert er die Umgebung und die Kommunikationskanäle zwischen den Workern und verwendet den CLI-Befehl torch.distributed.run, um verteiltes Training auf den Workerknoten auszuführen.

Die TorchDistributor-API unterstützt die Methoden, die in der folgenden Tabelle aufgeführt sind.

Methode und Signatur Beschreibung
init(self, num_processes, local_mode, use_gpu) Erstellen Sie eine Instanz von TorchDistributor.
run(self, main, *args) Führt das verteilte Training aus, indem main(**kwargs) aufgerufen wird, wenn Main eine Funktion ist, und führt den CLI-Befehl torchrun main *args aus, wenn Main ein Dateipfad ist.

Anforderungen

  • SPARK 3.4
  • Mindestens Databricks Runtime 13.0 ML

Entwicklungsworkflow für Notebooks

Wenn der Modellerstellungs- und Trainingsvorgang vollständig von einem Notebook auf Ihrem lokalen Computer oder einem Databricks-Notebook aus erfolgt, müssen Sie nur geringfügige Änderungen vornehmen, um Ihren Code für das verteilte Training vorzubereiten.

  1. Vorbereiten des Codes für einen einzelnen Knoten: Bereiten Sie den Einzelknotencodes mit PyTorch, PyTorch Lightning oder anderen Frameworks vor, die auf PyTorch/PyTorch Lightning wie der HuggingFace Trainer-API basieren, und testen Sie diese.

  2. Vorbereiten von Code für das verteilte Standardtraining: Sie müssen Ihr Einzelprozesstraining in verteiltes Training konvertieren. Fassen Sie diesen verteilten Code in einer einzigen Trainingsfunktion zusammen, die Sie mit TorchDistributor verwenden können.

  3. Verschieben von Importen innerhalb der Trainingsfunktion: Fügen Sie die erforderlichen Importe, z. B. import torch, innerhalb der Trainingsfunktion hinzu. Dadurch können Sie häufige Picklingfehler vermeiden. Darüber hinaus wird die device_id, an die Modelle und Daten gebunden sind, bestimmt durch:

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

  4. Starten des verteilten Trainings: Instanziieren Sie den TorchDistributor mit den gewünschten Parametern, und rufen Sie .run(*args) auf, um das Training zu starten.

Es folgt ein Trainingscodebeispiel:

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)

Migrieren von Schulungen aus externen Repositorys

Wenn Sie eine vorhandene verteilte Trainingsprozedur in einem externen Repository gespeichert haben, können Sie problemlos zu Azure Databricks migrieren, indem Sie die folgenden Schritte ausführen:

  1. Importieren des Repositorys: Importieren Sie das externe Repository als Databricks-Git-Ordner.
  2. Erstellen eines neuen Notebooks Initialisieren Sie ein neues Azure Databricks Notebook im Repository.
  3. Starten des verteilten Trainings Rufen Sie TorchDistributor in einer Notebookzelle wie folgt auf:
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)

Problembehandlung

Ein häufiger Fehler für den Notebookworkflow besteht darin, dass Objekte beim Ausführen des verteilten Trainings nicht gefunden werden können oder kein Pickling möglich ist. Dies kann geschehen, wenn die Bibliothekimportanweisungen nicht an andere Executors verteilt werden.

Um dieses Problem zu vermeiden, schließen Sie alle Importanweisungen (z. B. import torch) sowohl am Anfang der Trainingsfunktion, die mit TorchDistributor(...).run(<func>) aufgerufen wird, als auch in anderen benutzerdefinierten Funktionen, die in der Trainingsmethode aufgerufen werden, ein.

NCCL-Fehler: ncclInternalError: Internal check failed.

Wenn während der Schulung mit mehreren Knoten dieser Fehler auftritt, weist er in der Regel auf ein Problem mit der Netzwerkkommunikation zwischen GPUs hin. Dieses Problem tritt auf, wenn NCCL (NVIDIA Collective Communications Library) bestimmte Netzwerkschnittstellen für die GPU-Kommunikation nicht verwenden kann.

Um diesen Fehler zu beheben, fügen Sie den folgenden Codeausschnitt in Ihrem Schulungscode hinzu, um die primäre Netzwerkschnittstelle zu verwenden.

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

Beispielnotebooks

In den folgenden Notebookbeispielen wird veranschaulicht, wie verteiltes Training mit PyTorch ausgeführt wird.

Verteiltes End-to-End-Training auf Databricks-Notebook

Notebook abrufen

Verteilte Feinabstimmung eines Notebooks für das Hugging Face-Modell

Notebook abrufen

Verteiltes Training auf einem PyTorch File Notebook

Notebook abrufen

Verteiltes Training mit PyTorch Lightning Notebook

Notebook abrufen