Entrenamiento distribuido en cuadernos

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 Administrar versiones preliminares de Azure Databricks.

El decorador @distributed de la API de Python para GPU sin servidor es la forma más cómoda de ejecutar el entrenamiento distribuido desde un cuaderno de Databricks. Decora tu función de entrenamiento, llámala, y AI Runtime la ejecuta en todas las GPUs del nodo al que está conectado tu portátil. El mismo código escala de una sola GPU a multi-GPU sin clúster que provisionar ni lanzador distribuido que configurar.

Tip

• El @distributed decorador ejecuta una función de entrenamiento en cada GPU de tu nodo desde dentro de un cuaderno.
• Soporta PyTorch DDP, FSDP y DeepSpeed, y traslada código de una sola GPU a multi-GPU con cambios mínimos.
• Conecta tu portátil a un acelerador 8xH100 y pon en marcha gpus=8 para entrenamiento completo con múltiples GPU.

Inicio rápido

El serverless_gpu paquete se preinstala cuando tu portátil está conectado a una GPU serverless. Decora tu función de entrenamiento con @distributed, y luego llámala con .distributed():

from serverless_gpu import distributed

# gpus is the number of GPUs on the node. gpu_type is optional and
# auto-detected from the accelerator your notebook is connected to.
@distributed(gpus=8, gpu_type="H100")
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()

Cada llamada a .distributed() crea una ejecución de MLflow (o una ejecución secundaria anidada si ya hay una activa) e imprime un enlace a la ejecución en la salida de la celda. Para una guía completa y ejecutable, véase el ejemplo completo.

Marcos admitidos

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

• PyTorch Distributed Data Parallel (DDP): paralelismo de datos estándar en varias GPU.
• Paralelismo de datos completamente fragmentado (FSDP): entrenamiento con uso eficiente de la memoria para modelos grandes.
• DeepSpeed: Microsoft biblioteca de optimización para el entrenamiento de modelos de gran tamaño.

Para ver escenarios de entrenamiento reales en los que se utiliza cada biblioteca, consulte los ejemplos de notebooks.

Cómo trabaja el @distributed decorador

Cuando llamas a una función decorada con .distributed(), AI Runtime gestiona las mecánicas que de otro modo configurarías a mano con un lanzador distribuido:

• Serialización y expansión: La función se serializa y se inicia en cada uno de los gpus que solicite. Cada GPU ejecuta una copia de la función con los mismos argumentos.
• Sincronización del entorno: El entorno Python y sus dependencias se replican en todos los rangos, por lo que todos los procesos ejecutan el mismo código.
• Variables de entorno de rango: las variables estándar, como LOCAL_RANK, se asignan a cada proceso. Léelos dentro de tu función para colocar el modelo y los datos en el dispositivo correcto.
• Recogida de resultados: Los valores de retorno se recogen de todos los rangos y se devuelven al llamante.
• Seguimiento de MLflow: Cada llamada a .distributed() crea una ejecución de MLflow, o una ejecución hija anidada si ya hay una activa, por lo que las métricas registradas desde tu función se registran en la misma ejecución.
• Ciclo de vida y tiempo de espera: La ejecución distribuida se ejecuta dentro del ciclo de vida del portátil. Finalizar el notebook finaliza la ejecución. El decorador tiene un tiempo de espera predeterminado de 3 horas. Pasa timeout en segundos para cambiarlo o timeout=None para desactivarlo. Los tiempos de espera personalizados requieren entorno GPU v5 o superior.

La API se basa en las bibliotecas estándar de PyTorch: Distributed Data Parallel (DDP), Fully Sharded Data Parallel (FSDP) y DeepSpeed.

Proviene de TorchDistributor

Si hoy ejecutas PyTorch distribuido en Spark con TorchDistributor y tu carga de trabajo cabe en un solo nodo, la serverless_gpu@distributed API es el reemplazo recomendado para nuevas cargas de deep learning. Elimina el clúster Spark y te da la misma ruta de código de una sola GPU a multi-GPU.

Feature	serverless_gpu @distributed API	TorchDistributor
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 (UCVolumeDataset para transmitir datos de archivos)	Mediante Spark o sistema de archivos

Para migrar una carga de trabajo de un solo nodo:

• Sustituye la llamada a TorchDistributor(...).run(train_fn, ...) por el decorador @distributed en train_fn y luego ejecuta con train_fn.distributed(...).
• Elimina el clúster de Spark y la configuración del nodo de trabajo con GPU. Conecta tu portátil a un acelerador 8xH100 y establece gpus=8 en su lugar.
• Mueve la carga de datos dentro de la función decorada. Consulta Carga de datos.
• Conserva tu código de modelo DDP, FSDP o DeepSpeed existente. El decorador sostiene los tres.

@distributed funciona en un solo nodo (véase Limitaciones), por lo que no reemplaza todas las cargas de trabajo de TorchDistributor. Mantén las cargas de trabajo que dependen de la integración de Spark en TorchDistributor. Para ejecutar entrenamiento distribuido desde tu máquina local o entre varios nodos, usa la CLI de Tiempo de Ejecución de IA, que está en Vista Previa Pública. Consulte AI Runtime CLI.

Ejemplo completo

El siguiente ejemplo entrena un modelo de perceptrón multicapa (MLP) en 8 GPUs H100 desde un portátil.

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():
       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. Importe la serverless_gpu biblioteca y el distributed módulo.

   import serverless_gpu
   from serverless_gpu import distributed
   

3. Envuelve el código de entrenamiento del modelo en una función y decora dicha función con el decorador @distributed. La función decorada es el punto de entrada para la ejecución distribuida, así que define toda la lógica de entrenamiento, carga de datos e inicialización de modelos dentro de ella.

   @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. Ejecuta 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. Al ejecutarlo, se genera un enlace 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. Para más información sobre la personalización de nombres de experimentos, el seguimiento de métricas y la reanudación de ejecuciones, consulte Seguimiento de experimentos y observabilidad.

Carga de datos

Coloca el código de carga de datos dentro de la @distributed función. Un conjunto de datos puede superar el tamaño máximo permitido por pickle, por lo que generarlo o cargarlo dentro del decorador evita errores de serialización:

from serverless_gpu import distributed

# This may cause a pickle error because the dataset is captured by the function.
dataset = get_dataset(file_path)

@distributed(gpus=8, gpu_type='H100')
def run_train():
    # Load the dataset inside the decorated function instead.
    dataset = get_dataset(file_path)
    ...

Para datos basados en archivos almacenados en volúmenes de Unity Catalog, use UCVolumeDataset de serverless_gpu.data, que transmite los archivos con almacenamiento en caché local y los divide automáticamente entre rangos y trabajadores. Para crear un punto de control del entrenamiento distribuido en un volumen, use UCVolumeWriter y UCVolumeReader. Consulte Carga de datos en tiempo de ejecución de IA y Punto de control de modelos.

Limitations

• El entrenamiento distribuido se ejecuta entre las GPUs en el único nodo al que está conectado tu portátil. Para un entrenamiento completo multi-GPU, conéctate a un acelerador 8xH100, que proporciona un nodo con 8 GPUs, y configura gpus=8.
• El tipo de acelerador debe coincidir. Si estableces gpu_type en @distributed, debe coincidir con el acelerador al que está conectado tu notebook ("H100" o "A10"). Una descoordinación hace que la carga de trabajo falle. El parámetro es opcional y se detecta automáticamente cuando se omite.
• AI Runtime recomienda el entorno GPU v4 y superiores. Los tiempos de espera personalizados (el timeout parámetro) requieren entorno GPU v5 o superior.
• El decorador se agota después de 3 horas por defecto. Pasa timeout en segundos para cambiarlo o timeout=None para desactivarlo.
• La ejecución se ejecuta dentro del ciclo de vida del cuaderno. Finalizar el notebook finaliza la ejecución.

Aprende más

• Para el decorador @distributed, GPUType y las API de Ray, consulta la documentación de referencia de la API de Python para Serverless GPU.
• Para patrones que hagan tu pipeline de entrenamiento más eficiente y resiliente, consulta la guía de rendimiento y resiliencia.
• Para escenarios de entrenamiento de extremo a extremo, consulta los ejemplos de cuadernos.