Carico di lavoro multi-GPU

È possibile avviare carichi di lavoro distribuiti tra più GPU in un singolo nodo usando l'API Python GPU senza server <>c0. L'API fornisce un'interfaccia semplice e unificata che astrae i dettagli del provisioning gpu, della configurazione dell'ambiente e della distribuzione del carico di lavoro. Con modifiche minime al codice, è possibile passare facilmente dal training a GPU singola all'esecuzione distribuita su più GPU dallo stesso notebook.

Framework supportati

L'API @distributed si integra con le principali librerie di training distribuite:

• PyTorch Distributed Data Parallel (DDP): Parallelismo dei dati standard su più GPU.
• Fully Sharded Data Parallel (FSDP): addestramento efficiente in termini di memoria per modelli di grandi dimensioni.
• DeepSpeed: libreria di ottimizzazione di Microsoft per il training di modelli di grandi dimensioni.

API serverless_gpu vs TorchDistributor

La tabella seguente confronta l'API serverless_gpu@distributed con TorchDistributor:

L'API serverless_gpu è l'approccio consigliato per i nuovi carichi di lavoro di Deep Learning in Databricks. TorchDistributor rimane disponibile per i carichi di lavoro strettamente associati ai cluster Spark.

Avvio rapido

L'API serverless GPU per l'addestramento distribuito è preinstallata quando si è connessi a una GPU serverless nei notebook e nei processi di Databricks. È consigliabile usare l'ambiente GPU 4 e versioni successive. Per usarlo per il training distribuito, importa e usa il decoratore distributed per distribuire la funzione di training.

Includere il codice di addestramento del modello in una funzione e decorare la funzione con il decoratore @distributed. La funzione decorata diventa il punto di ingresso per l'esecuzione distribuita, pertanto tutte le logiche di training, il caricamento dei dati e l'inizializzazione del modello devono essere definite all'interno di questa funzione.

Per avviare l'esecuzione distribuita, chiamare la funzione decorata usando train_function.distributed(). Ogni chiamata crea automaticamente un'esecuzione dell'esperimento MLflow o un'esecuzione figlio annidata se ne è già attiva una.

Il frammento di codice seguente mostra l'utilizzo di base di @distributed:

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

Di seguito è riportato un esempio completo che esegue il training di un modello perceptron multilivello (MLP) su 8 GPU H100 da un notebook:

1. Configurare il modello e definire le funzioni di utilità.

   
   # 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. Importare la serverless_gpu libreria e il distributed modulo.

   import serverless_gpu
   from serverless_gpu import distributed
   

3. Includere il codice di addestramento del modello in una funzione e decorare la funzione con il decoratore @distributed.

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

4. Esegui il training distribuito chiamando la funzione distribuita con argomenti definiti dall'utente.

   run_train.distributed(num_epochs=3, batch_size=1)
   

5. Quando viene eseguito, nell'output della cella del notebook viene generato un collegamento di esecuzione MLflow. Fare clic sul collegamento MLflow run o trovarlo nel pannello Experiment per visualizzare i risultati dell'esecuzione. Per informazioni dettagliate sulla personalizzazione dei nomi degli esperimenti, il rilevamento delle metriche e la ripresa delle esecuzioni, vedere Rilevamento e osservabilità dell'esperimento.

Dettagli dell'esecuzione distribuita

L'API GPU serverless è costituita da diversi componenti chiave:

• Gestione calcolo: gestisce l'allocazione e la gestione delle risorse
• Ambiente di runtime: gestisce gli ambienti e le dipendenze Python
• Lanciatore: orchestra l'esecuzione e il monitoraggio dei lavori

Quando è in esecuzione in modalità distribuita:

• La funzione viene serializzata e distribuita nel numero specificato di GPU
• Ogni GPU esegue una copia della funzione con gli stessi parametri
• L'ambiente viene sincronizzato tra tutte le GPU
• I risultati vengono raccolti e restituiti da tutte le GPU
• Gestione del ciclo di vita: l'esecuzione distribuita viene eseguita all'interno del ciclo di vita del notebook. Quando il notebook termina, termina anche l'esecuzione. Il decoratore @distributed ha un timeout predefinito di 3 ore. Per impostare un timeout personalizzato, specificare timeout in secondi oppure timeout=None per disabilitarlo. L'impostazione di un timeout è disponibile nell'ambiente GPU v5 e versioni successive.

L'API supporta librerie di training parallele comuni, ad esempio Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP), DeepSpeed.

È possibile trovare scenari di training distribuiti più reali usando le varie librerie negli esempi di notebook.

Domande frequenti

Dove inserire il codice di caricamento dei dati?

Quando si utilizza l'API GPU serverless per l'addestramento distribuito, spostare il codice di caricamento dei dati all'interno del decoratore @distributed. Le dimensioni del dataset possono superare la dimensione massima consentita da pickle, pertanto è consigliabile generare il dataset all'interno del decoratore, come illustrato di seguito:

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

Per i dati basati su file archiviati nei volumi di Unity Catalog, usare UCVolumeDataset da serverless_gpu.data, che trasmette i file con memorizzazione nella cache locale e li partiziona automaticamente tra i ranghi e i ruoli di lavoro. Per eseguire il checkpoint di training distribuito a un volume, usare UCVolumeWriter e UCVolumeReader. Vedere Caricare i dati nel runtime di intelligenza artificiale e nel checkpoint del modello.

Ulteriori informazioni

Per la documentazione di riferimento dell'API, vedere la documentazione Serverless GPU Python API.