Training distribuito nei notebook

Il @distributed decorator dell'SDK Databricks AI Runtime (databricks-sdk-air) è il modo più comodo per eseguire l'addestramento distribuito da un notebook Databricks. Decora la tua funzione di addestramento, chiamala, e AI Runtime la esegue su tutte le GPU del nodo a cui è collegato il tuo notebook. Lo stesso codice scala da una sola GPU a multi-GPU senza cluster da fornire e senza launcher distribuito da configurare.

Punti chiave in questa pagina:

• Il decoratore @distributed esegue una funzione di addestramento su tutte le GPU del nodo all’interno di un notebook.
• Supporta PyTorch DDP, FSDP e DeepSpeed, e sposta il codice da singola GPU su multi-GPU con modifiche minime.
• Collega il tuo notebook a un acceleratore 8xH100 per un addestramento completo multi-GPU. Per impostazione predefinita, @distributed usa tutte le GPU del nodo.

Quickstart

Gli esempi in questa pagina utilizzano l'SDK Databricks AI Runtime, che è preinstallato nell'ambiente GPU versione 6 e superiore. Per la versione 5 e precedenti, usa invece il package serverless_gpu. Vedi Migrazione da serverless_gpu. Per collegare un notebook a un acceleratore, vedi Connessione a un notebook.

Decora la tua attività di allenamento con @distributed, poi chiamala con .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()

Il training distribuito richiede un acceleratore 8xH100, che effettua il provisioning di un singolo nodo con 8 GPU. Di default, @distributed usa tutte e otto le GPU. Impostate num_accelerators per usarne meno. Il parametro accelerator_type è opzionale e viene rilevato automaticamente dall'acceleratore a cui è collegato il tuo notebook.

Ogni chiamata a .distributed() crea una run di MLflow (o una run figlia annidata, se è già attiva una run) e stampa un collegamento alla run nell'output della cella. Per una guida completa ed eseguibile, vedi l'esempio completo.

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.

Per scenari di addestramento reali che utilizzano ogni libreria, vedi esempi di addestramento distribuito.

Come lavora il @distributed decoratore

Quando chiami una funzione decorata con .distributed(), AI Runtime gestisce le meccaniche che altrimenti configureresti manualmente con un launcher distribuito:

• Serializzazione e fan-out: La funzione viene serializzata ed eseguita su ogni GPU selezionata. Ogni GPU esegue una copia della funzione con gli stessi argomenti.
• Sincronizzazione dell'ambiente: L'ambiente Python e le dipendenze sono replicati su tutti i livelli, quindi ogni processo esegue lo stesso codice.
• Variabili di ambiente di rango: variabili standard come LOCAL_RANK quelle vengono popolate per ogni processo. Leggili nella tua funzione per posizionare il modello e i dati sul dispositivo corretto.
• Raccolta dei risultati: I valori restituiti vengono raccolti da tutti i ranghi e restituiti al chiamante.
• Tracciamento MLflow: Ogni .distributed() chiamata crea una run MLflow, o una run figlia annidata se già attiva, quindi le metriche registrate dalla tua funzione finiscono sulla stessa run.
• Ciclo di vita e timeout: L'esecuzione distribuita si svolge all'interno del ciclo di vita del notebook. Terminare il quaderno interrompe la run. Il decoratore ha un timeout predefinito di 3 ore. Passa timeout in pochi secondi per cambiarlo, o timeout=None per disattivarlo. I timeout personalizzati richiedono l'ambiente GPU v5 e superiore.

L'API si basa sulle librerie standard di PyTorch: Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP) e DeepSpeed.

Migra da serverless_gpu

Il serverless_gpu pacchetto è obsoleto e non riceve più nuove funzionalità. A partire dall'ambiente GPU versione 6, il Databricks AI Runtime SDK (databricks-sdk-air) lo sostituisce e fornisce API equivalenti, con argomenti rinominati, sotto il databricks.air namespace. Il pacchetto che usi dipende dalla versione dell'ambiente GPU:

• Versione 5 e inferiori: includere solo serverless_gpu. Continua a usare serverless_gpu su queste versioni.
• Versione 6: Include entrambi i pacchetti, ma solo l'AI Runtime SDK riceve aggiornamenti. La versione 6 è l'ultima versione dell'ambiente GPU che include serverless_gpu.

L'AI Runtime SDK è un pacchetto separato dall'SDK Databricks per Python (databricks-sdk).

Per migrare un notebook:

1. Collega il notebook all'ambiente GPU versione 6 o superiore, che include l'AI Runtime SDK.
2. Modifica le importazioni da serverless_gpu a databricks.air.
3. Rinomina gpus in num_accelerators e gpu_type in accelerator_type, e usa i nuovi valori di tipo acceleratore. I vecchi nomi degli argomenti e i valori dei tipi GPU funzionano ancora in databricks.air ma emettono un FutureWarning. Il supporto per la versione 7 dell'ambiente GPU terminerà.

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

Per i carichi di lavoro della CLI Databricks che utilizzano databricks.air.data, aggiungi databricks-sdk-air[data] al environment.dependencies del carico di lavoro.

Proveniente da TorchDistributor

Se esegui un addestramento distribuito usando TorchDistributor su cluster Spark e il tuo carico di lavoro si adatta a un singolo nodo, l'API @distributed è la sostituta raccomandata su AI Runtime. Il codice del tuo modello DDP, FSDP o DeepSpeed non cambia.

Per migrare un carico di lavoro a singolo nodo:

• Sostituisci la chiamata a TorchDistributor(...).run(train_fn, ...) con il decoratore @distributed in train_fn, quindi avvia con train_fn.distributed(...).
• Rimuovi il cluster Spark e la configurazione del worker GPU. Collega invece il tuo notebook a un acceleratore 8xH100. @distributed usa tutte e otto le GPU di default.
• Sposta il caricamento dei dati all'interno della funzione decorata. Vedi caricamento dati.
• Conserva il codice dei modelli DDP, FSDP o DeepSpeed esistente. Il decoratore sostiene tutti e tre.

Il @distributed decorator è in Public Preview e gira su un singolo nodo (vedi Limitazioni), quindi non sostituisce ogni carico di lavoro di TorchDistributor. Mantieni i carichi di lavoro che dipendono dall'integrazione con Spark su TorchDistributor. Per eseguire l'addestramento distribuito dalla tua macchina locale o su più nodi, usa i comandi della CLI Databricks per AI Runtime. Vedi Utilizzare la Databricks CLI con Runtime di AI.

Esempio completo

Il seguente esempio addestra un modello multilayer perceptron (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. Importa il distributed decorator dall'AI Runtime SDK.

   from databricks.air import distributed
   

3. Includere il codice di addestramento del modello in una funzione e decorare la funzione con il decoratore @distributed. La funzione decorata è il punto di ingresso per l'esecuzione distribuita, quindi definisci tutta la logica di addestramento, il caricamento dei dati e l'inizializzazione dei modelli al suo interno.

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

4. Esegui l'addestramento 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.

Caricamento dati

Inserisci il codice di caricamento dati all'interno della @distributed funzione. Un dataset può superare la dimensione massima consentita da pickle, quindi generandolo o caricandolo all'interno del decorator si evitano errori di serializzazione:

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

Per i dati basati su file archiviati nei volumi di Unity Catalog, usare UCVolumeDataset da databricks.air.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. Vedi Caricamento dati su AI Runtime e Checkpoint con Distributed Checkpoint (DCP).

Limitations

• @distributed Supporta solo un singolo nodo: viene eseguito sulle GPU del nodo a cui è collegato il tuo notebook. Per addestrare su più nodi, usa la CLI AI Runtime Databricks. Vedi Utilizzare la Databricks CLI con Runtime di AI.
• Se imposti accelerator_type, deve corrispondere all'acceleratore a cui il tuo notebook è collegato (ad esempio, "GPU_8xH100"). Una mancata corrispondenza provoca il fallimento del carico di lavoro. Il parametro è opzionale e viene rilevato automaticamente quando omesso.
• AI Runtime supporta l'ambiente GPU v4 e versioni successive. I timeout personalizzati (il timeout parametro) richiedono l'ambiente GPU v5 e superiore.
• Per impostazione predefinita, il decoratore va in timeout dopo 3 ore. Passa timeout in pochi secondi per cambiarlo, o timeout=None per disattivarlo.
• L'esecuzione si svolge all'interno del ciclo di vita del notebook. Terminare il quaderno interrompe la run.

Risorse aggiuntive

• Per le API @distributed decorator, AcceleratorType e Ray, vedi il riferimento all'API Databricks AI Runtime SDK.
• Per schemi che rendono la tua pipeline di allenamento più efficiente e resiliente, consulta la guida alla performance e alla resilienza.
• Per scenari di addestramento end-to-end, vedi esempi di addestramento distribuito.