Delen via

Workload met meerdere GPU's

Belangrijk

Deze functie bevindt zich in de bètaversie. Werkruimtebeheerders kunnen de toegang tot deze functie beheren vanaf de pagina Previews . Zie Azure Databricks-previews beheren.

U kunt gedistribueerde workloads op meerdere GPU's op één knooppunt starten met behulp van de Serverloze GPU Python-API. De API biedt een eenvoudige, geïntegreerde interface waarmee de details van GPU-inrichting, omgevingsinstallatie en workloaddistributie worden weggenomen. Met minimale codewijzigingen kunt u naadloos overstappen van training met één GPU naar gedistribueerde uitvoering met meerdere GPU's vanuit hetzelfde notebook.

Ondersteunde frameworks

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

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

serverless_gpu API versus TorchDistributor

De volgende tabel vergelijkt de serverless_gpu@distributed API met TorchDistributor:

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 Binnen decorator gebruikt u Unity Catalog Volumes Met Spark of via het bestandssysteem

De serverless_gpu API is de aanbevolen benadering voor nieuwe Deep Learning-workloads in Databricks. TorchDistributor blijft beschikbaar voor workloads die nauw zijn gekoppeld aan Spark-clusters.

Snel starten

De serverloze GPU-API voor gedistribueerde training is vooraf geïnstalleerd wanneer u bent verbonden met een serverloze GPU in Databricks-notebooks en -taken. We raden GPU-omgeving 4 en hoger aan. Als u deze wilt gebruiken voor gedistribueerde training, importeert en gebruikt u de distributed decorator om uw trainingsfunctie te distribueren.

Verpakt de modeltrainingscode in een functie en verfraaid de functie met de @distributed decorator. De gedecoreerde functie wordt het invoerpunt voor gedistribueerde uitvoering. Alle trainingslogica, het laden van gegevens en modelinitialisatie moeten binnen deze functie worden gedefinieerd.

Waarschuwing

De gpu_type parameter in @distributed moet overeenkomen met het type accelerator waaraan uw notebook is gekoppeld. Vereist bijvoorbeeld @distributed(gpus=8, gpu_type='H100') dat uw notebook is verbonden met een H100-accelerator. Als u een niet-overeenkomend acceleratortype gebruikt (zoals verbinding maken met A10 tijdens het opgeven van H100), mislukt de werkbelasting.

In het onderstaande codefragment ziet u het basisgebruik van @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():
    ...

Hieronder ziet u een volledig voorbeeld dat een MLP-model (multilayer perceptron) traint op 8 H100 GPU's van 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():
    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)

  2. Importeer de serverless_gpu-bibliotheek en de gedistribueerde module.

    import serverless_gpu
from serverless_gpu import distributed

  3. Verpakt de modeltrainingscode in een functie en verfraaid de functie met de @distributed decorator.

    @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 deze wordt uitgevoerd, wordt er een MLflow-uitvoeringskoppeling gegenereerd in de uitvoer van de notebookcel. Klik op de koppeling voor de MLflow-uitvoering of zoek deze in het deelvenster Experiment om de uitvoeringsresultaten te bekijken.

Details van gedistribueerde uitvoering

Serverloze GPU-API bestaat uit verschillende belangrijke onderdelen:

  • Compute Manager: verwerkt resourcetoewijzing en -beheer
  • Runtime-omgeving: Beheert Python-omgevingen en -afhankelijkheden
  • Startprogramma: uitvoering en bewaking van taken organiseren

Wanneer deze wordt uitgevoerd in de gedistribueerde modus:

  • De functie wordt geserialiseerd en verdeeld over het opgegeven aantal GPU's
  • Elke GPU voert een kopie van de functie uit met dezelfde parameters
  • De omgeving wordt gesynchroniseerd voor alle GPU's
  • Resultaten worden verzameld en geretourneerd door alle GPU's

De API ondersteunt populaire parallelle trainingsbibliotheken zoals Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP), DeepSpeed.

U vindt meer echte gedistribueerde trainingsscenario's met behulp van de verschillende bibliotheken in notebookvoorbeelden.

Veelgestelde vragen

Waar moet de code voor het laden van gegevens worden geplaatst?

Wanneer u de serverloze GPU-API gebruikt voor gedistribueerde training, verplaatst u de code voor het laden van gegevens binnen de @distributed decorator. De datasetgrootte kan groter zijn dan de maximale grootte die is toegestaan door de module pickle, dus het wordt aanbevolen om de dataset in de decorator te genereren, zoals hieronder getoond.

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

Meer informatie

Raadpleeg de documentatie van de Serverloze GPU Python-API voor de API-referentie.

  • Last updated on