Verteiltes Training mit TorchDistributor
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.
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.
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
TorchDistributorverwenden können.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 diedevice_id, an die Modelle und Daten gebunden sind, bestimmt durch:device_id = int(os.environ["LOCAL_RANK"])Starten des verteilten Trainings: Instanziieren Sie den
TorchDistributormit 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:
- Importieren des Repositorys: Importieren Sie das externe Repository als Databricks-Git-Ordner.
- Erstellen eines neuen Notebooks Initialisieren Sie ein neues Azure Databricks Notebook im Repository.
- Starten des verteilten Trainings Rufen Sie
TorchDistributorin 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.
Beispielnotebooks
In den folgenden Notebookbeispielen wird veranschaulicht, wie verteiltes Training mit PyTorch ausgeführt wird.