Charge de travail multi-GPU

Important

Cette fonctionnalité est en version bêta. Les administrateurs d’espace de travail peuvent contrôler l’accès à cette fonctionnalité à partir de la page Aperçus . Consultez Gérer les préversions d’Azure Databricks.

Vous pouvez lancer des charges de travail distribuées sur plusieurs GPU sur un seul nœud à l’aide de l’API Python GPU serverless. L’API fournit une interface simple et unifiée qui extrait les détails de l’approvisionnement gpu, de la configuration de l’environnement et de la distribution des charges de travail. Avec des modifications de code minimales, vous pouvez passer sans effort de l'entraînement sur GPU unique à l'exécution répartie sur plusieurs GPU dans le même notebook.

Frameworks pris en charge

L’API @distributed s’intègre à des bibliothèques de formation distribuées majeures :

PyTorch Distributed Data Parallel (DDP) : parallélisme de données multi-GPU standard.
Données entièrement partitionnées parallèles (FSDP) : entraînement à l’efficacité de la mémoire pour les grands modèles.
DeepSpeed : bibliothèque d’optimisation de Microsoft pour l’entraînement de modèles volumineux.

API serverless_gpu et TorchDistributor

Le tableau suivant compare l’API serverless_gpu@distributed à TorchDistributor :

Fonctionnalité	`serverless_gpu` `@distributed` API	DistributeurDeTorches
Infrastructure	Entièrement serverless, aucune gestion de cluster	Nécessite un cluster Spark avec des workers GPU
Paramétrage	Décorateur unique, configuration minimale	Nécessite l’installation d’un cluster Spark et de TorchDistributor
Support de plate-forme	PyTorch DDP, FSDP, DeepSpeed	Principalement PyTorch DDP
Chargement des données	À l'intérieur du décorateur, utilise des volumes Unity Catalog	Via Spark ou le système de fichiers

L’API serverless_gpu est l’approche recommandée pour les nouvelles charges de travail d’apprentissage profond sur Databricks. TorchDistributor reste disponible pour les charges de travail étroitement couplées aux clusters Spark.

Démarrage rapide

L'API GPU sans serveur pour l'entraînement distribué est préinstallée lorsque vous êtes connecté à un GPU sans serveur dans les notebooks et les tâches de Databricks. Nous vous recommandons l’environnement GPU 4 et versions ultérieures. Pour l’utiliser pour l’entraînement distribué, importez et utilisez le « décorateur » distributed afin de distribuer votre fonction d’entraînement.

Encapsulez le code d’entraînement du modèle dans une fonction et décorez la fonction avec le @distributed décorateur. La fonction décorée devient le point d’entrée pour l’exécution distribuée : toute logique d’entraînement, le chargement des données et l’initialisation du modèle doivent être définis à l’intérieur de cette fonction.

Avertissement

Le gpu_type paramètre dans @distributed doit correspondre au type d’accélérateur auquel votre notebook est connecté. Par exemple, @distributed(gpus=8, gpu_type='H100') nécessite que votre bloc-notes soit connecté à un accélérateur H100. L’utilisation d’un type d’accélérateur incompatible (par exemple, la connexion à A10 lors de la spécification de H100) entraîne l’échec de la charge de travail.

L’extrait de code ci-dessous montre l’utilisation de base de @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():
    ...

Voici un exemple complet qui entraîne un modèle de perceptron multicouche (MLP) sur 8 GPU H100 à partir d’un notebook :

Configurez votre modèle et définissez des fonctions utilitaires.


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

Importez la bibliothèque serverless_gpu et le module distribué .

import serverless_gpu
from serverless_gpu import distributed

Encapsulez le code d’entraînement du modèle dans une fonction et décorez la fonction avec le @distributed décorateur.

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

Exécutez l’entraînement distribué en appelant la fonction distribuée avec des arguments définis par l’utilisateur.
```
run_train.distributed(num_epochs=3, batch_size=1)
```
Lors de l’exécution, un lien d’exécution MLflow est généré dans la sortie de cellule du notebook. Cliquez sur le lien d’exécution MLflow ou recherchez-le dans le panneau Expérience pour afficher les résultats de l’exécution.

Détails de l’exécution distribuée

L’API GPU serverless se compose de plusieurs composants clés :

Gestionnaire de calcul : gère l’allocation et la gestion des ressources
Environnement d’exécution : gère les environnements et dépendances Python
Lanceur : orchestre l’exécution et la surveillance des travaux

Lors de l’exécution en mode distribué :

La fonction est sérialisée et distribuée sur le nombre spécifié de GPU
Chaque GPU exécute une copie de la fonction avec les mêmes paramètres
L’environnement est synchronisé sur toutes les GPU
Les résultats sont collectés et retournés à partir de toutes les GPU

L’API prend en charge les bibliothèques d’entraînement parallèles populaires telles que Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP), DeepSpeed.

Vous trouverez des scénarios d’entraînement distribués plus réalistes en utilisant différentes bibliothèques dans des exemples de notebooks.

Questions fréquentes (FAQ)

Où le code de chargement des données doit-il être placé ?

Lorsque vous utilisez l’API GPU serverless pour l’entraînement distribué, déplacez le code de chargement des données à l’intérieur du décorateur @distributed . La taille du jeu de données peut dépasser la taille maximale autorisée par pickle. Il est donc recommandé de générer le jeu de données à l’intérieur du décorateur, comme indiqué ci-dessous :

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

En savoir plus

Pour la référence API, consultez la documentation de l’API Python Serverless GPU.

Commentaires

Est-ce que cette page vous a été utile?

Last updated on 2026-03-21