Multi-GPU- und Multiknoten-Workload

Sie können verteilte Workload über mehrere GPUs hinweg starten – entweder innerhalb eines einzelnen Knotens oder über mehrere Knoten – mit der Serverless GPU Python-API. 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 Ausführung über Remote-GPUs aus demselben Notizbuch wechseln.

Schnellstart

Die serverlose GPU-API für verteiltes Training ist in serverlosen GPU-Computingumgebungen für Databricks-Notizbücher vorinstalliert. 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.

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, the GPU type,
# and whether or not the GPUs are remote
@distributed(gpus=8, gpu_type='A10', remote=True)
def run_train():
    ...

Im Folgenden finden Sie ein vollständiges Beispiel, das ein MLP-Modell (Multilayer Perceptron) auf 8 A10 GPU-Knoten mithilfe eines Notebooks 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='A10', remote=True)
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 node 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 Knoten synchronisiert.
Ergebnisse werden gesammelt und von allen GPUs zurückgegeben.

Wenn remote auf True eingestellt ist, wird die Workload auf den Remote-GPUs verteilt. Wenn remote auf False gesetzt ist, wird die Arbeitslast auf dem einzelnen GPU-Knoten ausgeführt, der mit dem aktuellen Notizbuch verbunden ist. Wenn der Knoten über mehrere GPU-Chips verfügt, werden alle verwendet.

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

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

Starten mit Ray

Die serverlose GPU-API unterstützt auch das Starten von verteiltem Training mit Ray mithilfe des @ray_launch-Dekorators, der auf @distributed aufgesetzt ist. Jede ray_launch Aufgabe bootet zuerst ein torch-distributed-Rendezvous, um den Ray-Head-Worker zu bestimmen und IPs zu sammeln. Die Rang-Null beginnt ray start --head (mit Metrikexport, falls aktiviert), legt ihre dekorierte Funktion als Ray-Treiber fest RAY_ADDRESSund führt sie aus. Andere Knoten treten über ray start --address bei und warten, bis der Treiber eine Abschlussmarkierung schreibt.

Zusätzliche Konfigurationsdetails:

Verwenden Sie RayMetricsMonitor zum Aktivieren der Ray-Systemmetrikenerfassung auf jedem Knoten mit remote=True.
Definieren Sie Ray-Laufzeitoptionen (Akteure, Datasets, Platzierungsgruppen und Terminplanung) in Ihrer eingerichteten Funktion mithilfe standardmäßiger Ray-APIs.
Verwalten Sie clusterweite Steuerelemente (GPU-Anzahl und Typ, Remote- und lokaler Modus, asynchrones Verhalten und Databricks-Poolumgebungsvariablen) außerhalb der Funktion in den Dekorierargumenten oder Notizbuchumgebungen.

Das folgende Beispiel zeigt, wie Sie Folgendes verwenden @ray_launch:

from serverless_gpu.ray import ray_launch
@ray_launch(gpus=16, remote=True, gpu_type='A10')
def foo():
    import os
    import ray
    print(ray.state.available_resources_per_node())
    return 1
foo.distributed()

Ein vollständiges Beispiel finden Sie in diesem Notizbuch, das Ray startet, um ein Resnet18-Neuronales Netzwerk auf mehreren A10-GPUs zu trainieren.

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, remote=True)
def run_train():
  # good practice
  dataset = get_dataset(file_path)
  ....

Kann ich reservierte GPU-Pools verwenden?

Wenn in Ihrem Arbeitsbereich ein reservierter GPU-Pool verfügbar ist (bitte wenden Sie sich an Ihren Administrator) und Sie remote auf True im @distributed-Dekorator angeben, wird die Arbeitsauslastung standardmäßig im reservierten GPU-Pool gestartet. Wenn Sie den GPU-Pool bei Bedarf verwenden möchten, legen Sie die Umgebungsvariable DATABRICKS_USE_RESERVED_GPU_POOLFalse vor dem Aufrufen der verteilten Funktion wie unten dargestellt fest:

import os
os.environ['DATABRICKS_USE_RESERVED_GPU_POOL'] = 'False'
@distributed(gpus=8, remote=True)
def run_train():
    ...

Learn more

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

Feedback

War diese Seite hilfreich?

Last updated on 2025-12-19