Entrenamiento distribuido con TorchDistributor
En este artículo se describe cómo realizar el entrenamiento distribuido en modelos de PyTorch ML mediante TorchDistributor.
TorchDistributor es un módulo de código abierto en PySpark que ayuda a los usuarios a hacer el entrenamiento distribuido con PyTorch en sus clústeres de Spark, por lo que permite iniciar trabajos de entrenamiento de PyTorch como trabajos de Spark. En segundo plano, inicializa el entorno y los canales de comunicación entre los trabajadores y utiliza el comando torch.distributed.run de la CLI para ejecutar el entrenamiento distribuido en los nodos de trabajo.
La API de TorchDistributor admite los métodos que se muestran en la tabla siguiente.
| Método y firma | Descripción |
|---|---|
init(self, num_processes, local_mode, use_gpu) |
Cree una instancia de TorchDistributor. |
run(self, main, *args) |
Ejecuta el entrenamiento distribuido invocando main(**kwargs) si main es una función y ejecuta el comando torchrun main *args de la CLI si main es una ruta de acceso de archivo. |
Requisitos
- Spark 3.4
- Databricks Runtime 13.0 ML o superior
Flujo de trabajo de desarrollo para cuadernos
Si el proceso de creación y entrenamiento del modelo se produce por completo desde un cuaderno en la máquina local o en Databricks Notebook, solo tiene que realizar pequeños cambios para preparar el código para el entrenamiento distribuido.
Preparar el código de nodo único: Prepare y pruebe el código de nodo único con PyTorch, PyTorch Lightning u otros marcos basados en PyTorch o PyTorch Lightning, como la API de HuggingFace Trainer.
Preparación del código para el entrenamiento distribuido estándar: Debe convertir el entrenamiento de proceso único en el entrenamiento distribuido. Haga que este código distribuido esté incluido en una función de entrenamiento que pueda usar con
TorchDistributor.Mover las importaciones dentro de la función de entrenamiento: Agregue las importaciones necesarias, como
import torch, dentro de la función de entrenamiento. Esto le permite evitar errores comunes de selección. Además, losdevice_ida los que se vinculan los modelos y los datos están determinados por:device_id = int(os.environ["LOCAL_RANK"])Iniciar el entrenamiento distribuido: Cree una instancia de
TorchDistributorcon los parámetros deseados y llame a.run(*args)para iniciar el entrenamiento.
A continuación se muestra un ejemplo de código de entrenamiento:
from pyspark.ml.torch.distributor import TorchDistributor
def train(learning_rate, use_gpu):
import torch
import torch.distributed as dist
import torch.nn.parallel.DistributedDataParallel as DDP
from torch.utils.data import DistributedSampler, DataLoader
backend = "nccl" if use_gpu else "gloo"
dist.init_process_group(backend)
device = int(os.environ["LOCAL_RANK"]) if use_gpu else "cpu"
model = DDP(createModel(), **kwargs)
sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, sampler=sampler)
output = train(model, loader, learning_rate)
dist.cleanup()
return output
distributor = TorchDistributor(num_processes=2, local_mode=False, use_gpu=True)
distributor.run(train, 1e-3, True)
Migración del entrenamiento desde repositorios externos
Si tiene un procedimiento de entrenamiento distribuido existente almacenado en un repositorio externo, puede migrar fácilmente a Azure Databricks haciendo lo siguiente:
- Importar el repositorio: Importar el repositorio externo como Carpeta Git de Databricks.
- Creación de un cuaderno: Inicialice un nuevo Azure Databricks Notebook en el repositorio.
- Inicio del entrenamiento distribuido: En una celda del cuaderno, llame a
TorchDistributorcomo se muestra:
from pyspark.ml.torch.distributor import TorchDistributor
train_file = "/path/to/train.py"
args = ["--learning_rate=0.001", "--batch_size=16"]
distributor = TorchDistributor(num_processes=2, local_mode=False, use_gpu=True)
distributor.run(train_file, *args)
Solución de problemas
Un error común para el flujo de trabajo del cuaderno es que los objetos no se pueden encontrar ni seleccionar al ejecutar el entrenamiento distribuido. Esto sucede cuando las instrucciones de importación de la biblioteca no se distribuyen a otros ejecutores.
Para evitar este problema, incluya todas las instrucciones de importación (por ejemplo, import torch) en la parte superior del método de entrenamiento al que se llama con TorchDistributor(...).run(<func>)y también dentro de cualquier otra función definida por el usuario que se haya llamado en el método de entrenamiento.
Error de NCCL: ncclInternalError: Internal check failed.
Cuando se produce este error durante el entrenamiento de varios nodos, normalmente indica un problema con la comunicación de red entre GPU. Este problema se produce cuando NCCL (biblioteca de comunicaciones colectiva de NVIDIA) no puede usar determinadas interfaces de red para la comunicación con GPU.
Para resolver este error, agregue el siguiente fragmento de código en el código de entrenamiento para usar la interfaz de red principal.
import os
os.environ["NCCL_SOCKET_IFNAME"] = "eth0"
Cuadernos de ejemplo
En los siguientes ejemplos de cuadernos se muestra cómo realizar el entrenamiento distribuido con PyTorch.