Distribuerad träning i anteckningsböcker

Important

Den här funktionen finns i Beta.

Decoratorn @distributed från Databricks AI Runtime SDK (databricks-sdk-air) är det mest bekväma sättet att köra distribuerad träning från en Databricks-notebook. Dekorera din träningsfunktion, anropa den, och AI Runtime kör den på alla GPU:er på den nod som din notebook är ansluten till. Samma kod skalar från en GPU till flera GPU:er utan kluster att provisionera och utan distribuerad launcher att konfigurera.

Viktiga punkter på denna sida:

• Dekoratören @distributed kör en träningsfunktion på alla GPU:er på din nod från en anteckningsboksmiljö.
• Den stöder PyTorch DDP, FSDP och DeepSpeed, och flyttar kod från ett enda GPU till flera GPU med minimala ändringar.
• Koppla din notebook till en 8xH100-accelerator för full multi-GPU-träning. Som standard använder @distributed alla GPU:er på noden.

Note

Denna sida täcker distribuerad träning från Databricks-notebooks med Databricks AI Runtime SDK. För att skicka distribuerade träningsarbetsbelastningar från din lokala maskin, använd Databricks CLI-kommandon för AI Runtime. Se Använd Databricks CLI med AI Runtime.

Quickstart

Exemplen på denna sida använder Databricks AI Runtime SDK, som är förinstallerad i GPU-miljö version 6 och uppåt. På version 5 och lägre, använd paketet serverless_gpu istället. Se migrering från serverless_gpu. För att koppla en notebook till en accelerator, se Steg 1: Koppla en notebook till en GPU.

Dekorera din träningsfunktion med @distributed, och kalla den sedan med .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()

Distribuerad träning kräver en 8xH100-accelerator som etablerar en enskild nod med 8 GPU:er. Som standard använder @distributed alla åtta grafikkort. Ställ in num_accelerators för att använda färre. Parametern accelerator_type är också valfri och upptäcks automatiskt från acceleratorn som din anteckningsbok är ansluten till.

Varje anrop av .distributed() skapar en MLflow-körning (eller en kapslad underordnad körning om det redan finns en aktiv körning) och visar en länk till körningen i cellens utdata. För en komplett, körbar genomgång, se Fullständigt exempel.

Ramverk som stöds

API:et @distributed integreras med större distribuerade utbildningsbibliotek:

• PyTorch Distributed Data Parallel (DDP): Standard multi-GPU data parallellism.
• Fullständigt fragmenterad dataparallell (FSDP): Minneseffektiv träning för stora modeller.
• DeepSpeed: Microsoft optimeringsbibliotek för stor modellträning.

För verkliga träningsscenarier som använder varje bibliotek, se exempel på distribuerad träning.

Hur dekoratören @distributed arbetar

När du anropar en dekorerad funktion med .distributed(), hanterar AI Runtime det praktiska som du annars skulle behöva konfigurera manuellt med ett distribuerat startprogram:

• Sekventiell körning och fan-out: Funktionen körs sekventiellt och startas på varje valt grafikkort. Varje GPU kör en kopia av funktionen med samma argument.
• Miljösynkronisering: Python-miljön och beroendena replikeras över alla nivåer, så varje process kör samma kod.
• Miljövariabler för rank: Standardvariabler som LOCAL_RANK fylls i för varje process. Läs dem i din funktion för att placera modellen och data på rätt enhet.
• Resultatinsamling: Returvärden samlas in från alla nivåer och återlämnas till den som ringer.
• MLflow-spårning: Varje .distributed() anrop skapar en MLflow-körning, eller en nästlad barnkörning om en redan är aktiv, så att mätvärden som loggas från din funktion hamnar på samma körning.
• Livscykel och timeout: Distribuerad exekvering körs inom anteckningsbokens livscykel. Om du avslutar anteckningsboken avslutas körningen. Dekoratorn har en standardtidsgräns på 3 timmar. Passera timeout på några sekunder för att ändra det, eller timeout=None inaktivera det. Anpassade tidsgränser kräver GPU-miljö v5 eller senare.

API:et bygger vidare på de standardiserade PyTorch-biblioteken: Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP) och DeepSpeed.

Migrera från serverless_gpu

Paketet serverless_gpu är föråldrat och får inte längre nya funktioner. Från och med GPU-miljö version 6 ersätter Databricks AI Runtime SDK (databricks-sdk-air) den och tillhandahåller motsvarande API:er, med omdöpta argument, under databricks.air namnrymden. Paketet du använder beror på GPU-miljöns version:

• Version 5 och lägre: Inkludera endast serverless_gpu. Fortsätt använda serverless_gpu på dessa versioner.
• Version 6: Inkluderar båda paketen, men endast AI Runtime SDK får uppdateringar. Version 6 är den sista GPU-miljöversionen som inkluderar serverless_gpu.

AI Runtime SDK är ett separat paket från Databricks SDK för Python (databricks-sdk).

	serverless_gpu (inaktuell)	AI Runtime SDK (databricks.air)
GPU-miljöversioner	6 eller lägre	6 och uppåt
Import	from serverless_gpu import distributed, ray_init	from databricks.air import distributed, ray_init
Argument för antal GPU:er	gpus	num_accelerators
argument för GPU-typ	gpu_type (GPUType)	accelerator_type (AcceleratorType)
GPU-typvärden	"A10", "H100", "B300"	"GPU_1xA10", "GPU_1xH100", "GPU_8xH100", "GPU_8xB300"
Modul för dataladdning och checkpointing	serverless_gpu.data	databricks.air.data. Se Ladda data på AI Runtime.

För att migrera en anteckningsbok:

1. Koppla notebooken till GPU-miljön version 6 eller högre, vilket inkluderar AI Runtime SDK.
2. Ändra import från serverless_gpu till databricks.air.
3. Byt namn på gpus till num_accelerators och gpu_type till accelerator_type, och använd de nya acceleratortypvärdena. De gamla argumentnamnen och GPU-typvärdena fungerar fortfarande i databricks.air men genererar en FutureWarning. Stödet upphör för GPU-miljö version 7.

Tip

Be Genie Code att göra detta åt dig:

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

För Databricks CLI-arbetsbelastningar som använder databricks.air.data, lägg till databricks-sdk-air[data] i arbetsbelastningens environment.dependencies.

Kommer från TorchDistributor

Om du kör distribuerad träning med TorchDistributor på Spark-kluster och din arbetsbelastning får plats på en enda nod, är @distributed-API:t den rekommenderade ersättningen på AI Runtime. Din DDP-, FSDP- eller DeepSpeed-modellkod ändras inte.

Feature	@distributed API (applikationsprogrammeringsgränssnitt)	Lyktdistributör
Infrastructure	Fullständigt serverlös, ingen klusterhantering	Kräver ett Spark-kluster med GPU-arbetare
Konfigurera	Enkel dekoratör, minimal konfiguration	Kräver konfiguration av Spark-kluster och TorchDistributor
Stöd för ramverk	PyTorch DDP, FSDP, DeepSpeed	Främst PyTorch DDP
Datainläsning	Inuti dekoratorn används Unity-katalogvolymer (UCVolumeDataset för strömmande fildata).	Via Spark eller filsystem

För att migrera en arbetsbelastning med en enda nod:

• Byt ut anropet TorchDistributor(...).run(train_fn, ...) mot dekoratören @distributed på train_fn, och starta sedan med train_fn.distributed(...).
• Ta bort Spark-klustret och GPU-worker-konfigurationen. Koppla istället din notebook till en 8xH100-accelerator. @distributed använder alla åtta grafikkort som standard.
• Flytta data som laddas in i den dekorerade funktionen. Se Data loading.
• Behåll din befintliga DDP-, FSDP- eller DeepSpeed-modellkod. Dekoratören har stöd för alla tre.

Dekoratorn @distributed är i Beta och körs på en enda nod (se Begränsningar), så den ersätter inte varje TorchDistributor-arbetsbelastning. Behåll arbetsbelastningar som är beroende av Spark-integration på TorchDistributor. För att köra distribuerad träning från din lokala maskin eller över flera noder, använd Databricks CLI-kommandon för AI Runtime. Se Använd Databricks CLI med AI Runtime.

Fullständigt exempel

Följande exempel tränar en flerlagerperceptronmodell (MLP) på 8 H100-GPU:er i en notebook.

1. Konfigurera din modell och definiera verktygsfunktioner.

   
   # 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. Importera dekoratören distributed från AI Runtime SDK.

   from databricks.air import distributed
   

3. Lägg in modellträningskoden i en funktion och använd dekoratören @distributed på funktionen. Den dekorerade funktionen är ingångspunkten för distribuerad exekvering, så definiera all träningslogik, dataladdning och modellinitiering inuti den.

   @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. Kör den distribuerade träningen genom att anropa den distribuerade funktionen med användardefinierade argument.

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

5. När den körs genereras en MLflow-körningslänk i notebook-cellens utdata. Klicka på MLflow-körningslänken eller hitta den i panelen Experiment för att se körningsresultatet. Mer information om hur du anpassar experimentnamn, spårar mått och återupptar körningar finns i Experimentspårning och observerbarhet.

Dataladdning

Placera dataladdningskod i @distributed funktionen. En datamängd kan överstiga den maximala storlek som tillåts av pickle, så att generera eller ladda in den i dekoratören undviker serialiseringsfel:

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

För filbaserade data som lagras i Unity Catalog-volymer använder du UCVolumeDataset från databricks.air.data, som strömmar filer med lokal cachelagring och partitionerar dem mellan led och arbetare automatiskt. Om du vill spara kontrollpunkter för distribuerad träning till en volym använder du UCVolumeWriter och UCVolumeReader. Se Ladda data på AI Runtime och Checkpoint med Distributed Checkpoint (DCP).

Limitations

• @distributed stöder bara en enda nod: den körs över GPU:erna på noden som din notebook är ansluten till. För att träna över flera noder, använd AI Runtime Databricks CLI. Se Använd Databricks CLI med AI Runtime.
• Om du sätter accelerator_type, måste den matcha acceleratorn som din Notebook är ansluten till (till exempel "GPU_8xH100"). En konflikt leder till att arbetslasten misslyckas. Parametern är valfri och upptäcks automatiskt när den utelämnas.
• AI Runtime stödjer GPU-miljö v4 och uppåt. Anpassade tidsgränser (parametern timeout) kräver GPU-miljö v5 eller senare.
• Dekoratören går ut efter 3 timmar som standard. Passera timeout på några sekunder för att ändra det, eller timeout=None inaktivera det.
• Exekveringen sker inom anteckningsbokens livscykel. Om du avslutar anteckningsboken avslutas körningen.

Ytterligare resurser

• För @distributed decorator-, AcceleratorType, och Ray-API:erna, se Databricks AI Runtime SDK API-referensen.
• För mönster som gör din träningspipeline mer effektiv och motståndskraftig, se guiden för prestation och resiliens.
• För end-to-end-träningsscenarier, se exempel på distribuerad träning.