Gedistribueerd trainen in notebooks

De @distributed decorator van de Databricks AI Runtime SDK (databricks-sdk-air) is de meest handige manier om gedistribueerde training uit te voeren vanaf 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 accelerator voor volledige multi-GPU training. Standaard gebruikt @distributed elke GPU op de node.

Quickstart

De voorbeelden op deze pagina gebruiken de Databricks AI Runtime SDK, die vooraf is geïnstalleerd in GPU-omgeving versie 6 en hoger. In versie 5 en lager gebruik je in plaats daarvan het serverless_gpu pakket. Zie Migreren van serverless_gpu. Om een notebook aan een accelerator te koppelen, zie Verbinden met een notebook.

Versier je trainingsfunctie met @distributed, noem het dan met .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()

Gedistribueerde training vereist een 8xH100-accelerator, die één knooppunt met 8 GPU's inricht. Standaard gebruikt @distributed alle acht GPU's. Stel num_accelerators in om minder te gebruiken. De parameter accelerator_type is ook optioneel en wordt automatisch gedetecteerd vanuit de accelerator waarop je notebook is aangesloten.

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 fan-out: De functie wordt geserialiseerd en gestart op elke geselecteerde GPU. 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.

Migreren van serverless_gpu

Het serverless_gpu pakket is verouderd en krijgt geen nieuwe functies meer. Vanaf GPU-omgeving versie 6 vervangt de Databricks AI Runtime SDK (databricks-sdk-air) deze en biedt equivalente API's, met hernoemde argumenten, onder de databricks.air naamruimte. Het pakket dat je gebruikt hangt af van de GPU-omgevingsversie:

• Versie 5 en lager: Neem alleen serverless_gpu op. Blijf serverless_gpu gebruiken op deze versies.
• Versie 6: Bevat beide pakketten, maar alleen de AI Runtime SDK ontvangt updates. Versie 6 is de laatste GPU-omgevingsversie die serverless_gpu bevat.

De AI Runtime SDK is een apart pakket van de Databricks SDK voor Python (databricks-sdk).

Om een notebook te migreren:

1. Verbind het notebook met GPU-omgeving versie 6 of hoger, die de AI Runtime SDK bevat.
2. Imports wijzigen van serverless_gpu naar databricks.air.
3. Hernoem gpus naar num_accelerators en gpu_type naar accelerator_type, en gebruik de nieuwe versneller-type-waarden. De oude argumentnamen en GPU-typewaarden werken nog steeds in databricks.air , maar geven een FutureWarning. De ondersteuning voor GPU-omgeving versie 7 wordt beëindigd.

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

Voor Databricks CLI-workloads die gebruikmaken van databricks.air.data, voeg databricks-sdk-air[data] toe aan de workload's environment.dependencies.

Afkomstig van TorchDistributor

Als je gedistribueerde training uitvoert met TorchDistributor op Spark-clusters en je workload past op één node, is de @distributed API de aanbevolen vervanger op AI Runtime. Je DDP-, FSDP- of DeepSpeed-modelcode verandert niet.

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 in plaats daarvan aan op een 8xH100 accelerator. @distributed gebruikt standaard alle acht GPU's.
• 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.

De @distributed decorator staat in Public Preview en draait op één node (zie Beperkingen), dus hij 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 de Databricks CLI-commando's voor AI Runtime. 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 distributed decorator uit de AI Runtime SDK.

   from databricks.air 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()
   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. 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 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)
    ...

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

• @distributed ondersteunt slechts één node: het wordt uitgevoerd op alle GPU's van de node waar je notebook op is aangesloten. Om over meerdere knooppunten te trainen, gebruik je de AI Runtime Databricks CLI. Zie Gebruik de Databricks CLI met AI Runtime.
• Als je accelerator_type instelt, moet het overeenkomen met de accelerator waarop je notebook is aangesloten (bijvoorbeeld "GPU_8xH100"). Een mismatch zorgt ervoor dat de werklast faalt. De parameter is optioneel en wordt automatisch gedetecteerd wanneer deze wordt weggelaten.
• AI Runtime ondersteunt GPU-omgeving v4 en hoger. 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

• Voor de @distributed decorator, AcceleratorType, en Ray API's, zie de referentie van de Databricks AI Runtime SDK API.
• Voor patronen die je trainingspijplijn efficiënter en veerkrachtiger maken, zie de handleiding voor prestaties en veerkracht.
• Voor end-to-end trainingsscenario's, zie voorbeelden van gedistribueerde trainingen.