Multi-GPU-arbetsbelastning

Viktigt!

Den här funktionen finns i Beta. Arbetsyteadministratörer kan styra åtkomsten till den här funktionen från sidan Förhandsversioner . Se Hantera förhandsversioner av Azure Databricks.

Du kan starta distribuerade arbetsbelastningar över flera GPU:er på en enda nod med hjälp av Serverlös GPU Python API. API:et tillhandahåller ett enkelt, enhetligt gränssnitt som sammanfattar informationen om GPU-etablering, miljökonfiguration och arbetsbelastningsdistribution. Med minimala kodändringar kan du sömlöst gå från enkel GPU-träning till distribuerad körning med flera GPU:er från samma notebook-fil.

Ramverk som stöds

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

PyTorch Distributed Data Parallel (DDP) – Standard för dataparallellism med flera GPU:er.
Fullständigt shardat dataparallellt (FSDP) – Minneseffektiv träning för stora modeller.
DeepSpeed – Microsofts optimeringsbibliotek för stor modellträning.

serverless_gpu API jämfört med TorchDistributor

I följande tabell jämförs API:et serverless_gpu@distributed med TorchDistributor:

Feature	`serverless_gpu` `@distributed` API	Lyktdistributör
Infrastruktur	Fullständigt serverlös, ingen klusterhantering	Kräver ett Spark-kluster med GPU-arbetare
Inställningar	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	Inredare, använder Unity-katalogvolymer	Via Spark eller filsystem

serverless_gpu API:et är den rekommenderade metoden för nya djupinlärningsarbetsbelastningar i Databricks. TorchDistributor är fortfarande tillgängligt för arbetsbelastningar nära kopplade till Spark-kluster.

Snabbstart

Det serverlösa GPU-API:et för distribuerad träning är förinstallerat när du är ansluten till en serverlös GPU i Databricks notebook-filer och jobb. Vi rekommenderar GPU-miljö 4 och senare. Om du vill använda den för distribuerad träning ska du importera och använda dekoratören distributed för att distribuera träningsfunktionen.

Lägg in modellträningskoden i en funktion och använd dekoratören @distributed på funktionen. Den dekorerade funktionen blir startpunkt för distribuerad körning – all träningslogik, datainläsning och modellinitialisering bör definieras inom denna funktion.

Varning

Parametern gpu_type i @distributed måste matcha den acceleratortyp som din notebook är ansluten till. Till exempel @distributed(gpus=8, gpu_type='H100') kräver att din dator är ansluten med en H100-accelerator. Om du använder en felmatchad acceleratortyp (till exempel att ansluta till A10 när du anger H100) misslyckas arbetsbelastningen.

Kodfragmentet nedan visar den grundläggande användningen av @distributed:

# Import the distributed decorator
from serverless_gpu import distributed

# Decorate your training function with @distributed and specify the number of GPUs and GPU type
@distributed(gpus=8, gpu_type='H100')
def run_train():
    ...

Nedan visas ett fullständigt exempel som tränar en MLP-modell på 8 H100 GPU:er från en notebook.

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

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)

Importera serverless_gpu-biblioteket och den distribuerade modulen.
```
import serverless_gpu
from serverless_gpu import distributed
```

Lägg in modellträningskoden i en funktion och använd dekoratören @distributed på funktionen.

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

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)
```
När en MLflow-körning utförs genereras en länk i notebookcellens utdata. Klicka på MLflow-körningslänken eller hitta den i panelen Experiment för att se körningsresultatet.

Detaljer om distribuerad körning

Serverlöst GPU-API består av flera viktiga komponenter:

Compute Manager: Hanterar resursallokering och hantering
Körningsmiljö: Hanterar Python-miljöer och beroenden
Launcher: Samordnar uppdragskörning och övervakning

När du kör i distribuerat läge:

Funktionen serialiseras och distribueras över det angivna antalet GPU:er
Varje GPU kör en kopia av funktionen med samma parametrar
Miljön synkroniseras över alla GPU:er
Resultat samlas in och returneras från alla GPU:er

API:et stöder populära parallella träningsbibliotek som DDP (Distributed Data Parallel ), Fullständigt fragmenterad dataparallell (FSDP), DeepSpeed.

Du hittar fler verkliga distribuerade träningsscenarier med hjälp av de olika biblioteken i notebook-exempel.

Vanliga frågor och svar

Var ska datainläsningskoden placeras?

När du använder det serverlösa GPU-API: et för distribuerad träning flyttar du datainläsningskoden i @distributed dekoratören. Datamängdens storlek kan överskrida den maximala storlek som tillåts av pickle, så vi rekommenderar att du genererar datamängden inuti dekoratören enligt nedan:

from serverless_gpu import distributed

# this may cause pickle error
dataset = get_dataset(file_path)
@distributed(gpus=8, gpu_type='H100')
def run_train():
  # good practice
  dataset = get_dataset(file_path)
  ....

Lära sig mer

Api-referensen finns i dokumentationen för Serverlös GPU Python API .

Feedback

Var den här sidan till hjälp?

Last updated on 2026-03-21