Formation distribuée dans les notebooks

info

Bêta

Cette fonctionnalité est en Bêta. Les administrateurs du Workspace peuvent contrôler l'accès à cette fonctionnalité à partir de la page Previews . Consultez Gérer les aperçus Databricks.

Le décorateur @distributed de l’ API Python GPU Serverless 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 passe d’un seul GPU à plusieurs GPU sans aucun cluster à provisionner ni aucun lanceur distribué à configurer.

astuce

Principaux points à retenir

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 et définissez gpus=8 pour un entraînement multi-GPU complet.

Démarrage rapide

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

Python
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 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 des scénarios d’entraînement réels qui utilisent chaque bibliothèque, consultez les exemples de notebooks.

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 répartition (fan-out) : la fonction est sérialisée et lancée sur chacun des gpus que vous demandez. 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.

Provenant de TorchDistributor

Si vous exécutez 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 remplacement recommandé pour les nouvelles charges de travail de deep learning. Il supprime le cluster Spark et vous offre le même chemin de code du GPU unique au multi-GPU.

Fonctionnalité	`serverless_gpu` `@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é	`serverless_gpu` `@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 des worker GPU. Connectez plutôt votre notebook à un accélérateur 8xH100 et définissez gpus=8.
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 (voir Limitations) et 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 AI Runtime CLI, qui est en préversion publique. Voir AI Runtime CLI.

Exemple complet

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

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)

Importez la bibliothèque serverless_gpu et le module distributed.
Python
```
import serverless_gpu
from serverless_gpu import distributed
```

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(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.
Python
```
run_train.distributed(num_epochs=3, batch_size=1)
```
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 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 basées sur des fichiers stockées dans les volumes Unity Catalog, utilisez UCVolumeDataset de serverless_gpu.data, qui transmet les fichiers en continu avec mise en cache locale et les partitionne automatiquement entre les rangs et les workers. Pour créer un point de contrôle de la formation distribuée vers un volume, utilisez UCVolumeWriter et UCVolumeReader. Consultez Chargement des données sur l'AI Runtime et Point de contrôle de modèle.

Limitations

L’entraînement distribué s’exécute sur les GPU du nœud unique auquel votre notebook est connecté. Pour un entraînement multi-GPU complet, connectez-vous à un accélérateur 8xH100, qui provisionne un nœud avec 8 GPU, et définissez gpus=8.
Le type d'accélérateur doit correspondre. Si vous définissez gpu_type dans @distributed, cette valeur doit correspondre à l'accélérateur auquel votre notebook est connecté ("H100" ou "A10"). Une incompatibilité entraîne l'échec du workload. Le parameter est optionnel et détecté automatiquement en cas d'omission.
AI Runtime recommande l'environnement GPU v4 ou supérieur. Les délais d'attente personnalisés (le parameter timeout) nécessitent l'environnement GPU v5 ou supérieur.
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.

En savoir plus

Pour le décorateur @distributed, GPUType et les APIs Ray, consultez la documentation de référence Serverless GPU Python API.
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 des scénarios d’entraînement de bout en bout, consultez les exemples de notebooks.

Démarrage rapide​

Frameworks pris en charge​

Fonctionnement du décorateur @distributed​

Provenant de TorchDistributor​

Exemple complet​

Chargement des données​

Limitations​

En savoir plus​