Gedistribueerd trainen in notebooks

Belangrijk

Deze functie bevindt zich in de bètaversie.

De @distributed decorator van de Serverless GPU Python API is de meest handige manier om gedistribueerde training uit te voeren vanuit een Databricks-notebook. Decoreer je trainingsfunctie, noem het zo, en AI Runtime draait het op alle GPU's op de node waar je notebook op verbonden is. Dezelfde code schaalt van single-GPU naar multi-GPU zonder cluster om te provisioneren en zonder gedistribueerde launcher om te configureren.

Belangrijke punten op deze pagina:

  • De @distributed decorator voert vanuit een notebook een trainingsfunctie uit op alle GPU's in je node.
  • Het ondersteunt PyTorch DDP, FSDP en DeepSpeed, en verplaatst single-GPU code naar multi-GPU met minimale wijzigingen.
  • Sluit je notebook aan op een 8xH100-versneller en stel gpus=8 in voor volledige multi-GPU-training.

Note

Deze pagina bevat gedistribueerde training van Databricks-notebooks met de serverloze GPU-Python-API. Om gedistribueerde trainingsworkloads vanaf je lokale machine in te dienen, gebruik je de Databricks CLI-commando's voor AI Runtime, die in Public Preview staan. Zie Gebruik de Databricks CLI met AI Runtime.

Quickstart

Het serverless_gpu pakket is vooraf geïnstalleerd wanneer je notebook is aangesloten op een serverless GPU. Versier je trainingsfunctie met @distributed, noem het dan met .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()

Gedistribueerde training vereist een 8xH100-accelerator, die één knooppunt met 8 GPU's inricht. Wanneer u de @distributed decorator gebruikt, stelt u het in gpus=8. De gpu_type parameter is optioneel en automatisch gedetecteerd vanuit de accelerator waar uw notebook mee is verbonden.

Elke .distributed()-aanroep maakt een MLflow-run aan (of een geneste onderliggende run als er al een run actief is) en toont een link naar de run in de celuitvoer. Voor een volledige, uitvoerbare walkthrough, zie Volledig voorbeeld.

Ondersteunde frameworks

De @distributed API kan worden geïntegreerd met belangrijke gedistribueerde trainingsbibliotheken:

  • PyTorch Distributed Data Parallel (DDP): Standaard multi-GPU-gegevensparallelisme.
  • Fully Sharded Data Parallel (FSDP): Geheugenefficiënte training voor grote modellen.
  • DeepSpeed: Microsoft optimalisatiebibliotheek voor grote modellen trainen.

Voor echte trainingsscenario's voor elke bibliotheek, zie voorbeelden van gedistribueerde training.

Hoe de @distributed decorateur werkt

Wanneer je een versierde functie aanroept met .distributed(), behandelt AI Runtime de mechanieken die je anders met de hand zou configureren met een gedistribueerde launcher:

  • Serialisatie en uitschalen: De functie wordt geserialiseerd en gestart op elk van de gpus die je aanvraagt. Elke GPU draait een kopie van de functie met dezelfde argumenten.
  • Omgevingssynchronisatie: De Python-omgeving en afhankelijkheden worden over alle rangen gerepliceerd, zodat elk proces dezelfde code draait.
  • Omgevingsvariabelen voor de rang: Standaardvariabelen zoals LOCAL_RANK worden voor elk proces ingesteld. Lees ze in je functie om het model en de gegevens op het juiste apparaat te plaatsen.
  • Resultaatverzameling: Retourwaarden worden verzameld van alle rangen en teruggegeven aan de caller.
  • MLflow-tracering: Elke .distributed() aanroep maakt een MLflow-run aan, of een geneste onderliggende run als er al een run actief is, zodat de metrische gegevens die vanuit je functie worden gelogd in dezelfde run terechtkomen.
  • Levenscyclus en time-out: Gedistribueerde uitvoering wordt uitgevoerd binnen de levenscyclus van het notebook. Het beëindigen van het notitieboekje beëindigt de run. De decorateur heeft standaard een timeout van 3 uur. Geef timeout op in seconden om dit te wijzigen, of timeout=None om het uit te schakelen. Aangepaste time-outs vereisen een GPU-omgeving van v5 of hoger.

De API bouwt voort op de standaard PyTorch-bibliotheken: Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP) en DeepSpeed.

Afkomstig van TorchDistributor

Als je vandaag gedistribueerde PyTorch op Spark draait met TorchDistributor en je workload past op één node, is de serverless_gpu@distributed API de aanbevolen vervanging voor nieuwe deep learning-workloads. Het verwijdert de Spark-cluster en geeft je hetzelfde codepad van single-GPU naar multi-GPU.

Feature serverless_gpu @distributed API Zaklampverdeler
Infrastructuur Volledig serverloos, geen clusterbeheer Vereist een Spark-cluster met GPU-werkrollen
Installatie Eén enkele decorator, minimale configuratie Vereist Spark-cluster en installatie van TorchDistributor
Framework-ondersteuning PyTorch DDP, FSDP, DeepSpeed Voornamelijk PyTorch DDP
Gegevens laden In de decorator worden Unity Catalog-volumes gebruikt (UCVolumeDataset voor het streamen van bestandsgegevens). Met Spark of via het bestandssysteem

Om een single-node-workload te migreren:

  • Vervang de TorchDistributor(...).run(train_fn, ...) aanroep door de @distributed decorator op train_fn, en start vervolgens met train_fn.distributed(...).
  • Verwijder het Spark-cluster en de GPU-workerconfiguratie. Sluit je notebook aan op een 8xH100-accelerator en stel in plaats daarvan gpus=8 in.
  • Verplaats het laden van gegevens naar de gedecoreerde functie. Zie Gegevens laden.
  • Behoud je bestaande DDP-, FSDP- of DeepSpeed-modelcode. De decorateur ondersteunt alle drie.

@distributed draait op één enkele node (zie Beperkingen), dus het vervangt niet elke TorchDistributor-werklast. Houd workloads die afhankelijk zijn van Spark-integratie op TorchDistributor draaien. Om gedistribueerde training uit te voeren vanaf je lokale machine of over meerdere nodes, gebruik je in plaats daarvan de AI Runtime CLI, die in Public Preview staat. Zie Gebruik de Databricks CLI met AI Runtime.

Volledig voorbeeld

Het volgende voorbeeld traint een multilayer perceptron (MLP) model op 8 H100 GPU's vanaf een notebook.

  1. Stel uw model in en definieer hulpprogrammafuncties.

    
# 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. Importeer de serverless_gpu bibliotheek en de distributed module.

    import serverless_gpu
from serverless_gpu import distributed

  3. Verpakt de modeltrainingscode in een functie en verfraaid de functie met de @distributed decorator. De gedecoreerde functie is het instappunt voor gedistribueerde uitvoering, dus definieer alle trainingslogica, het laden van gegevens en de modelinitialisatie daarbinnen.

    @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. Voer de gedistribueerde training uit door de gedistribueerde functie aan te roepen met door de gebruiker gedefinieerde argumenten.

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

  5. Wanneer dit wordt uitgevoerd, wordt in de uitvoer van de notebookcel een link naar een MLflow-run gegenereerd. Klik op de koppeling voor de MLflow-uitvoering of zoek deze in het deelvenster Experiment om de uitvoeringsresultaten te bekijken. Zie Experimenttracking en waarneembaarheid voor meer informatie over het aanpassen van experimentnamen, het bijhouden van metrische gegevens en het hervatten van uitvoeringen.

Data-laden

Plaats data-laadcode in de @distributed functie. Een dataset kan de maximale grootte overschrijden die door pickle is toegestaan, dus door deze binnen de decorator te genereren of te laden, worden serialisatiefouten voorkomen:

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)
    ...

Voor bestandsgebaseerde gegevens die zijn opgeslagen in Unity Catalog-volumes, gebruikt u UCVolumeDataset van databricks.air.data, waarmee bestanden worden gestreamd met lokale caching en automatisch over ranks en workers worden verdeeld. Als u een controlepunt van gedistribueerde training op een volume wilt opslaan, gebruikt u UCVolumeWriter en UCVolumeReader. Zie Gegevens laden op AI Runtime en Checkpoint met Distributed Checkpoint (DCP).

Limitations

  • Gedistribueerde training draait over de GPU's op de enkele node waar je notebook mee verbonden is. Voor volledige multi-GPU training, sluit aan op een 8xH100 accelerator, die één node met 8 GPU's inricht, en stel gpus=8.
  • Het type versneller moet overeenkomen. Als je gpu_type instelt in "H100", moet dit overeenkomen met de accelerator waarmee je notebook is verbonden ("A10" of @distributed). Een mismatch zorgt ervoor dat de werklast faalt. De parameter is optioneel en wordt automatisch gedetecteerd wanneer deze wordt weggelaten.
  • AI Runtime raadt GPU-omgeving v4 en hoger aan. Aangepaste timeouts (de timeout parameter) vereisen GPU-omgeving v5 en hoger.
  • De decorator heeft standaard een time-out van 3 uur. Geef timeout op in seconden om dit te wijzigen, of timeout=None om het uit te schakelen.
  • De uitvoering verloopt binnen de levenscyclus van het notitieboek. Het beëindigen van het notitieboekje beëindigt de run.

Aanvullende bronnen

  • Laatst bijgewerkt om