Compartir a través de

Carga de trabajo de varias GPU

Importante

Esta característica se encuentra en su versión beta. Los administradores del área de trabajo pueden controlar el acceso a esta característica desde la página Vistas previas . Consulte Administración de versiones preliminares de Azure Databricks.

Puede iniciar cargas de trabajo distribuidas entre varias GPU en un solo nodo mediante la API de Python de GPU sin servidor. La API proporciona una interfaz sencilla y unificada que abstrae los detalles del aprovisionamiento de GPU, la configuración del entorno y la distribución de cargas de trabajo. Con cambios mínimos en el código, puede pasar sin problemas del entrenamiento de una sola GPU a la ejecución distribuida de varias GPU desde el mismo cuaderno.

Marcos admitidos

La @distributed API se integra con las principales bibliotecas de entrenamiento distribuidas:

  • PyTorch Distributed Data Parallel (DDP) — El paralelismo de datos multi-GPU estándar.
  • Paralelismo de Datos Totalmente Fragmentado (FSDP): entrenamiento eficiente en cuanto a memoria para modelos de gran tamaño.
  • DeepSpeed : biblioteca de optimización de Microsoft para el entrenamiento de modelos de gran tamaño.

API de serverless_gpu frente a TorchDistributor

En la tabla siguiente se compara la serverless_gpu@distributed API con TorchDistributor:

Feature serverless_gpu @distributed API Distribuidor de Antorchas
Infraestructura Totalmente sin servidor, sin administración de clústeres Requiere un clúster de Spark con trabajadores GPU
Configuración Decorador único, configuración mínima Requiere el clúster de Spark y la configuración de TorchDistributor
Soporte de marco PyTorch DDP, FSDP, DeepSpeed Principalmente PyTorch DDP
Carga de datos Dentro del decorador, utiliza Volúmenes del Catálogo de Unity Mediante Spark o sistema de archivos

La serverless_gpu API es el enfoque recomendado para las nuevas cargas de trabajo de aprendizaje profundo en Databricks. TorchDistributor sigue estando disponible para cargas de trabajo estrechamente acopladas con clústeres de Spark.

Inicio rápido

La API de GPU sin servidor para el entrenamiento distribuido está preinstalada cuando te conectas a una GPU sin servidor dentro de los cuadernos y trabajos de Databricks. Se recomienda el entorno de GPU 4 y versiones posteriores. Para usarlo para el entrenamiento distribuido, importe y use el decorador distributed para distribuir su función de entrenamiento.

Envuelva el código de entrenamiento del modelo en una función y decore la función con el decorador @distributed. La función decorada se convierte en el punto de entrada para la ejecución distribuida; toda la lógica de entrenamiento, la carga de datos y la inicialización del modelo deben definirse dentro de esta función.

Advertencia

El parámetro gpu_type en @distributed debe coincidir con el tipo de acelerador al que está conectado el notebook. Por ejemplo, @distributed(gpus=8, gpu_type='H100') requiere que tu portátil esté conectado a un acelerador H100. El uso de un tipo de acelerador incompatible (por ejemplo, conectarse a A10 al especificar H100) hará que falle la carga de trabajo.

El fragmento de código siguiente muestra el uso básico de @distributed:

# Import the distributed decorator
from serverless_gpu import distributed

# Decorate your training function with @distributed and specify the number of GPUs and GPU type
@distributed(gpus=8, gpu_type='H100')
def run_train():
    ...

A continuación se muestra un ejemplo completo que entrena un modelo de perceptron multicapa (MLP) en 8 GPUs H100 desde un notebook.

  1. Configure el modelo y defina las funciones de utilidad.

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

  2. Importe la biblioteca serverless_gpu y el módulo distribuido .

    import serverless_gpu
from serverless_gpu import distributed

  3. Envuelva el código de entrenamiento del modelo en una función y decore la función con el decorador @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. Ejecute el entrenamiento distribuido llamando a la función distribuida con argumentos definidos por el usuario.

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

  5. Cuando se ejecuta, se genera un vínculo de ejecución de MLflow en la salida de la celda del cuaderno. Haga clic en el vínculo Ejecutar de MLflow o busque en el panel Experimento para ver los resultados de la ejecución.

Detalles de ejecución distribuida

La API de GPU sin servidor consta de varios componentes clave:

  • Administrador de computación: gestiona la asignación y administración de recursos.
  • Entorno en tiempo de ejecución: administra entornos y dependencias de Python.
  • Lanzador: orquesta la ejecución y la supervisión de las tareas

Cuando se ejecuta en modo distribuido:

  • La función se serializa y distribuye entre el número especificado de GPU.
  • Cada GPU ejecuta una copia de la función con los mismos parámetros
  • El entorno se sincroniza en todas las GPU.
  • Los resultados se recopilan y devuelven de todas las GPU

La API admite bibliotecas de entrenamiento paralelas populares, como Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP), DeepSpeed.

Puede encontrar escenarios de entrenamiento distribuidos más reales mediante las distintas bibliotecas de ejemplos de cuadernos.

FAQs

¿Dónde debe colocarse el código de carga de datos?

Al usar la API sin servidor de GPU para el entrenamiento distribuido, mueva el código de carga de datos dentro del decorador @distributed. El tamaño del conjunto de datos puede superar el tamaño máximo permitido por pickle, por lo que se recomienda generar el conjunto de datos dentro del decorador, como se muestra a continuación:

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

Aprende más

Para obtener la referencia de API, consulte la documentación de la API de Python de GPU sin servidor .

  • Last updated on