Formation distribuée dans des notebooks

info

Cette fonctionnalité est en version Bêta.

Le décorateur @distributed du SDK Databricks AI Runtime (databricks-sdk-air) 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 l'exécute sur tous les GPU du nœud auquel votre notebook est connecté. Le même code monte en charge d'un seul GPU à plusieurs GPU sans cluster à provisionnement et sans lanceur distribué à configurer.

Points clés de cette page :

• Le décorateur @distributed exécute une fonction d’entraînement sur chaque GPU de votre nœud à partir d’un notebook.
• Il prend en charge PyTorch DDP, FSDP et DeepSpeed, et permet de passer du code monopuce à du code multipuce avec un minimum de modifications.
• Connectez votre notebook à un accélérateur 8xH100 pour un entraînement multi-GPU complet. Par default, @distributed utilise tous les GPU du nœud.

remarque

Cette page couvre l'entraînement distribué à partir des notebooks Databricks avec le SDK Databricks AI Runtime. Pour soumettre des charges de travail d'entraînement distribuées depuis votre machine locale, utilisez les commandes CLI Databricks pour AI Runtime. Consultez Utiliser la CLI Databricks avec AI Runtime.

Démarrage rapide

Les exemples sur cette page utilisent le SDK Databricks AI Runtime, qui est préinstallé dans l’environnement GPU version 6 et ultérieure. Sur la version 5 et inférieure, utilisez plutôt le package serverless_gpu. Voir Migrer depuis serverless_gpu. Pour connecter un notebook à un accélérateur, voir Étape 1 : associer un notebook à un GPU.

Décorez votre fonction d’entraînement avec @distributed, puis appelez-la avec .distributed():

Python

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

L'entraînement distribué nécessite un accélérateur 8xH100 ou 8xB300, ce qui provisionne un seul nœud avec 8 GPU. default, @distributed utilise les huit GPU. Définissez num_accelerators sur utiliser moins. Le parameter accelerator_type est également facultatif et est automatiquement détecté à partir de l’accélérateur auquel votre Notebook est connecté.

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 imprime un Link d’exécution dans la sortie de la cellule. Pour une présentation complète et exécutable, voir Exemple complet.

Frameworks pris en charge

L’API @distributed s’intègre aux principales bibliothèques d’entraînement distribué :

• PyTorch Distributed Data Parallel (DDP) : parallélisme de données multi-GPU standard.
• Fully Sharded Data Parallel (FSDP) : entraînement économiseur de mémoire pour les modèles de grande taille.
• DeepSpeed : bibliothèque d'optimisation de Microsoft pour l'entraînement de grands modèles.

Pour connaître les scénarios d'entraînement réels qui utilisent chaque bibliothèque, consultez les exemples de formation distribuée.

Fonctionnement du décorateur @distributed

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

• Sérialisation et fan-out : la fonction est sérialisée et lancée sur chaque GPU sélectionné. Chaque GPU exécute une copie de la fonction avec les mêmes arguments.
• Synchronisation de l’environnement : l’environnement et les dépendances Python sont répliqués sur tous les rangs, afin que chaque processus exécute le même code.
• Rank des variables d'environnement : les variables standard telles que LOCAL_RANK sont renseignées pour chaque processus. Lisez-les dans votre fonction pour placer le modèle et les données sur le périphérique approprié.
• Collecte des résultats : les valeurs de retour sont collectées sur tous les rangs et renvoyées à l’appelant.
• Suivi MLflow : chaque appel .distributed() crée une exécution MLflow ou une exécution enfant imbriquée si une exécution est déjà active. Ainsi, les métriques enregistrées à partir de votre fonction seront associées à la même exécution.
• Cycle de vie et délai d'expiration : l'exécution distribuée s'effectue pendant le cycle de vie du notebook. La terminaison du notebook met fin à l’exécution. Le décorateur a un délai d’expiration default de 3 heures. Passez timeout en secondes pour le modifier, ou timeout=None pour le désactiver. Les délais d'expiration personnalisés nécessitent l'environnement GPU v5 et supérieur.

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

Migrer depuis serverless_gpu

Le package serverless_gpu est obsolète et ne reçoit plus de nouvelles fonctionnalités. À partir de la version 6 de l’environnement GPU, le Databricks AI Runtime SDK (databricks-sdk-air) le remplace et fournit des APIs équivalentes, avec des arguments renommés, sous l’espace de noms databricks.air. Le package que vous utilisez dépend de la version de l’environnement GPU :

• Version 5 et inférieure : inclure uniquement serverless_gpu. Continuez à utiliser serverless_gpu sur ces versions.
• Version 6 : comprend les deux packages, mais seul le SDK d’AI Runtime reçoit des mises à jour. La version 6 est la dernière version de l’environnement GPU qui comprend serverless_gpu.

Le SDK AI Runtime est un package distinct du SDK Databricks pour Python (databricks-sdk).

	serverless_gpu (obsolète)	SDK AI Runtime (databricks.air)
Versions de l'environnement GPU	6 et inférieur	6 et plus
Importer	from serverless_gpu import distributed, ray_init	from databricks.air import distributed, ray_init
Argument du nombre de GPU	gpus	num_accelerators
Argument de type GPU	gpu_type (GPUType)	accelerator_type (AcceleratorType)
Valeurs du type de GPU	"A10", "H100", "B300"	"GPU_1xA10", "GPU_1xH100", "GPU_8xH100", "GPU_8xB300"
Module de chargement de données et de checkpointing	serverless_gpu.data	databricks.air.data. Consultez Charger des données dans AI Runtime.

	serverless_gpu (obsolète)	SDK AI Runtime (databricks.air)
Versions de l'environnement GPU	6 et inférieur	6 et plus
Importer	from serverless_gpu import distributed, ray_init	from databricks.air import distributed, ray_init
Argument du nombre de GPU	gpus	num_accelerators
Argument de type GPU	gpu_type (GPUType)	accelerator_type (AcceleratorType)
Valeurs du type de GPU	"A10", "H100", "B300"	"GPU_1xA10", "GPU_1xH100", "GPU_8xH100", "GPU_8xB300"
Module de chargement de données et de checkpointing	serverless_gpu.data	databricks.air.data. Consultez Charger des données dans AI Runtime.

Pour migrer un notebook :

1. Connectez le notebook à la version 6 ou supérieure de l’environnement GPU, qui inclut le SDK Runtime AI.
2. Remplacez les importations de serverless_gpu par databricks.air.
3. Renommez gpus en num_accelerators et gpu_type en accelerator_type, et utilisez les nouvelles valeurs de type d'accélérateur. Les anciens noms d’argument et valeurs de type GPU fonctionnent toujours dans databricks.air mais émettent un FutureWarning. La prise en charge se terminera avec la version 7 de l’environnement GPU.

prompt

Dites à Genie Code de le faire pour vous :

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

Pour les charges de travail de type CLI Databricks utilisant databricks.air.data, ajoutez databricks-sdk-air[data] au environment.dependencies de la charge de travail.

Issu de TorchDistributor

Si vous utilisez l'entraînement distribué avec TorchDistributor sur des clusters Spark et que votre charge de travail tient sur un seul nœud, l'API @distributed est le remplacement recommandé sur AI Runtime. Votre code de modèle DDP, FSDP ou DeepSpeed ne change pas.

Fonctionnalité	@distributed API	TorchDistributor
Infrastructure	Entièrement serverless, aucune gestion de cluster	Nécessite un cluster Spark avec des GPU Worker
Installer	Décorateur unique, configuration minimale	Nécessite la configuration d'un cluster Spark et de TorchDistributor
Prise en charge des frameworks	PyTorch DDP, FSDP, DeepSpeed	Principalement PyTorch DDP
Chargement de données	À l’intérieur du décorateur, utilise des volumes Unity Catalog (UCVolumeDataset pour les données de fichier en streaming)	Via Spark ou le système de fichiers

Fonctionnalité	@distributed API	TorchDistributor
Infrastructure	Entièrement serverless, aucune gestion de cluster	Nécessite un cluster Spark avec des GPU Worker
Installer	Décorateur unique, configuration minimale	Nécessite la configuration d'un cluster Spark et de TorchDistributor
Prise en charge des frameworks	PyTorch DDP, FSDP, DeepSpeed	Principalement PyTorch DDP
Chargement de données	À l’intérieur du décorateur, utilise des volumes Unity Catalog (UCVolumeDataset pour les données de fichier 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 sur train_fn, puis lancez avec train_fn.distributed(...).
• Supprimez la configuration du cluster Spark et du worker GPU. Connectez plutôt votre notebook à un accélérateur 8xH100. @distributed utilise les huit GPU default.
• 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. Le décorateur prend en charge les trois.

Le décorateur @distributed est en bêta et s’exécute sur un seul nœud (voir Limites), il ne remplace donc pas toutes les charges de travail de TorchDistributor. Conservez les charges de travail dépendant de l’intégration Spark sur TorchDistributor. Pour effectuer un entraînement distribué à partir de votre machine locale ou sur plusieurs nœuds, utilisez les commandes de la CLI Databricks pour AI Runtime. Voir Utiliser la CLI Databricks avec AI Runtime.

Exemple complet

L'exemple suivant entraîne un modèle perceptron multicouche (MLP) sur 8 GPU H100 à partir d'un notebook.

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

   Python

   
   # 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 le décorateur distributed à partir du SDK AI Runtime.

   Python

   from databricks.air import distributed

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

   Python

   @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. Exécutez l'entraînement distribué en appelant la fonction distribuée avec des arguments définis par l'utilisateur.

   Python

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

5. Lors de son exécution, un Link d'exécution MLflow est généré dans la sortie de la cellule du Notebook. Cliquez sur le Link d’exécution MLflow ou trouvez-le dans le panneau Experimentation pour voir les résultats de l’exécution. Pour plus de détails sur la personnalisation des noms d’expérimentations, le suivi des métriques et la reprise des exécutions, consultez Suivi et observabilité des expérimentations.

Chargement des données

Placez le code de chargement des données dans la fonction @distributed. Un dataset 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 :

Python

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

Pour les données basées sur des fichiers stockées dans des volumes Unity Catalog, utilisez UCVolumeDataset de databricks.air.data, qui diffuse des fichiers avec mise en cache locale et les partitionne automatiquement entre les rangs et les Worker. Pour créer un point de contrôle de la formation distribuée sur un volume, utilisez UCVolumeWriter et UCVolumeReader. Voir Charger des données sur AI Runtime et Point de contrôle avec Distributed Checkpoint (DCP).

Limitations

• @distributed ne prend en charge qu'un seul nœud : il s'exécute sur les GPU du nœud auquel votre notebook est connecté. Pour effectuer un entraînement sur plusieurs nœuds, utilisez la CLI AI Runtime Databricks. Consultez Utiliser la CLI Databricks avec AI Runtime.
• Si vous définissez accelerator_type, il doit correspondre à l'accélérateur auquel votre notebook est connecté (par exemple, "GPU_8xH100"). Une inadéquation entraîne l’échec du workload. Le parameter est facultatif et automatiquement détecté en cas d'omission.
• AI Runtime prend en charge l’environnement GPU v4 et supérieur. Les délais d’expiration personnalisés (parameter timeout) nécessitent un environnement GPU v5 et supérieur.
• Le décorateur expire après 3 heures default. Passez timeout en secondes pour la modifier, ou timeout=None pour la désactiver.
• Les exécutions se déroulent pendant le cycle de vie du notebook. L'arrêt du notebook arrête l'exécution.

Ressources supplémentaires

• Pour le décorateur @distributed, AcceleratorType et les APIs Ray, consultez la Référence de l'API SDK Databricks AI Runtime.
• Pour des modèles qui rendent votre pipeline d'entraînement plus efficace et résilient, consultez le guide des performances et de la résilience.
• Pour des scénarios d’entraînement de bout en bout, consultez les exemples d’entraînement distribué.