Entraînement distribué dans les carnets

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.

Le décorateur @distributed de l’API Python Serverless GPU est le moyen le plus pratique d’exécuter un entraînement distribué à partir d’un notebook Databricks. Décorez votre fonction d’entraînement, appelez-la, et AI Runtime la fait tourner sur tous les GPU du nœud auquel votre notebook est connecté. Le même code évolue d’un seul GPU à un multi-GPU sans cluster à provisionner et sans lanceur distribué à configurer.

Tip

• Le décorateur @distributed exécute une fonction d’entraînement sur tous les GPU de votre nœud depuis un notebook.
• Il prend en charge PyTorch DDP, FSDP et DeepSpeed, et déplace le code GPU unique vers multi-GPU avec un minimum de modifications.
• Connectez votre notebook à un accélérateur 8xH100 et définissez gpus=8 pour un entraînement multi-GPU complet.

Quickstart

Le serverless_gpu package est préinstallé lorsque votre ordinateur portable est connecté à un GPU serverless. Décorez votre activité d’entraînement avec @distributed, puis appelez-la avec .distributed():

from serverless_gpu import distributed

# gpus is the number of GPUs on the node. gpu_type is optional and
# auto-detected from the accelerator your notebook is connected to.
@distributed(gpus=8, gpu_type="H100")
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()

Chaque appel à .distributed() crée une exécution MLflow (ou une exécution enfant imbriquée si une exécution est déjà active) et affiche un lien vers l’exécution dans la sortie de la cellule. Pour un guide complet et exécutable, voir l’exemple complet.

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.
• Fully Sharded Data Parallel (FSDP) : entraînement économe en mémoire pour les modèles de grande taille.
• DeepSpeed : bibliothèque d'optimisation de Microsoft pour l'entraînement de modèles volumineux.

Pour des scénarios d’entraînement réels utilisant chaque bibliothèque, voir les exemples de carnets.

Comment travaille le @distributed décorateur

Lorsque vous appelez une fonction décorée avec .distributed(), AI Runtime gère les mécaniques que vous configureriez autrement à la main avec un lanceur distribué :

• Sérialisation et distribution : la fonction est sérialisée et lancée sur chacun des gpus que vous demandez. Chaque carte graphique exécute une copie de la fonction avec les mêmes arguments.
• Synchronisation de l’environnement : L’environnement Python et ses dépendances sont répliqués à tous les niveaux, donc chaque processus exécute le même code.
• Variables d’environnement de classement : variables standard telles que LOCAL_RANK sont remplies pour chaque processus. Lisez-les dans votre fonction pour placer le modèle et les données sur le bon appareil.
• Collecte des résultats : Les valeurs de retour sont collectées de tous les rangs et retournées à l’appelant.
• Suivi MLflow : Chaque .distributed() appel crée une exécution MLflow, ou une exécution enfant imbriquée si une exécution est déjà active, de sorte que les métriques enregistrées depuis votre fonction se retrouvent sur la même exécution.
• Cycle de vie et délai d’expiration : L’exécution distribuée s’exécute dans le cycle de vie du notebook. Terminer le notebook met fin au cours. Le décorateur a un délai d’expiration par défaut de 3 heures. Passe timeout en quelques secondes pour le changer, ou timeout=None le désactiver. Les délais personnalisés nécessitent un environnement GPU v5 et supérieur.

L’API s’appuie sur les bibliothèques standard PyTorch : Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP) et DeepSpeed.

Venant de TorchDistributor

Si vous lancez PyTorch distribué sur Spark aujourd’hui avec TorchDistributor et que votre charge de travail tient sur un seul nœud, l’API serverless_gpu@distributed est le remplaçant recommandé pour les nouvelles charges de travail en apprentissage profond. Cela supprime le cluster Spark et vous donne le même chemin de code entre un seul GPU et plusieurs GPU.

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	Dans l’élément décoratif, utilise les volumes Unity Catalog (UCVolumeDataset pour les données de fichiers en streaming)	Via Spark ou le système de fichiers

Pour migrer une charge de travail à nœud unique :

• Remplacez l’appel TorchDistributor(...).run(train_fn, ...) par le décorateur @distributed dans train_fn, puis lancez avec train_fn.distributed(...).
• Supprime le cluster Spark et la configuration du GPU Worker. Connectez votre ordinateur portable à un accélérateur 8xH100 et réglez gpus=8 à la place.
• Déplacez le chargement des données à l’intérieur de la fonction décorée. Voir Chargement des données.
• Conservez votre code de modèle DDP, FSDP ou DeepSpeed existant. L’élément décoratif prend en charge les trois.

@distributed fonctionne sur un seul nœud (voir Limitations), il ne remplace donc pas toutes les charges de travail TorchDistributor. Gardez les charges de travail qui dépendent de l’intégration Spark sur TorchDistributor. Pour exécuter une formation distribuée depuis votre machine locale ou sur plusieurs nœuds, utilisez plutôt la CLI d’exécution IA, qui se trouve en Aperçu public. Consultez l’interface CLI d’AI Runtime.

Exemple complet

L’exemple suivant entraîne un modèle de perceptron multicouche (MLP) sur 8 GPU H100 depuis un ordinateur portable.

1. 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():
       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. Importez la serverless_gpu bibliothèque et le distributed module.

   import serverless_gpu
   from serverless_gpu import distributed
   

3. Encapsulez le code d’entraînement du modèle dans une fonction et décorez la fonction avec l’élément décoratif @distributed. La fonction décorée est le point d’entrée pour l’exécution distribuée, donc définissez toute la logique d’entraînement, le chargement des données et l’initialisation des modèles à l’intérieur de celle-ci.

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

5. Lors de l’exécution, un lien d’exécution MLflow est généré dans la sortie de la cellule de 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. Pour plus d’informations sur la personnalisation des noms d’expériences, le suivi des métriques et la reprise des exécutions, consultez Suivi des expériences et observabilité.

Chargement des données

Placez le code de chargement des données dans la @distributed fonction. Un jeu de données peut dépasser la taille maximale autorisée par pickle, donc le générer ou le charger dans le décorateur évite les erreurs de sérialisation :

from serverless_gpu import distributed

# This may cause a pickle error because the dataset is captured by the function.
dataset = get_dataset(file_path)

@distributed(gpus=8, gpu_type='H100')
def run_train():
    # Load the dataset inside the decorated function instead.
    dataset = get_dataset(file_path)
    ...

Pour les données stockées sous forme de fichiers dans des volumes d’Unity Catalog, utilisez UCVolumeDataset de serverless_gpu.data, qui permet de diffuser les fichiers avec mise en cache locale et de les partitionner automatiquement entre les rangs et les workers. Pour point de contrôle de l’entraînement distribué sur un volume, utilisez UCVolumeWriter et UCVolumeReader. Consultez Charger des données sur AI Runtime et Enregistrement de point de contrôle du modèle.

Limitations

• L’entraînement distribué s’exécute entre les GPU sur le nœud unique auquel votre ordinateur portable est connecté. Pour un entraînement complet multi-GPU, connectez-vous à un accélérateur 8xH100, qui fournit un nœud avec 8 GPU, et réglez gpus=8.
• Le type d’accélérateur doit correspondre. Si vous définissez gpu_type dans @distributed, il doit correspondre à l’accélérateur auquel votre ordinateur portable est connecté ("H100" ou "A10"). Une incompatibilité entraîne l’échec de la charge de travail. Le paramètre est optionnel et détecté automatiquement lorsqu’il est omis.
• AI Runtime recommande l’environnement GPU v4 et les versions supérieures. Les délais personnalisés (le timeout paramètre) nécessitent l’environnement GPU v5 et supérieur.
• Le décorateur expire après 3 heures par défaut. Passe timeout en quelques secondes pour le changer, ou timeout=None le désactiver.
• L’exécution s’effectue dans le cycle de vie du carnet. Terminer le notebook met fin au cours.

En savoir plus

• Pour le décorateur @distributed, GPUType et les API Ray, consultez la documentation de référence des API Python Serverless GPU.
• Pour des schémas qui rendent votre pipeline d’entraînement plus efficace et résilient, consultez le guide de la performance et de la résilience.
• Pour des scénarios d’entraînement de bout en bout, voir les exemples de carnets.