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.

remarque

Cette page couvre la formation distribuée à partir des notebooks Databricks avec l’API Python GPU serverless. Pour soumettre des charges de travail d'entraînement distribuées depuis votre machine locale, utilisez le CLI AI Runtime, qui est en Public Preview. Voir CLI d'AI Runtime.

Vous pouvez lancer des charges de travail distribuées sur plusieurs GPU sur un seul nœud à l’aide de l’ API Python Serverless GPU. L’API fournit une interface simple et unifiée qui fait abstraction des détails du provisionnement du GPU, de la configuration de l’environnement et de la distribution des charges de travail. Avec des changements de code minimes, vous pouvez passer de manière transparente de l’entraînement sur un seul GPU à l’exécution distribuée sur plusieurs GPU depuis le même Notebook.

remarque

L'entraînement distribué nécessite un accélérateur 8xH100, qui provisionne un nœud unique avec 8 GPU. Lorsque vous utilisez le décorateur @distributed, définissez gpus=8. Le paramètre gpu_type est facultatif et automatiquement détecté à partir de l'accélérateur auquel votre Notebook est connecté.

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.

API serverless_gpu vs. TorchDistributor

Le tableau suivant compare l'API serverless_gpu @distributed avec TorchDistributor:

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

L'API serverless_gpu est l'approche recommandée pour les nouvelles charges de travail de deep learning sur Databricks. TorchDistributor reste disponible pour les charges de travail étroitement couplées aux clusters Spark.

Quick start

L'API GPU serverless pour la formation distribuée est préinstallée lorsque vous êtes connecté à un GPU serverless au sein des notebooks et des jobs Databricks. Nous recommandons l'environnement GPU 4 et supérieur. Pour l'utiliser pour la formation distribuée, importez et utilisez le décorateur distributed pour distribuer votre fonction de formation.

Encapsulez 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 devient le point d'entrée pour l'exécution distribuée, ainsi, toute la logique d'entraînement, le chargement des données et l'initialisation du modèle doivent être définis à l'intérieur de cette fonction.

Pour lancer l'exécution distribuée, appelez votre fonction décorée en utilisant train_function.distributed(). Chaque appel crée automatiquement une exécution d'Expérimentation MLflow, ou une exécution enfant imbriquée si une exécution est déjà active.

attention

Si vous définissez gpu_type dans @distributed, assurez-vous qu'il correspond au type d'accélérateur auquel votre Notebook est connecté ("H100" ou "A10"). Spécifier le mauvais type d'accélérateur entraînera l'échec de la charge de travail.

L'extrait de code ci-dessous montre l'utilisation de base de @distributed:

Python
from serverless_gpu import distributed

# Decorate your training function with @distributed and specify the number of GPUs.
# gpu_type='H100' is optional and will be auto-detected if not set.
@distributed(gpus=8, gpu_type='H100')
def run_train():
    ...

run_train.distributed()

Voici un exemple complet qui 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
```

Encapsulez le code d'entraînement du modèle dans une fonction et décorez la fonction avec le décorateur @distributed.

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.

Détails de l'exécution distribuée

L'API GPU Serverless se compose de plusieurs composants clés :

Gestionnaire de compute : gère l’allocation et la gestion des ressources
Environnement Runtime : Gère les environnements Python et les dépendances.
Launcher : orchestre l'exécution et le monitoring des Jobs

Lors de l'exécution en mode distribué :

La fonction est sérialisée et distribuée sur le nombre spécifié de GPU
Chaque GPU exécute une copie de la fonction avec les mêmes paramètres.
L’environnement est synchronisé sur tous les GPU
Les résultats sont collectés et renvoyés de tous les GPU.
Gestion du cycle de vie : l’exécution distribuée s’exécute dans le cycle de vie du Notebook. Lorsque le Notebook se termine, l’exécution aussi. Le décorateur @distributed a un délai d’expiration par default de 3 heures. Pour définir un délai d’expiration personnalisé, transmettez timeout en secondes ou timeout=None pour le désactiver. Le paramétrage d’un délai d’expiration est disponible dans l’environnement GPU v5 et versions ultérieures.

L'API prend en charge des bibliothèques d'entraînement parallèle populaires telles que Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP), DeepSpeed.

Vous trouverez d'autres scénarios réels de formation distribuée utilisant les différentes bibliothèques dans les exemples de notebooks.

FAQ

Où le code de chargement des données doit-il être placé ?

Lorsque vous utilisez l'API GPU Serverless pour l'entraînement distribué, déplacez le code de chargement des données à l'intérieur du décorateur @distributed. La taille du dataset peut dépasser la taille maximale autorisée par pickle, il est donc recommandé de générer le dataset à l'intérieur du décorateur, comme indiqué ci-dessous :

Python
from serverless_gpu import distributed

# this may cause pickle error
dataset = get_dataset(file_path)
@distributed(gpus=8, gpu_type='H100')
def run_train():
  # good practice
  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.

En savoir plus

Pour la référence de l'API, consultez la documentation de l'API Python Serverless GPU.

Frameworks pris en charge​

API serverless_gpu vs. TorchDistributor​

Quick start​

Détails de l'exécution distribuée​

FAQ​

Où le code de chargement des données doit-il être placé ?​

En savoir plus​