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.

Nota:

Esta página trata sobre el entrenamiento distribuido desde notebooks de Databricks con la API de Python para GPU sin servidor. Para enviar cargas de trabajo de entrenamiento distribuidas desde la máquina local, use la CLI del entorno de ejecución de IA, que se encuentra en versión preliminar pública. Consulte AI Runtime CLI.

Puede iniciar cargas de trabajo distribuidas entre varias GPU en un solo nodo mediante el Serverless GPU Python API. 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.

Nota:

El entrenamiento distribuido requiere un acelerador de 8xH100, que aprovisiona un solo nodo con 8 GPU. Al usar el @distributed decorador, establezca gpus=8. El gpu_type parámetro es opcional y se detecta automáticamente desde el acelerador al que está conectado el cuaderno.

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.

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	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	Decorador interno, utiliza Unity Catalog Volumes (`UCVolumeDataset` para la transmisión de datos de archivos)	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.

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 se convierte en el punto de entrada para la ejecución distribuida, por lo que toda la lógica de entrenamiento, la carga de datos y la inicialización del modelo deben definirse dentro de esta función.

Para iniciar la ejecución distribuida, llame a la función decorada mediante train_function.distributed(). Cada llamada crea automáticamente una ejecución de un experimento de MLflow o una ejecución secundaria anidada si ya hay una activa.

Advertencia

Si establece gpu_type en @distributed, asegúrese de que coincide con el tipo de acelerador al que está conectado el cuaderno ("H100" o "A10"). Si se especifica el tipo de acelerador incorrecto, se producirá un error en la carga de trabajo.

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

from serverless_gpu import distributed

# Decorate your training function with @distributed and specify the number of GPUs.
# gpu_type='H100' is optional and will be auto-detected if not set.
@distributed(gpus=8, gpu_type='H100')
def run_train():
    ...

run_train.distributed()

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

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)

Importe la serverless_gpu biblioteca y el distributed módulo.

import serverless_gpu
from serverless_gpu import distributed

Envuelve el código de entrenamiento del modelo en una función y decora dicha 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()

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

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
Administración del ciclo de vida: la ejecución distribuida se ejecuta dentro del ciclo de vida del cuaderno. Cuando el cuaderno finaliza, también lo hace la ejecución. El decorador @distributed tiene un tiempo de espera predeterminado de 3 horas. Para establecer un tiempo de espera personalizado, pasa timeout en segundos, o timeout=None para desactivarlo. La configuración de un tiempo de espera está disponible en el entorno de GPU v5 y versiones posteriores.

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

Puedes encontrar más escenarios reales de entrenamiento distribuido utilizando las diversas bibliotecas en 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)
  ....

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.

Aprende más

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

Comentarios

¿Le ha resultado útil esta página?

Last updated on 2026-07-17