Charge de travail multi-GPU et multi-nœuds

Vous pouvez lancer une charge de travail distribuée sur plusieurs GPU ( au sein d’un nœud unique ou sur plusieurs nœuds) à l’aide de l’API Python GPU serverless. L’API fournit une interface simple et unifiée qui extrait les détails de l’approvisionnement gpu, de la configuration de l’environnement et de la distribution des charges de travail. Avec des modifications de code minimales, vous pouvez passer de l’entraînement mono-GPU à l’exécution distribuée entre les GPU distants à partir du même notebook.

Démarrage rapide

L'API GPU sans serveur pour l'entraînement distribué est préinstallée dans les environnements de calcul GPU sans serveur pour les notebooks Databricks. Nous vous recommandons l’environnement GPU 4 et versions ultérieures. Pour l’utiliser pour l’entraînement distribué, importez et utilisez le « décorateur » distributed afin de distribuer votre fonction d’entraînement.

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

# Import the distributed decorator
from serverless_gpu import distributed

# Decorate your training function with @distributed and specify the number of GPUs, the GPU type,
# and whether or not the GPUs are remote
@distributed(gpus=8, gpu_type='A10', remote=True)
def run_train():
    ...

Voici un exemple complet qui entraîne un modèle de perceptron multicouche (MLP) sur 8 nœuds GPU A10 depuis un notebook :

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():
    dist.init_process_group("nccl")
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))

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 distribué .

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 @distributed décorateur.

@distributed(gpus=8, gpu_type='A10', remote=True)
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 node 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.
```
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 cellule du 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.

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

L’API GPU serverless se compose de plusieurs composants clés :

Gestionnaire de calcul : gère l’allocation et la gestion des ressources
Environnement d’exécution : gère les environnements et dépendances Python
Lanceur : orchestre l’exécution et la surveillance des travaux

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 nœuds
Les résultats sont collectés et retournés à partir de toutes les GPU

Si remote est défini sur True, la charge de travail est distribuée sur les GPU distants. Si remote est défini sur False, la charge de travail s’exécute sur le nœud GPU unique connecté au notebook actuel. Si le nœud a plusieurs puces GPU, tous seront utilisés.

L’API prend en charge les bibliothèques d’entraînement parallèles populaires telles que Distributed Data Parallel (DDP), Données Parallèles Entièrement Fragmentées (FSDP), DeepSpeed et Ray.

Vous trouverez des scénarios d’entraînement distribués plus réalistes en utilisant différentes bibliothèques dans des exemples de notebooks.

Lancer avec Ray

L’API GPU serverless prend également en charge le lancement d’un entraînement distribué à l’aide de Ray, utilisant le décorateur @ray_launch, qui est superposé sur @distributed. Chaque ray_launch tâche initialise d'abord un rendez-vous distribué avec PyTorch pour déterminer le travailleur principal de Ray et rassembler des adresses IP. Le classement zéro démarre ray start --head (avec l’exportation des métriques si activé), définit RAY_ADDRESSet exécute votre fonction décorée en tant que pilote Ray. D'autres nœuds se joignent via ray start --address et attendent que le pilote écrive un marqueur d’achèvement.

Détails de configuration supplémentaires :

Pour activer la collecte de métriques système Ray sur chaque nœud, utilisez RayMetricsMonitor avec remote=True.
Définissez les options de runtime Ray (acteurs, jeux de données, groupes de placement et planification) à l’intérieur de votre fonction décorée à l’aide d’API Ray standard.
Gérez les contrôles à l’échelle du cluster (nombre et type GPU, mode distant ou mode local, comportement asynchrone et variables d’environnement de pool Databricks) en dehors de la fonction dans les arguments du décorateur ou l'environnement du notebook.

L’exemple ci-dessous montre comment utiliser @ray_launch:

from serverless_gpu.ray import ray_launch
@ray_launch(gpus=16, remote=True, gpu_type='A10')
def foo():
    import os
    import ray
    print(ray.state.available_resources_per_node())
    return 1
foo.distributed()

Pour obtenir un exemple complet, consultez ce notebook, qui lance Ray pour entraîner un réseau neuronal Resnet18 sur plusieurs GPU A10.

Questions fréquentes (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 jeu de données peut dépasser la taille maximale autorisée par pickle. Il est donc recommandé de générer le jeu de données à l’intérieur du décorateur, comme indiqué ci-dessous :

from serverless_gpu import distributed

# this may cause pickle error
dataset = get_dataset(file_path)
@distributed(gpus=8, remote=True)
def run_train():
  # good practice
  dataset = get_dataset(file_path)
  ....

Puis-je utiliser des pools GPU réservés ?

Si le pool GPU réservé est disponible (vérifiez auprès de votre administrateur) dans votre espace de travail et que vous spécifiez remoteTrue dans le @distributed décorateur, la charge de travail est lancée sur le pool GPU réservé par défaut. Si vous souhaitez utiliser le pool GPU à la demande, définissez la variable DATABRICKS_USE_RESERVED_GPU_POOLFalse d’environnement avant d’appeler la fonction distribuée, comme indiqué ci-dessous :

import os
os.environ['DATABRICKS_USE_RESERVED_GPU_POOL'] = 'False'
@distributed(gpus=8, remote=True)
def run_train():
    ...

Learn more

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

Commentaires

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

Last updated on 2025-12-19