Entraînement distribué dans les carnets

Important

Cette fonctionnalité est en version bêta.

Le @distributed décorateur du SDK Databricks AI Runtime (databricks-sdk-air) est la manière la plus pratique d’exécuter l’entraînement distribué depuis 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.

Points clés sur cette page :

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 ordinateur portable à un accélérateur 8xH100 pour un entraînement complet multi-GPU. Par défaut, @distributed utilise tous les GPU du nœud.

Note

Cette page couvre l’entraînement distribué depuis les notebooks Databricks avec le SDK Databricks AI Runtime. Pour soumettre des charges d’entraînement distribuées depuis votre machine locale, utilisez les commandes CLI Databricks pour AI Runtime. Voir Utiliser l’interface de ligne de commande Databricks avec AI Runtime.

Quickstart

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

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

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, qui provisionne un seul nœud avec 8 GPU. Par défaut, @distributed utilise les huit GPU. Réglez num_accelerators pour en utiliser moins. Le accelerator_type paramètre est également optionnel et détecté automatiquement par 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 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 exemples d’entraînement distribué.

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 fan-out : La fonction est sérialisée et lancée sur chaque GPU sélectionné. 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.

Migrer depuis `serverless_gpu`

Le serverless_gpu package est obsolète et ne reçoit plus de nouvelles fonctionnalités. À partir de l’environnement GPU version 6, le SDK d’exécution IA Databricks (databricks-sdk-air) le remplace et fournit des API équivalentes, avec des arguments renommés, sous l’espace databricks.air de noms. 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 : Inclut les deux packages, mais seul le SDK d’exécution IA reçoit les mises à jour. La version 6 est la dernière version de l’environnement GPU qui inclut serverless_gpu.

Le SDK d’exécution IA est un paquet distinct du SDK Databricks pour Python (databricks-sdk).

	`serverless_gpu` (déconseillé)	SDK d’exécution IA (`databricks.air`)
Versions de l’environnement GPU	6 ans et moins	6 ans et plus
Import	`from serverless_gpu import distributed, ray_init`	`from databricks.air import distributed, ray_init`
Argument de nombre de GPU	`gpus`	`num_accelerators`
Argument de type de 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 sauvegarde d’état	`serverless_gpu.data`	`databricks.air.data`. Consultez Charger des données sur AI Runtime.

Pour migrer un notebook :

Connectez le notebook à l’environnement GPU version 6 ou supérieure, qui inclut le SDK d’exécution IA.
Changer les importations de serverless_gpu vers databricks.air.
Renommer gpus en num_accelerators et gpu_type en accelerator_type, et utiliser les nouvelles valeurs de type d’accélérateur. Les anciens noms d’arguments et valeurs de type de GPU fonctionnent toujours dans databricks.air mais produisent un FutureWarning. Le support prendra fin pour la version 7 de l’environnement GPU.

Tip

Indiquez au code Genie de procéder ainsi pour vous :

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

Pour les workloads Databricks CLI qui utilisent databricks.air.data, ajouter databricks-sdk-air[data] au environment.dependencies du workload.

Venant de TorchDistributor

Si vous exécutez un entraînement distribué via 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.

Feature	`@distributed` Interface de Programmation d'Applications (API)	DistributeurDeTorches
Infrastructure	Entièrement serverless, aucune gestion de cluster	Nécessite un cluster Spark avec des workers GPU
Configuration	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. Connecte ton ordinateur portable à un accélérateur 8xH100 à la place. @distributed utilise les huit GPU par défaut.
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.

Le @distributed décorateur est en version Bêta et 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 un entraînement distribué depuis votre machine locale ou sur plusieurs nœuds, utilisez les commandes CLI Databricks pour AI Runtime. Voir Utiliser l’interface de ligne de commande Databricks avec AI Runtime.

Exemple complet

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

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)

Importez le distributed décorateur depuis le SDK d’exécution IA.
```
from databricks.air import distributed
```

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

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 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 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 stockées sous forme de fichiers dans des volumes d’Unity Catalog, utilisez UCVolumeDataset de databricks.air.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. Voir Charger des données dans AI Runtime et Checkpoint avec Distributed Checkpoint (DCP).

Limitations

@distributed Il 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 entraîner sur plusieurs nœuds, utilisez la CLI AI Runtime Databricks. Voir Utiliser l’interface de ligne de commande Databricks avec AI Runtime.
Si vous définissez accelerator_type, il doit correspondre à l’accélérateur auquel votre carnet est connecté (par exemple, "GPU_8xH100"). 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 prend en charge l’environnement GPU v4 et supérieur. 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.

Ressources additionnelles

Pour le décorateur @distributed, AcceleratorType et les API Ray, voir la référence de l’API du SDK Databricks AI Runtime.
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 d’entraînement distribué.

Commentaires

Cette page a-t-elle été utile ?

Dernière mise à jour : 2026-10-22