Entrenamiento distribuido en cuadernos

Importante

Esta característica se encuentra en su versión beta.

El @distributed decorador del SDK de Runtime de IA de Databricks (databricks-sdk-air) es la forma más cómoda de ejecutar 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.

Puntos clave en esta página:

• 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 para un entrenamiento completo con múltiples GPU. Por defecto, @distributed usa todas las GPUs del nodo.

Note

Esta página cubre el entrenamiento distribuido desde notebooks de Databricks con Databricks AI Runtime SDK. Para enviar cargas de trabajo de entrenamiento distribuidas desde tu máquina local, utiliza los comandos de la CLI de Databricks para AI Runtime, que están en Vista previa pública. Consulta Usar la CLI de Databricks con AI Runtime.

Inicio rápido

Los ejemplos de esta página utilizan el SDK de Runtime de IA de Databricks, que está preinstalado en el entorno GPU versión 6 y superior. En la versión 5 y anteriores, usa el serverless_gpu paquete en su lugar. Véase Migrar desde serverless_gpu. Para conectar un notebook a un acelerador, consulta Conectar a un notebook.

Decora tu función de entrenamiento con @distributed, y luego llámala 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()

El entrenamiento distribuido requiere un acelerador de 8xH100, que aprovisiona un solo nodo con 8 GPU. Por defecto, @distributed usa las ocho GPUs. Establezca num_accelerators para usar menos. El parámetro accelerator_type también es opcional y se detecta automáticamente a partir del acelerador al que está conectado su cuaderno.

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 escenarios reales de entrenamiento que usan cada biblioteca, véanse ejemplos de entrenamiento distribuido.

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 dispersión: La función se serializa y se lanza en cada GPU seleccionada. 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.

Migrar desde serverless_gpu

El serverless_gpu paquete está obsoleto y ya no recibe nuevas funciones. A partir de la versión 6 del entorno GPU, el SDK de Runtime de IA de Databricks (databricks-sdk-air) lo reemplaza y proporciona APIs equivalentes, con argumentos renombrados, bajo el databricks.air espacio de nombres. El paquete que utilices depende de la versión del entorno de la GPU:

• Versión 5 e inferiores: Incluir solo serverless_gpu. Sigue usando serverless_gpu en estas versiones.
• Versión 6: Incluye ambos paquetes, pero solo el SDK de Tiempo de Ejecución de IA recibe actualizaciones. La versión 6 es la última versión del entorno GPU que incluye serverless_gpu.

El SDK de tiempo de ejecución de IA es un paquete separado del SDK de Databricks para Python (databricks-sdk).

	serverless_gpu (en desuso)	SDK de Tiempo de Ejecución de IA (databricks.air)
Versiones del entorno GPU	6 o menos	6 en adelante
Importar	from serverless_gpu import distributed, ray_init	from databricks.air import distributed, ray_init
Argumento de recuento de GPU	gpus	num_accelerators
Argumento de tipo de GPU	gpu_type (GPUType)	accelerator_type (AcceleratorType)
Valores del tipo de GPU	"A10", "H100", "B300"	"GPU_1xA10", "GPU_1xH100", , "GPU_8xH100", "GPU_8xB300"
Módulo de carga de datos y puntos de control	serverless_gpu.data	databricks.air.data. Consulte Carga de datos en el entorno de ejecución de IA.

Para migrar un cuaderno:

1. Conecta el cuaderno al entorno GPU versión 6 o superior, que incluye el SDK de AI Runtime.
2. Cambiar importaciones de serverless_gpu a databricks.air.
3. Renombra gpus a num_accelerators y gpu_type a accelerator_type, y usa los nuevos valores de tipo acelerador. Los nombres de argumentos antiguos y los valores de tipo de GPU siguen funcionando en databricks.air, pero emiten un FutureWarning. La compatibilidad finalizará en la versión 7 del entorno de GPU.

Tip

Indique a Genie Code que lo haga por usted:

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

Para cargas de trabajo de la CLI de Databricks que usen databricks.air.data, agregue databricks-sdk-air[data] al environment.dependencies de la carga de trabajo.

Proviene de TorchDistributor

Si ejecutas entrenamiento distribuido usando TorchDistributor en clústeres de Spark y tu carga de trabajo cabe en un solo nodo, la @distributed API es el reemplazo recomendado en AI Runtime. Tu código de modelo DDP, FSDP o DeepSpeed no cambia.

Feature	@distributed Interfaz de Programación de Aplicaciones (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, se utilizan volúmenes de Unity Catalog (UCVolumeDataset para la transmisión de 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 en su lugar. @distributed usa las ocho GPUs por defecto.
• 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 es compatible con 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. Consulta Usar la CLI de Databricks con AI Runtime.

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. Importa el decorador distributed desde el AI Runtime SDK.

   from databricks.air 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()
   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. 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 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)
    ...

Para datos basados en archivos almacenados en volúmenes de Unity Catalog, use UCVolumeDataset de databricks.air.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. Consulta Cargar datos en AI Runtime y Checkpoint con Distributed Checkpoint (DCP).

Limitations

• @distributed Solo soporta un único nodo: se ejecuta a través de las GPUs del nodo al que está conectado tu notebook. Para entrenar en varios nodos, utiliza la CLI de AI Runtime Databricks. Consulta Usar la CLI de Databricks con AI Runtime.
• Si configuras accelerator_type, debe coincidir con el acelerador al que está conectado tu portátil (por ejemplo, "GPU_8xH100"). Una discrepancia provoca que la carga de trabajo falle. El parámetro es opcional y se detecta automáticamente cuando se omite.
• AI Runtime soporta el entorno GPU v4 y superiores. Los tiempos de espera personalizados (el timeout parámetro) requieren entorno GPU v5 o superior.
• El decorador caduca tras 3 horas de forma predeterminada. 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.

Recursos adicionales

• Para el decorador @distributed, AcceleratorType y las API de Ray, consulta la referencia de la API del SDK de Databricks AI Runtime.
• 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, véanse ejemplos de entrenamiento distribuido.