Compartir a través de

Carga de trabajo de varias GPU y varios nodos

Puede iniciar la carga de trabajo distribuida en varias GPU (ya sea dentro de un solo nodo o en varios nodos) 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 entre GPU remotas desde el mismo cuaderno.

Inicio rápido

La API de GPU sin servidor para el entrenamiento distribuido está preinstalada en los entornos de cálculo de GPU sin servidor para los cuadernos 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.

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, the GPU type,
# and whether or not the GPUs are remote
@distributed(gpus=8, gpu_type='A10', remote=True)
def run_train():
    ...

A continuación se muestra un ejemplo completo que entrena un modelo de perceptron multicapa (MLP) en 8 nodos de GPU A10 desde un cuaderno:

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

  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.

    Salida en la celda del cuaderno

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 todos los nodos.
  • Los resultados se recopilan y devuelven de todas las GPU

Si remote se configura como True, la carga de trabajo se distribuye en las GPU remotas. Si se establece Falseen remote, la carga de trabajo se ejecuta en el único nodo de GPU conectado a la libreta actual. Si el nodo tiene varios chips de GPU, se usarán todos ellos.

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

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

Inicio con Ray

La API de GPU sin servidor también admite el inicio de entrenamiento distribuido con Ray empleando el decorador @ray_launch, que se superpone a @distributed. Cada ray_launch tarea arranca primero una cita distribuida por antorchas para decidir al trabajador principal de Ray y reunir direcciones IP. El rango cero comienza ray start --head (con la exportación de métricas si está habilitada), establece RAY_ADDRESSy ejecuta la función decorada como controlador Ray. Otros nodos se unen a través de ray start --address y esperan hasta que el controlador escriba un marcador de finalización.

Detalles de configuración adicionales:

  • Para habilitar la recopilación de métricas del sistema Ray en cada nodo, use RayMetricsMonitor con remote=True.
  • Defina las opciones de tiempo de ejecución de Ray (actores, conjuntos de datos, grupos de ubicación y programación) dentro de la función decorada mediante la API estándar de Ray.
  • Administrar controles en todo el clúster (recuento y tipo de GPU, modo remoto o local, comportamiento asincrónico y variables de entorno del grupo de Databricks) fuera de la función en los argumentos del decorador o en el entorno del notebook.

En el ejemplo siguiente se muestra cómo usar @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()

Para obtener un ejemplo completo, consulte este notebook, que inicia Ray para entrenar una red neuronal ResNet18 en varias tarjetas GPU A10.

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, remote=True)
def run_train():
  # good practice
  dataset = get_dataset(file_path)
  ....

¿Puedo usar grupos de GPU reservados?

Si el grupo de GPU reservado está disponible (consulte con el administrador) en el área de trabajo y especifique remote en True el @distributed decorador, la carga de trabajo se iniciará en el grupo de GPU reservado de forma predeterminada. Si desea usar el grupo de GPU a petición, establezca la variable DATABRICKS_USE_RESERVED_GPU_POOL de entorno en False antes de llamar a la función distribuida, como se muestra a continuación:

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

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