Verteilte Schulung in Notizbüchern

Von Bedeutung

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

Der @distributed Dekorateur der Serverless GPU Python API ist die bequemste Möglichkeit, verteiltes Training aus einem Databricks-Notebook auszuführen. Wenn man seine Trainingsfunktion so gestaltet, nennt man es, und AI Runtime läuft sie über alle GPUs auf dem Knoten, mit dem Ihr Notebook verbunden ist. Der gleiche Code skaliert von einer einzelnen GPU zu mehreren GPUs, ohne Cluster zur Bereitstellung und ohne verteilten Launcher zum Konfigurieren.

Tip

  • Der @distributed Decorator führt eine Trainingsfunktion über jede GPU deines Knotens aus dem Inneren eines Notebooks aus.
  • Es unterstützt PyTorch DDP, FSDP und DeepSpeed und überträgt mit minimalen Änderungen den Single-GPU-Code auf Multi-GPUs.
  • Schließen Sie Ihr Notebook an einen 8xH100-Beschleuniger an und stellen Sie ein vollständiges Multi-GPU-Training ein.gpus=8

Quickstart

Das serverless_gpu Paket ist vorinstalliert, wenn dein Notebook an eine serverlose GPU angeschlossen ist. Dekoriere deine Trainingsveranstaltung mit @distributed, und nenne sie dann mit .distributed():

from serverless_gpu import distributed

# gpus is the number of GPUs on the node. gpu_type is optional and
# auto-detected from the accelerator your notebook is connected to.
@distributed(gpus=8, gpu_type="H100")
def train():
    import os
    import torch
    import torch.distributed as dist

    # Bind this process to its own GPU before training.
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")
    dist.init_process_group("nccl")
    # ... build the model and data on `device`, then run your training loop ...
    dist.destroy_process_group()

train.distributed()

Jeder .distributed() Aufruf erstellt einen MLflow-Run (oder einen verschachtelten Kind-Run, falls einer bereits aktiv ist) und druckt einen Run-Link in der Zellausgabe aus. Für eine vollständige, laufbare Anleitung siehe vollständiges Beispiel.

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: Microsoft Optimierungsbibliothek für umfangreiche Modellschulungen.

Für reale Trainingsszenarien, die jede Bibliothek verwenden, siehe Notizbuchbeispiele.

Wie der @distributed Dekorateur arbeitet

Wenn Sie eine dekorierte Funktion mit .distributed()aufrufen, übernimmt AI Runtime die Mechaniken, die Sie sonst manuell mit einem verteilten Launcher konfigurieren würden:

  • Serialisierung und Verteilung: Die Funktion wird serialisiert und für jede der angeforderten gpus gestartet. Jede GPU führt eine Kopie der Funktion mit denselben Argumenten aus.
  • Umgebungssynchronisation: Die Python-Umgebung und die Abhängigkeiten werden über alle Ränge hinweg repliziert, sodass jeder Prozess denselben Code ausführt.
  • Umgebungsvariablen für den Rang: Standardvariablen wie LOCAL_RANK werden für jeden Prozess gesetzt. Lesen Sie sie in Ihrer Funktion, um das Modell und die Daten auf dem richtigen Gerät zu platzieren.
  • Ergebnissammlung: Rückgabewerte werden aus allen Rängen gesammelt und an den Anrufer zurückgegeben.
  • MLflow-Tracking: Jeder .distributed() Aufruf erstellt einen MLflow-Run oder einen verschachtelten Kind-Run, falls einer bereits aktiv ist, sodass die von Ihrer Funktion protokollierten Metriken auf demselben Run landen.
  • Lebenszyklus und Timeout: Die verteilte Ausführung läuft innerhalb des Lebenszyklus des Notizbuchs. Das Beenden des Notebooks beendet die Ausführung. Der Dekorateur hat standardmäßig eine Timeout von 3 Stunden. Geben Sie timeout in Sekunden an, um den Wert zu ändern, oder timeout=None, um ihn zu deaktivieren. Benutzerdefinierte Timeouts erfordern GPU-Umgebung v5 und höher.

Die API baut auf den Standard-PyTorch-Bibliotheken auf: Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP) und DeepSpeed.

Stammt von TorchDistributor

Wenn Sie heute verteiltes PyTorch auf Spark mit TorchDistributor ausführen und Ihre Arbeitslast auf einen einzigen Knoten passt, ist die serverless_gpu@distributed API der empfohlene Ersatz für neue Deep-Learning-Workloads. Es entfernt den Spark-Cluster und gibt dir denselben Codepfad von einer einzelnen GPU zu mehreren GPUs.

Funktion serverless_gpu @distributed API TorchDistributor
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 Im Decorator werden Unity-Katalog-Volumes verwendet (UCVolumeDataset zum Streaming von Dateidaten). Über Spark oder das Dateisystem

Um eine Workload mit nur einem Knoten zu migrieren:

  • Ersetze den TorchDistributor(...).run(train_fn, ...) Aufruf durch den @distributed Decorator auf train_fn, und starte dann mit train_fn.distributed(...).
  • Entferne den Spark-Cluster und die GPU-Worker-Konfiguration. Verbinde dein Notebook stattdessen mit einem 8xH100-Beschleuniger und setze gpus=8.
  • Verlagern Sie das Laden von Daten in die dekorierte Funktion. Siehe Datenladen.
  • Behalte deinen bestehenden DDP-, FSDP- oder DeepSpeed-Modellcode. Der Dekorateur unterstützt alle drei.

@distributed läuft auf einem einzelnen Knoten (siehe Einschränkungen), sodass er nicht jede TorchDistributor-Workload ersetzt. Behalte Workloads, die von der Spark-Integration abhängen, auf TorchDistributor. Um verteiltes Training von Ihrem lokalen Rechner oder über mehrere Knoten hinweg auszuführen, verwenden Sie stattdessen die AI Runtime CLI, die in der öffentlichen Vorschau verfügbar ist. Siehe AI Runtime CLI.

Vollständiges Beispiel

Das folgende Beispiel trainiert in einem Notebook ein Multilayer-Perzeptron-Modell (MLP) mithilfe von 8 H100-GPUs.

  1. 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():
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
    dist.init_process_group("nccl")

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)

  2. Importieren Sie die serverless_gpu Bibliothek und das distributed Modul.

    import serverless_gpu
from serverless_gpu import distributed

  3. Schließen Sie den Modellschulungscode in eine Funktion ein und dekorieren Sie die Funktion mit dem @distributed Dekorator. Die dekorierte Funktion ist der Einstiegspunkt für die verteilte Ausführung, daher sollten Sie die gesamte Trainingslogik, das Laden der Daten und die Modellinitialisierung innerhalb dieser Funktion definieren.

    @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()

  4. Führen Sie das verteilte Training aus, indem Sie die verteilte Funktion mit benutzerdefinierten Argumenten aufrufen.

    run_train.distributed(num_epochs=3, batch_size=1)

  5. Bei Ausführung wird eine MLflow-Ausführungsverknüpfung in der Notebookzellenausgabe generiert. Klicken Sie auf den MLflow-Ausführungslink, oder suchen Sie ihn im Experimentbereich , um die Ausführungsergebnisse anzuzeigen. Weitere Informationen zum Anpassen von Experimentnamen, zum Nachverfolgen von Metriken und zum Fortsetzen von Ausführungen finden Sie unter Experimentnachverfolgung und Beobachtbarkeit.

Datenladung

Fügen Sie den Datenladecode in die @distributed Funktion ein. Ein Datensatz kann die von pickle maximal zulässige Größe überschreiten, sodass das Generieren oder Laden innerhalb des Dekorators Serialisierungsfehler vermeidet:

from serverless_gpu import distributed

# This may cause a pickle error because the dataset is captured by the function.
dataset = get_dataset(file_path)

@distributed(gpus=8, gpu_type='H100')
def run_train():
    # Load the dataset inside the decorated function instead.
    dataset = get_dataset(file_path)
    ...

Verwenden Sie für dateibasierte Daten, die in Unity Catalog-Volumes gespeichert sind, UCVolumeDataset aus serverless_gpu.data, das Dateien mit lokalem Caching streamt und sie automatisch über Ranks und Worker partitioniert. Um Prüfpunkte für verteiltes Training in einem Volume zu speichern, verwenden Sie UCVolumeWriter und UCVolumeReader. Siehe Daten in die AI Runtime laden und Modell-Checkpointing.

Limitations

  • Verteiltes Training läuft über die GPUs auf dem einzelnen Knoten, mit dem dein Notebook verbunden ist. Für das vollständige Multi-GPU-Training verbinden Sie sich mit einem 8xH100-Beschleuniger, der einen Knoten mit 8 GPUs bereitstellt, und setzen Sie gpus=8.
  • Der Beschleunigertyp muss übereinstimmen. Wenn du @distributed in gpu_type festlegst, muss es mit dem Accelerator übereinstimmen, mit dem dein Notebook verbunden ist ("H100" oder "A10"). Eine Nichtübereinstimmung führt dazu, dass der Workload fehlschlägt. Der Parameter ist optional und wird automatisch erkannt, wenn er weggelassen wird.
  • AI Runtime empfiehlt GPU-Umgebung v4 und höher. Benutzerdefinierte Timeouts (der Parameter timeout ) erfordern GPU-Umgebung ab Version 5.
  • Der Dekorateur läuft standardmäßig nach 3 Stunden ab. Geben Sie timeout in Sekunden an, um den Wert zu ändern, oder timeout=None, um ihn zu deaktivieren.
  • Die Ausführung läuft innerhalb des Lebenszyklus des Notizbuchs. Das Beenden des Notebooks beendet die Ausführung.

Weitere Informationen

  • Last updated on