Multi-GPU-Workload

Von Bedeutung

Dieses Feature befindet sich in der Betaversion. Arbeitsbereichsadministratoren können den Zugriff auf dieses Feature über die Vorschauseite steuern. Siehe Verwalten von Azure Databricks-Vorschauen.

Sie können verteilte Workloads über mehrere GPUs auf einem einzelnen Knoten mithilfe der SERVERless GPU Python-API starten. Die API bietet eine einfache, einheitliche Schnittstelle, die die Details der GPU-Bereitstellung, der Umgebungseinrichtung und der Workloadverteilung abstrahiert. Mit minimalen Codeänderungen können Sie nahtlos von der Einzel-GPU-Schulung zur verteilten Multi-GPU-Ausführung aus demselben Notizbuch wechseln.

Unterstützte Frameworks

Die @distributed API ist in wichtige verteilte Schulungsbibliotheken integriert:

PyTorch Distributed Data Parallel (DDP) – Standardmäßige Multi-GPU-Datenparallelität.
Fully Sharded Data Parallel (FSDP) – Speichereffiziente Schulung für große Modelle.
DeepSpeed – Die Optimierungsbibliothek von Microsoft für umfangreiche Modellschulungen.

serverless_gpu API vs. TorchDistributor

In der folgenden Tabelle wird die serverless_gpu@distributed API mit TorchDistributor verglichen:

Funktion	`serverless_gpu` `@distributed` API	Fackelverteiler
Infrastruktur	Vollständig serverlos, keine Clusterverwaltung	Erfordert einen Spark-Cluster mit GPU-Workern
Konfiguration	Einzelner Dekorierer, minimale Konfiguration	Erfordert Spark-Cluster und TorchDistributor-Einrichtung
Framework-Unterstützung	PyTorch DDP, FSDP, DeepSpeed	Hauptsächlich PyTorch DDP
Laden von Daten	Innendekoror verwendet Unity-Katalogvolumes	Über Spark oder das Dateisystem

Die serverless_gpu API ist der empfohlene Ansatz für neue Deep Learning-Workloads auf Databricks. TorchDistributor bleibt für Workloads verfügbar, die eng mit Spark-Clustern gekoppelt sind.

Schnellstart

Die serverlose GPU-API für verteilte Schulungen wird vorinstalliert, wenn Sie mit einer serverlosen GPU in Databricks-Notizbüchern und -Aufträgen verbunden sind. Wir empfehlen GPU-Umgebung 4 und höher. Um es für verteiltes Training zu nutzen, importieren und verwenden Sie den distributed Dekorateur für die Verteilung Ihrer Trainingsfunktion.

Schließen Sie den Modellschulungscode in eine Funktion ein und dekorieren Sie die Funktion mit dem @distributed Dekorator. Die dekorierte Funktion wird zum Einstiegspunkt für die verteilte Ausführung – alle Schulungslogik, das Laden von Daten und die Modellinitialisierung sollten in dieser Funktion definiert werden.

Warnung

Der gpu_type Parameter in @distributed muss mit dem Beschleunigertyp übereinstimmen, mit dem Ihr Notebook verbunden ist. Beispielsweise erfordert @distributed(gpus=8, gpu_type='H100'), dass Ihr Notebook mit einem H100-Beschleuniger verbunden ist. Die Verwendung eines nicht übereinstimmenden Beschleunigertyps (z. B. das Herstellen einer Verbindung mit A10 beim Angeben von H100) bewirkt, dass die Arbeitslast fehlschlägt.

Der folgende Codeausschnitt zeigt die grundlegende Verwendung von @distributed:

# Import the distributed decorator
from serverless_gpu import distributed

# Decorate your training function with @distributed and specify the number of GPUs and GPU type
@distributed(gpus=8, gpu_type='H100')
def run_train():
    ...

Im Folgenden finden Sie ein vollständiges Beispiel, das ein Multilayer-Perceptron (MLP)-Modell auf 8 H100 GPUs in einem Notebook trainiert:

Richten Sie Ihr Modell ein, und definieren Sie Hilfsfunktionen.


# Define the model
import os
import torch
import torch.distributed as dist
import torch.nn as nn

def setup():
    dist.init_process_group("nccl")
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))

def cleanup():
    dist.destroy_process_group()

class SimpleMLP(nn.Module):
    def __init__(self, input_dim=10, hidden_dim=64, output_dim=1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, output_dim)
        )

    def forward(self, x):
        return self.net(x)

Importieren Sie die serverless_gpu-Bibliothek und das verteilte Modul.
```
import serverless_gpu
from serverless_gpu import distributed
```

Schließen Sie den Modellschulungscode in eine Funktion ein und dekorieren Sie die Funktion mit dem @distributed Dekorator.

@distributed(gpus=8, gpu_type='H100')
def run_train(num_epochs: int, batch_size: int) -> None:
    import mlflow
    import torch.optim as optim
    from torch.nn.parallel import DistributedDataParallel as DDP
    from torch.utils.data import DataLoader, DistributedSampler, TensorDataset

    # 1. Set up multi-GPU environment
    setup()
    device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")

    # 2. Apply the Torch distributed data parallel (DDP) library for data-parellel training.
    model = SimpleMLP().to(device)
    model = DDP(model, device_ids=[device])

    # 3. Create and load dataset.
    x = torch.randn(5000, 10)
    y = torch.randn(5000, 1)

    dataset = TensorDataset(x, y)
    sampler = DistributedSampler(dataset)
    dataloader = DataLoader(dataset, sampler=sampler, batch_size=batch_size)

    # 4. Define the training loop.
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    loss_fn = nn.MSELoss()

    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)
        model.train()
        total_loss = 0.0
        for step, (xb, yb) in enumerate(dataloader):
            xb, yb = xb.to(device), yb.to(device)
            optimizer.zero_grad()
            loss = loss_fn(model(xb), yb)
            # Log loss to MLflow metric
            mlflow.log_metric("loss", loss.item(), step=step)

            loss.backward()
            optimizer.step()
            total_loss += loss.item() * xb.size(0)

        mlflow.log_metric("total_loss", total_loss)
        print(f"Total loss for epoch {epoch}: {total_loss}")

    cleanup()

Führen Sie die verteilte Schulung aus, indem Sie die verteilte Funktion mit benutzerdefinierten Argumenten aufrufen.
```
run_train.distributed(num_epochs=3, batch_size=1)
```
Bei Ausführung wird in der Notizbuchzellenausgabe eine MLflow-Ausführungsverknüpfung generiert. Klicken Sie auf den MLflow-Ausführungslink, oder suchen Sie ihn im Experimentbereich , um die Ausführungsergebnisse anzuzeigen.

Details zur verteilten Ausführung

Serverlose GPU-API besteht aus mehreren wichtigen Komponenten:

Compute-Manager: Behandelt ressourcenzuordnung und -verwaltung
Laufzeitumgebung: Verwaltet Python-Umgebungen und Abhängigkeiten
Launcher: Koordiniert die Ausführung und Überwachung von Aufträgen

Beim Ausführen im verteilten Modus:

Die Funktion wird serialisiert und über die angegebene Anzahl von GPUs verteilt.
Jede GPU führt eine Kopie der Funktion mit denselben Parametern aus.
Die Umgebung wird über alle GPUs synchronisiert.
Ergebnisse werden gesammelt und von allen GPUs zurückgegeben.

Die API unterstützt beliebte parallele Schulungsbibliotheken wie Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP), DeepSpeed.

Mithilfe der verschiedenen Bibliotheken in Notizbuchbeispielen finden Sie weitere real verteilte Schulungsszenarien.

Häufig gestellte Fragen

Wo sollte der Datenladecode platziert werden?

Wenn Sie die Serverless GPU-API für verteiltes Training verwenden, verschieben Sie den Datenladecode innerhalb des Dekorators @distributed. Die Datensatzgröße kann die maximal zulässige Größe des Pickle-Moduls überschreiten. Daher wird empfohlen, den Datensatz innerhalb des Dekorators zu erzeugen, wie im Folgenden gezeigt:

from serverless_gpu import distributed

# this may cause pickle error
dataset = get_dataset(file_path)
@distributed(gpus=8, gpu_type='H100')
def run_train():
  # good practice
  dataset = get_dataset(file_path)
  ....

Weitere Informationen

Die API-Referenz finden Sie in der Dokumentation zur Serverless GPU Python API .

Feedback

War diese Seite hilfreich?

Last updated on 2026-03-21