es parte de una serie sobre IA distribuida en múltiples GPU:
Parte 1: Comprender el paradigma del host y del dispositivo Parte 2: Operaciones colectivas y punto a punto (este artículo) Parte 3: Cómo se comunican las GPU (próximamente) Parte 4: Acumulación de gradiente y paralelismo de datos distribuidos (DDP) (próximamente) Parte 5: ZeRO (próximamente) Parte 6: Paralelismo tensorial (próximamente)
Introducción
En la publicación anterior, establecimos el paradigma del dispositivo host e introdujimos el concepto de rangos para cargas de trabajo de múltiples GPU. Ahora, exploraremos los patrones de comunicación específicos proporcionados por el módulo torch.distributed de PyTorch para coordinar el trabajo e intercambiar datos entre estos rangos. Estas operaciones, conocidas como colectivas, son los componentes básicos de las cargas de trabajo distribuidas.
Aunque PyTorch expone estas operaciones, en última instancia llama a un marco de backend que realmente implementa la comunicación. Para las GPU NVIDIA, es NCCL (Biblioteca de comunicaciones colectivas de NVIDIA), mientras que para AMD es RCCL (Biblioteca de comunicaciones colectivas de ROCm).
NCCL implementa primitivas de comunicación multi-GPU y multi-nodo optimizadas para GPU y redes NVIDIA. Detecta automáticamente la topología actual (canales de comunicación como PCIe, NVLink, InfiniBand) y selecciona el más eficiente.
Descargo de responsabilidad 1: dado que las GPU NVIDIA son las más comunes, nos centraremos en el backend de NCCL para esta publicación.
Descargo de responsabilidad 2: por motivos de brevedad, el código que se presenta a continuación solo proporciona los argumentos principales de cada método en lugar de todos los argumentos disponibles.
Descargo de responsabilidad 3: Para simplificar, no mostramos la desasignación de memoria de los tensores, pero operaciones como la dispersión no liberarán automáticamente la memoria del rango de origen (si no entiendes lo que quiero decir, está bien, quedará claro muy pronto).
Comunicación: bloqueo versus no bloqueo
Para trabajar juntas, las GPU deben intercambiar datos. La CPU inicia la comunicación poniendo en cola los núcleos NCCL en secuencias CUDA (si no sabe qué son las secuencias CUDA, consulte la primera publicación de blog de esta serie), pero la transferencia de datos real ocurre directamente entre las GPU a través de la interconexión, sin pasar por la memoria principal de la CPU. Lo ideal es que las GPU estén conectadas con una interconexión de alta velocidad como NVLink o InfiniBand (estas interconexiones se tratan en la tercera publicación de esta serie).
Esta comunicación puede ser sincrónica (con bloqueo) o asincrónica (sin bloqueo), lo cual exploramos a continuación.
Comunicación síncrona (bloqueo)
Comportamiento: cuando llama a un método de comunicación síncrono, el proceso del host se detiene y espera hasta que el kernel NCCL se ponga en cola con éxito en la secuencia CUDA activa actual. Una vez en cola, la función regresa. Esto suele ser sencillo y fiable. Tenga en cuenta que el host no está esperando a que se complete la transferencia, solo a que la operación se ponga en cola. Sin embargo, bloquea esa secuencia específica para que no pase a la siguiente operación hasta que el kernel NCCL se ejecute hasta su finalización.
Comunicación asincrónica (sin bloqueo)
Comportamiento: cuando llama a un método de comunicación asincrónica, la llamada regresa inmediatamente y la operación de puesta en cola ocurre en segundo plano. No se pone en cola en la secuencia activa actual, sino en una secuencia NCCL interna dedicada por dispositivo. Esto permite que su CPU continúe con otras tareas, una técnica conocida como cálculo superpuesto con comunicación. La API asincrónica es más compleja porque puede generar un comportamiento indefinido si no usa correctamente .wait() (que se explica a continuación) y modifica los datos mientras se transfieren. Sin embargo, dominarlo es clave para desbloquear el máximo rendimiento en la capacitación distribuida a gran escala.
Punto a punto (uno a uno)
Estas operaciones no se consideran colectivas, pero son primitivas de comunicación fundamentales. Facilitan la transferencia directa de datos entre dos rangos específicos y son fundamentales para tareas en las que una GPU necesita enviar información específica a otra.
Síncrono (bloqueo): el proceso host espera a que la operación se ponga en cola en la secuencia CUDA antes de continuar. El kernel se pone en cola en la secuencia activa actual. torch.distributed.send(tensor, dst): envía un tensor a un rango de destino específico. torch.distributed.recv(tensor, src): recibe un tensor de un rango fuente. Al tensor receptor se le debe asignar previamente la forma y el tipo correctos. Asíncrono (sin bloqueo): el proceso del host inicia la operación de puesta en cola e inmediatamente continúa con otras tareas. El kernel se pone en cola en un flujo NCCL interno dedicado por dispositivo, lo que permite superponer la comunicación con el cálculo. Estas operaciones devuelven una solicitud (técnicamente un objeto de trabajo) que se puede utilizar para rastrear el estado de la cola. request = torch.distributed.isend(tensor, dst): inicia una operación de envío asincrónica. request = torch.distributed.irecv(tensor, src): inicia una operación de recepción asincrónica. request.wait(): bloquea el host solo hasta que la operación se haya puesto en cola con éxito en la secuencia CUDA. Sin embargo, bloquea la secuencia CUDA actualmente activa para que no ejecute núcleos posteriores hasta que se complete esta operación asincrónica específica. request.wait(timeout): si proporciona un argumento de tiempo de espera, el comportamiento del host cambia. Bloqueará el subproceso de la CPU hasta que se complete el trabajo de NCCL o se agote el tiempo de espera (generando una excepción). En casos normales, los usuarios no necesitan establecer el tiempo de espera. request.is_completed(): devuelve True si la operación se ha puesto en cola correctamente en una secuencia CUDA. Puede usarse para realizar encuestas. No garantiza que los datos reales hayan sido transferidos.
Cuando PyTorch inicia un kernel NCCL, automáticamente inserta una dependencia (es decir, fuerza una sincronización) entre su flujo activo actual y el flujo NCCL. Esto significa que la secuencia NCCL no comenzará hasta que finalice todo el trabajo previamente puesto en cola en la secuencia activa, lo que garantiza que el tensor que se envía ya contiene los valores finales.
De manera similar, llamar a req.wait() inserta una dependencia en la otra dirección. Cualquier trabajo que ponga en cola en la secuencia actual después de req.wait() no se ejecutará hasta que se complete la operación NCCL, por lo que puede usar de forma segura los tensores recibidos.
Principales “errores” en la NCCL
Si bien el envío y la recepción están etiquetados como "sincrónicos", su comportamiento en NCCL puede resultar confuso. Una llamada sincrónica en un tensor CUDA bloquea el subproceso de la CPU del host solo hasta que el núcleo de transferencia de datos se pone en cola en la secuencia, no hasta que se completa la transferencia de datos. Luego, la CPU queda libre para poner en cola otras tareas.
Hay una excepción: la primera llamada a torch.distributed.recv() en un proceso realmente bloquea y espera a que finalice la transferencia, probablemente debido a los procedimientos internos de preparación de NCCL. Las llamadas posteriores solo se bloquearán hasta que la operación esté en cola.
Considere este ejemplo donde el rango 1 se bloquea porque la CPU intenta acceder a un tensor que la GPU aún no ha recibido:
rango = torch.distributed.get_rank() si rango == 0: t = torch.tensor([1,2,3], dtype=torch.float32, dispositivo=dispositivo) # torch.distributed.send(t, dst=1) # No se realiza ninguna operación de envío en caso contrario: # rango == 1 (suponiendo solo 2 rangos) t = torch.empty(3, dtype=torch.float32, dispositivo=dispositivo) torch.distributed.recv(t, src=0) # Bloquea solo hasta que se pone en cola (después de la primera ejecución) print("Esto se imprimirá si NCCL se calienta") print print("Esto NO se imprimirá")
El proceso de la CPU en el rango 1 se atasca en la impresión
Si ejecuta este código varias veces, observe que Esto se imprimirá si NCCL se calienta y no se imprimirá en las ejecuciones posteriores, ya que la CPU todavía está bloqueada al imprimir.
Colectivos
Cada función de operación colectiva admite operaciones de sincronización y asincrónicas a través del argumento async_op. El valor predeterminado es Falso, es decir, operaciones sincrónicas.
Colectivos uno a todos
Estas operaciones implican que un rango envíe datos a todos los demás rangos del grupo.
Transmisión
torch.distributed.broadcast(tensor, src): copia un tensor de un rango de fuente única (src) a todos los demás rangos. Cada proceso termina con una copia idéntica del tensor. El parámetro tensor tiene dos propósitos: (1) cuando el rango del proceso coincide con el src, el tensor son los datos que se envían; (2) de lo contrario, se utiliza tensor para guardar los datos recibidos. rango = torch.distributed.get_rank() si rango == 0: # tensor de rango de origen = torch.tensor([1,2,3], dtype=torch.int64, dispositivo=dispositivo) else: # rangos de destino tensor = torch.empty(3, dtype=torch.int64, dispositivo=dispositivo) torch.distributed.broadcast(tensor, src=0)
Dispersión
torch.distributed.scatter (tensor, scatter_list, src): distribuye fragmentos de datos de un rango de origen en todos los rangos. La lista de dispersión en el rango de origen contiene múltiples tensores, y cada rango (incluido el origen) recibe un tensor de esta lista en su variable tensorial. Las clasificaciones de destino simplemente pasan Ninguna para scatter_list. # La lista de dispersión debe ser Ninguna para todos los rangos que no sean de origen. scatter_list = Ninguno si rango! = 0 else [torch.tensor([i, i+1]).to(dispositivo) para i en rango(0,4,2)] tensor = torch.empty(2, dtype=torch.int64).to(dispositivo) torch.distributed.scatter(tensor, scatter_list, src=0) print(f'Rango {rango} recibido: {tensor}')
Colectivos todo a uno
Estas operaciones recopilan datos de todos los rangos y los consolidan en un único rango de destino.
Reducir
torch.distributed.reduce(tensor, dst, op): toma un tensor de cada rango, aplica una operación de reducción (como SUM, MAX, MIN) y almacena el resultado final solo en el rango de destino (dst). rango = torch.distributed.get_rank() tensor = torch.tensor([rango+1, rango+2, rango+3], dispositivo=dispositivo) torch.distributed.reduce(tensor, dst=0, op=torch.distributed.ReduceOp.SUM) print(tensor)
Recolectar
torch.distributed.gather(tensor, together_list, dst): reúne un tensor de cada rango en una lista de tensores en el rango de destino. La lista_reunión debe ser una lista de tensores (con el tamaño y tipo correcto) en el destino y Ninguno en el resto. # La lista_reunión debe ser Ninguna para todos los rangos que no sean de destino. rango = torch.distributed.get_rank() world_size = torch.distributed.get_world_size() together_list = Ninguno si rango! = 0 else [torch.zeros(3, dtype=torch.int64).to(dispositivo) para _ en rango(world_size)] t = torch.tensor([0+rank, 1+rank, 2+rank], dtype=torch.int64).to(dispositivo) torch.distributed.gather(t, together_list, dst=0) print(f'Después de la operación, el rango {rank} tiene: {gather_list}')
La variable world_size es el número total de rangos. Se puede obtener con torch.distributed.get_world_size(). Pero no se preocupe por los detalles de implementación por ahora, lo más importante es comprender los conceptos.
Colectivos de todos a todos
En estas operaciones, cada rango envía y recibe datos de todos los demás rangos.
Todo Reducir
torch.distributed.all_reduce(tensor, op): Igual que reducir, pero el resultado se almacena en cada rango en lugar de en un solo destino. # Ejemplo de torch.distributed.all_reduce rango = torch.distributed.get_rank() tensor = torch.tensor([rango+1, rango+2, rango+3], dtype=torch.float32, dispositivo=dispositivo) torch.distributed.all_reduce(tensor, op=torch.distributed.ReduceOp.SUM) print(f"Rango {rango} después de all_reduce: {tensor}")
Todos se reúnen
torch.distributed.all_gather(tensor_list, tensor): Igual que recopilar, pero la lista recopilada de tensores está disponible en todos los rangos. # Ejemplo de torch.distributed.all_gather rango = torch.distributed.get_rank() world_size = torch.distributed.get_world_size() input_tensor = torch.tensor([rank], dtype=torch.float32, dispositivo=dispositivo) tensor_list = [torch.empty(1, dtype=torch.float32, dispositivo=dispositivo) para _ en rango(world_size)] torch.distributed.all_gather(tensor_list, input_tensor) print(f"Rango {rank} recopilado: {[t.item() for t in tensor_list]}")
Reducir la dispersión
torch.distributed.reduce_scatter(output, input_list): equivalente a realizar una operación de reducción en una lista de tensores y luego dispersar los resultados. Cada rango recibe una parte diferente de la producción reducida. # Ejemplo de torch.distributed.reduce_scatter rango = torch.distributed.get_rank() world_size = torch.distributed.get_world_size() input_list = [torch.tensor([rank + i], dtype=torch.float32, dispositivo=dispositivo) para i en rango(world_size)] salida = torch.empty(1, dtype=torch.float32, dispositivo=dispositivo) torch.distributed.reduce_scatter(output, input_list, op=torch.distributed.ReduceOp.SUM) print(f"Rango {rank} valor reducido recibido: {output.item()}")
Sincronización
Las dos operaciones más utilizadas son request.wait() y torch.cuda.synchronize(). Es crucial comprender la diferencia entre estos dos:
request.wait(): se utiliza para operaciones asincrónicas. Sincroniza el flujo CUDA actualmente activo para esa operación, asegurando que el flujo espere a que se complete la comunicación antes de continuar. En otras palabras, bloquea la transmisión CUDA actualmente activa hasta que finalice la transferencia de datos. En el lado del host, solo hace que el host espere hasta que el kernel esté en cola; el host no espera a que se complete la transferencia de datos. torch.cuda.synchronize(): este es un comando más contundente que pausa el subproceso de la CPU del host hasta que hayan finalizado todas las tareas previamente puestas en cola en la GPU. Garantiza que la GPU esté completamente inactiva antes de que la CPU funcione, pero puede crear cuellos de botella en el rendimiento si se usa incorrectamente. Siempre que necesite realizar mediciones comparativas, debe utilizar esto para asegurarse de capturar el momento exacto en que las GPU están listas.
Conclusión
¡Felicitaciones por llegar hasta el final! En esta publicación, aprendiste sobre:
Operaciones punto a punto Sincronización y asíncrono en NCCL Operaciones colectivas Métodos de sincronización
En la próxima publicación del blog, profundizaremos en PCIe, NVLink y otros mecanismos que permiten la comunicación en un entorno distribuido.