Verteilte Schulung in Notizbüchern

Important

Dieses Feature befindet sich in der Betaversion.

Der @distributed Decorator aus dem Databricks AI Runtime SDK (databricks-sdk-air) 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.

Wichtige Punkte auf dieser Seite:

  • 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.
  • Verbinden Sie Ihr Notebook mit einem 8xH100-Beschleuniger für ein vollständiges Multi-GPU-Training. Standardmäßig verwendet @distributed jede GPU auf dem Knoten.

Note

Diese Seite behandelt verteiltes Training aus Databricks-Notebooks mit dem Databricks AI Runtime SDK. Um verteilte Trainings-Workloads von Ihrem lokalen Rechner zu übermitteln, verwenden Sie die Databricks CLI-Befehle für AI Runtime. Siehe Verwenden der Databricks-CLI mit AI Runtime.

Quickstart

Die Beispiele auf dieser Seite verwenden das Databricks AI Runtime SDK, das in der GPU-Umgebung Version 6 und höher vorinstalliert ist. In Version 5 und älter solltest du stattdessen das serverless_gpu Paket verwenden. Siehe Migration von serverless_gpu. Um ein Notebook an einen Beschleuniger anzuschließen, siehe Schritt 1: Ein Notebook an eine GPU anschließen.

Dekoriere deine Trainingsveranstaltung mit @distributed, und nenne sie dann mit .distributed():

from databricks.air import distributed

# By default, @distributed uses every GPU on the node your notebook is connected
# to. To use fewer, pass num_accelerators (for example, num_accelerators=2).
@distributed()
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()

Für verteiltes Training ist ein 8xH100-Beschleuniger erforderlich, der einen einzelnen Knoten mit 8 GPUs bereitstellt. Standardmäßig verwendet @distributed alle acht GPUs. Stellen Sie num_accelerators so ein, dass weniger verwendet wird. Der Parameter accelerator_type ist ebenfalls optional und wird automatisch vom Beschleuniger, mit dem Ihr Notebook verbunden ist, erkannt.

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, in denen die einzelnen Bibliotheken verwendet werden, siehe verteilte Trainingsbeispiele.

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 Fan-Out: Die Funktion wird serialisiert und auf jeder ausgewählten GPU 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.

Von serverless_gpu migrieren

Das serverless_gpu Paket ist veraltet und erhält keine neuen Funktionen mehr. Ab der GPU-Umgebung Version 6 ersetzt das Databricks AI Runtime SDK (databricks-sdk-air) diese und stellt entsprechende APIs mit umbenannten Argumenten unter dem databricks.air Namensraum bereit. Das verwendete Paket hängt von der GPU-Umgebungsversion ab:

  • Version 5 oder niedriger: Nur serverless_gpu einfügen. Verwenden Sie serverless_gpu weiterhin auf diesen Versionen.
  • Version 6: Enthält beide Pakete, aber nur das AI Runtime SDK erhält Updates. Version 6 ist die letzte GPU-Umgebungsversion, die serverless_gpu enthält.

Das AI Runtime SDK ist ein separates Paket vom Databricks SDK für Python (databricks-sdk).

  serverless_gpu (veraltet) AI Runtime SDK (databricks.air)
GPU-Umgebungsversionen 6 und darunter 6 und höher
Warenimport from serverless_gpu import distributed, ray_init from databricks.air import distributed, ray_init
Argument zur GPU-Anzahl gpus num_accelerators
GPU-Typ-Argument gpu_type (GPUType) accelerator_type (AcceleratorType)
GPU-Typen "A10", "H100", "B300" "GPU_1xA10", "GPU_1xH100", "GPU_8xH100", "GPU_8xB300"
Modul für das Laden von Daten und Checkpoints serverless_gpu.data databricks.air.data. Siehe Laden von Daten auf AI-Runtime.

Um ein Notebook zu migrieren:

  1. Verbinden Sie das Notebook mit der GPU-Umgebung Version 6 oder höher, die das AI Runtime SDK enthält.
  2. Importe von serverless_gpu zu databricks.air ändern.
  3. Umbennen gpus zu num_accelerators und gpu_type zu accelerator_type, und verwenden Sie die neuen Beschleuniger-Typenwerte. Die alten Argumentnamen und GPU-Typwerte funktionieren weiterhin in databricks.air, geben aber eine FutureWarning aus. Der Support wird in GPU-Umgebung Version 7 eingestellt.

Tip

Bitten Sie Genie Code, dies für Sie zu erledigen:

Migrate this notebook from the serverless_gpu SDK to the AI Runtime SDK.

Für Databricks-CLI-Workloads, die databricks.air.data verwenden, fügen Sie databricks-sdk-air[data] zum environment.dependencies des Workloads hinzu.

Stammt von TorchDistributor

Wenn Sie verteiltes Training mit TorchDistributor auf Spark-Clustern ausführen und Ihre Arbeitslast auf einen einzigen Knoten passt, ist die @distributed API der empfohlene Ersatz für KI-Runtime. Dein DDP-, FSDP- oder DeepSpeed-Modellcode ändert sich nicht.

Feature @distributed-API TorchDistributor
Infrastructure Vollständig serverlos, keine Clusterverwaltung Erfordert einen Spark-Cluster mit GPU-Workern
Setup 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. @distributed verwendet standardmäßig alle acht GPUs.
  • 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.

Der Decorator @distributed befindet sich in Beta und läuft auf einem einzigen Knoten (siehe Einschränkungen), daher ersetzt er nicht jede TorchDistributor-Arbeitslast. Behalte Workloads, die von der Spark-Integration abhängen, auf TorchDistributor. Um verteiltes Training von Ihrem lokalen Rechner oder über mehrere Knoten hinweg durchzuführen, verwenden Sie die Databricks CLI-Befehle für AI Runtime. Siehe Verwenden der Databricks-CLI mit AI Runtime.

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. Importiere den distributed Decorator aus dem AI Runtime SDK.

    from databricks.air 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()
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-parallel 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 der Ausführung wird in der Notebookzellenausgabe ein Link zur MLflow-Ausführung 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 databricks.air import distributed

# Don't do this: loading the dataset outside the function means it's captured
# and pickled with the function, which can raise a pickle error.
# dataset = get_dataset(file_path)

@distributed()
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 databricks.air.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 AI Runtime laden und Checkpoint mit verteiltem Checkpoint (DCP).

Limitations

  • @distributed unterstützt nur einen einzelnen Knoten: es läuft über die GPUs des Knotens, mit dem dein Notebook verbunden ist. Um über mehrere Knoten hinweg zu trainieren, verwenden Sie die AI Runtime Databricks CLI. Siehe Verwenden der Databricks-CLI mit AI Runtime.
  • Wenn Sie accelerator_type festlegen, muss es mit dem Beschleuniger übereinstimmen, mit dem Ihr Notebook verbunden ist (zum Beispiel "GPU_8xH100"). 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 unterstützt 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 Ressourcen

  • Zuletzt aktualisiert um