Formation distribuée dans les notebooks

info

Bêta

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 provisionnement de clusters 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 d’un code mono-GPU à un code multi-GPU avec un minimum de modifications.
• Connectez votre notebook à un accélérateur 8xH100 pour un entraînement multi-GPU complet. Par default, @distributed utilise chaque GPU du nœud.

remarque

Cette page décrit l'entraînement distribué à partir des notebooks Databricks avec le SDK Databricks AI Runtime. Pour soumettre des workloads d’entraînement distribué depuis votre machine locale, utilisez les commandes CLI de Databricks pour AI Runtime, qui sont en préversion publique. Consultez Utiliser la CLI Databricks avec AI Runtime.

Démarrage rapide

Les exemples de cette page utilisent le SDK Databricks AI Runtime, qui est préinstallé dans l’environnement GPU version 6 et supérieure. Sur la version 5 et inférieure, utilisez le package serverless_gpu à la place. Voir Migrer depuis serverless_gpu. Pour connecter un notebook à un accélérateur, voir Connecter à un notebook.

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, qui provisionne un seul nœud avec 8 GPU. default, @distributed utilise les huit GPU. Définissez num_accelerators pour en utiliser moins. Le parameter accelerator_type est également facultatif et automatiquement détecté à partir de l’accélérateur auquel votre Notebook est connecté.

Chaque appel à .distributed() crée un run MLflow (ou un run enfant imbriqué si un run est déjà actif) et affiche un Link vers le run dans le résultat de la cellule. Pour consulter un exemple complet et exécutable, consultez Exemple complet.

Frameworks pris en charge

L'API @distributed s'intègre aux principales bibliothèques de formation distribuée :

• PyTorch Distributed Data Parallel (DDP) : parallélisme de données multi-GPU standard.
• **Fully Sharded Data Parallel (FSDP)** : entraînement efficace en mémoire pour les grands modèles.
• DeepSpeed : La 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 d’entraînement distribué.

Fonctionnement du décorateur @distributed

Lorsque vous appelez une fonction décorée avec .distributed(), AI Runtime gère les mécanismes que vous auriez autrement configurés manuellement 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 Python et les dépendances sont répliqués sur tous les rangs, de sorte que chaque processus exécute le même code.
• Variables d'environnement de rang : 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 bon périphérique.
• Collecte des résultats : les valeurs de retour sont collectées à partir de tous les rangs et renvoyées à l'appelant.
• Suivi MLflow : chaque appel de .distributed() crée un run MLflow, ou un run enfant imbriqué si un run est déjà actif, de sorte que les métriques enregistrées depuis votre fonction aboutissent sur le même run.
• Lifecycle and timeout : l’exécution distribuée s’exécute pendant le cycle de vie du notebook. La fermeture du notebook met fin à l’exécution. Le décorateur a un default timeout de 3 heures. Transmettez 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 ou supérieur.

L'API s'appuie 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 bénéficie plus de nouvelles fonctionnalités. À partir de la version 6 de l’environnement GPU, le SDK Databricks AI Runtime (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 dans ces versions.
• Version 6 : comprend les deux packages, mais seul le SDK de l'environnement d'exécution de l'IA 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 séparé 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 de compte GPU	gpus	num_accelerators
Argument de type GPU	gpu_type (GPUType)	accelerator_type (AcceleratorType)
Valeurs de type 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. Voir Charger des données sur 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 de compte GPU	gpus	num_accelerators
Argument de type GPU	gpu_type (GPUType)	accelerator_type (AcceleratorType)
Valeurs de type 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. Voir Charger des données sur AI Runtime.

Pour migrer un notebook :

1. Connectez le notebook à la version 6 ou supérieure de l’environnement GPU, qui inclut le SDK AI Runtime.
2. Modifier les importations de serverless_gpu à 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’arguments et valeurs de type GPU fonctionnent toujours dans databricks.air mais émettent un FutureWarning. La prise en charge prendra fin dans la version 7 de l'environnement GPU.

prompt

Indiquez à 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 CLI Databricks qui utilisent databricks.air.data, ajouter databricks-sdk-air[data] au environment.dependencies de la charge de travail.

Provenant de TorchDistributor

Si vous effectuez un entraînement distribué à l'aide de TorchDistributor sur des clusters Spark et que votre charge de travail s'adapte sur un nœud unique, l'API @distributed est le remplacement recommandé sur AI Runtime. Le code de votre modèle DDP, FSDP ou DeepSpeed ne change pas.

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

Fonctionnalité	@distributed API	TorchDistributor
Infrastructure	Entièrement serverless, aucune gestion de clusters	Nécessite un cluster Spark avec des Worker GPU
Installer	Un seul décorateur, une configuration minimale	Nécessite la configuration d'un cluster Spark et de TorchDistributor
Prise en charge du framework	PyTorch DDP, FSDP, DeepSpeed	Principalement PyTorch DDP
Chargement des données	À l'intérieur du décorateur, utilise des 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 sur train_fn, puis lancez avec train_fn.distributed(...).
• Supprimez le cluster Spark et la configuration du worker GPU. Connectez votre notebook à un accélérateur 8xH100 à la place. @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.

@distributed s’exécute sur un seul nœud (consultez les Limitations) ; il ne remplace donc pas toutes les charges de travail TorchDistributor. Conservez sur TorchDistributor les charges de travail qui dépendent de l’intégration Spark. Pour exécuter un entraînement distribué à partir de votre machine locale ou sur plusieurs nœuds, utilisez plutôt la CLI AI Runtime, qui est en préversion publique. Consultez Utiliser la CLI Databricks avec AI Runtime.

Exemple complet

L'exemple suivant entraîne un modèle de 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 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 en son sein.

   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. Lorsqu'il est exécuté, un Link d'exécution MLflow est généré dans la sortie de la cellule du notebook. Cliquez sur le MLflow run Link ou recherchez-le dans le panneau Expérimentation pour voir les résultats de l'exécution. Pour plus de détails sur la personnalisation des noms d'expérimentation, le suivi des métriques et la reprise des exécutions, consultez le suivi et l'observabilité de l'Experimentation.

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. Par conséquent, le générer ou le charger à l’intérieur du décorateur permet d’éviter 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)
    ...

For file-based data stored in Unity Catalog volumes, use UCVolumeDataset from databricks.air.data, which streams files with local caching and partitions them across ranks and workers automatically. Pour enregistrer l'état de l'entraînement distribué dans 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 prend en charge un seul nœud : il s'exécute sur les GPU du nœud auquel votre notebook est connecté. Pour entraîner sur plusieurs nœuds, utilisez la CLI Databricks AI Runtime. 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 incohérence entraîne l’échec du workload. Le parameter est facultatif et détecté automatiquement en cas d’omission.
• AI Runtime prend en charge la version 4 et supérieure de l'environnement GPU. Les délais d’attente personnalisés (le parameter timeout) nécessitent la version 5 et supérieure de l’environnement GPU.
• Le temps d'exécution du décorateur expire par default au bout de 3 heures. Passez timeout en secondes pour le modifier, ou timeout=None pour le désactiver.
• L'exécution se déroule pendant le cycle de vie du notebook. La fermeture du notebook met fin à 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 découvrir les modèles qui rendent votre pipeline d'entraînement plus efficace et résilient, consultez le guide de performance et de résilience.
• Pour les scénarios de formation de bout en bout, consultez les exemples de formation distribuée.