Permite un entrenamiento más rápido con la biblioteca de paralelismo de datos de Amazon SageMaker

Optimiza tu entrenamiento con la potente biblioteca de paralelismo de datos de Amazon SageMaker

El entrenamiento de grandes modelos de lenguaje (LLM, por sus siglas en inglés) se ha vuelto cada vez más popular en el último año con el lanzamiento de varios modelos disponibles al público, como Llama2, Falcon y StarCoder. Ahora los clientes están entrenando LLMs de un tamaño sin precedentes que varía desde 1 billón hasta más de 175 mil millones de parámetros. El entrenamiento de estos LLMs requiere recursos informáticos significativos y tiempo, ya que cientos a miles de unidades de procesamiento gráfico (GPU, por sus siglas en inglés) deben usarse para manejar los vastos conjuntos de datos de entrenamiento y los tamaños del modelo de hoy en día. Un cuello de botella en el entrenamiento distribuido puede ser la comunicación entre GPU manejada por la Biblioteca de Comunicación Colectiva de NVIDIA (NCCL). En algunos trabajos de entrenamiento distribuido a gran escala, se puede gastar más tiempo en la comunicación entre GPU que en la computación real de las GPU. Para aliviar el cuello de botella de la comunicación de la GPU y permitir un entrenamiento más rápido, Amazon SageMaker se complace en anunciar una operación colectiva AllGather optimizada como parte de la biblioteca paralela distribuida de datos de SageMaker (SMDDP, por sus siglas en inglés). AllGather es la operación colectiva más utilizada en soluciones populares de paralelismo de datos eficientes en memoria, como DeepSpeed Zero Redundancy Optimizer (ZeRO) y Fully Sharded Data Parallelism (FSDP), y es el principal contribuyente a la sobrecarga de comunicación de la GPU. En esta publicación, mostramos una descripción general de alto nivel de cómo funciona SMDDP, cómo puedes habilitar SMDDP en tus scripts de entrenamiento de Amazon SageMaker y las mejoras de rendimiento que puedes esperar.

Descripción general de la solución

El entrenamiento paralelo de datos tradicional involucra replicar un modelo completo en múltiples GPUs, cada uno de los cuales entrena en diferentes fragmentos de datos del conjunto de datos. Durante el paso hacia atrás, los gradientes se promedian entre los trabajadores de GPU para que cada réplica del modelo se actualice con los mismos valores de gradiente a pesar de que se estén entrenando con fragmentos de datos diferentes. Esta técnica permite un entrenamiento mucho más rápido en vastos conjuntos de datos al paralelizar el consumo de datos de entrenamiento. Sin embargo, algunos de los modelos grandes de hoy en día (por ejemplo, Llama2 70B) son demasiado grandes para caber completamente en la memoria de la GPU, lo que hace que el paralelismo de datos tradicional no se pueda utilizar. Para seguir aprovechando los beneficios del paralelismo de datos y superar la limitación de memoria GPU, han ganado popularidad soluciones de paralelismo de datos fragmentados como DeepSpeed ZeRO, PyTorch FSDP y la biblioteca de paralelismo de modelos de SageMaker de Amazon.

En el paralelismo de datos fragmentados, en lugar de replicar todo el modelo en trabajadores de GPU, los parámetros del modelo, los gradientes y los estados del optimizador se dividen y se distribuyen (es decir, se dividen en fragmentos) en las GPU del trabajo de entrenamiento. Para realizar los cálculos de paso hacia adelante y hacia atrás, los parámetros se recopilan de los fragmentos en otras GPU para formar una o más capas del modelo. Después de realizar los cálculos, estas capas se liberan de la memoria para permitir que se recopilen las siguientes capas. Hay variantes del paralelismo de datos fragmentados en las que solo se dividen los estados del optimizador y los gradientes, pero no los parámetros del modelo. En este tipo de paralelismo de datos fragmentados, todavía se utiliza AllGather, pero solo antes de los cálculos de paso hacia adelante para recopilar los parámetros del modelo que han sido actualizados por diferentes fragmentos de gradiente o estados del optimizador de otras GPU. Consulta las diferentes etapas de DeepSpeed ZeRO y la estrategia de fragmentación SHARD_GRAD_OP de FSDP para obtener más detalles.

Se realiza una operación colectiva AllGather cada vez que se desfragmentan los parámetros, y NCCL proporciona la implementación de código abierto estándar de esta rutina. Como se muestra a continuación, cada trabajador de GPU involucre en el AllGather comienza con un búfer de entrada y termina con todos los búferes de entrada de otros trabajadores concatenados. Cuando se utiliza AllGather en paralelismo de datos fragmentados, los búferes de entrada contienen los fragmentos de parámetros del modelo y los búferes de salida grandes contienen una o más capas del modelo materializadas a partir de los otros fragmentos.

• Ciberseguridad protege la alimentación y la agricultura
• Utilizando LLMs para codificar nuevas tareas para los robots
• Enfoque de un Arquitecto Ejecutivo de FinOps Cómo la Inteligencia Artificial y la Automatización Simplifican la Gestión de Datos

Aunque NCCL se usa típicamente para AllGather en el entrenamiento distribuido, su implementación subyacente a nivel de bajo nivel no está adaptada a la infraestructura de red de las instancias de Amazon Elastic Compute Cloud (Amazon EC2), por lo que su rendimiento puede disminuir la velocidad del entrenamiento de extremo a extremo. La biblioteca SMDDP es una biblioteca de comunicación colectiva para las unidades de procesamiento gráfico (GPU) de NVIDIA que sirve como un reemplazo directo de NCCL y proporciona un mejor rendimiento para trabajos de entrenamiento distribuido con PyTorch. Específicamente, SMDDP ofrece una implementación optimizada de AllGather para los tipos de instancias p4d/p4de.

Dado que las operaciones colectivas como AllGather bloquean la computación de pase hacia adelante y hacia atrás, una ejecución más rápida de estas operaciones se traduce directamente en un tiempo de entrenamiento de extremo a extremo más corto sin efectos secundarios en la convergencia. Otras operaciones colectivas que se utilizan con menos frecuencia en el entrenamiento paralelo de datos fragmentados se manejan volviendo a NCCL.

Recorrido

AllGather optimizado para AWS

AllGather optimizado para AWS utiliza las siguientes técnicas para lograr un mejor rendimiento en la infraestructura de AWS en comparación con NCCL:

1. Movemos datos entre instancias a través de la red de Adaptador de Fibra Elástica (EFA) con un patrón de comunicación de todos a todos. EFA es la solución de red de baja latencia y alto rendimiento de AWS, y un patrón de todos a todos para la comunicación de red entre nodos es más específico para las características de EFA y la infraestructura de red de AWS, ya que requiere menos saltos de paquetes en comparación con el patrón de comunicación en anillo o en árbol de NCCL.
2. GDRCopy se utiliza para coordinar el tráfico local de NVLink y la red EFA. GDRCopy es una biblioteca que proporciona comunicación de baja latencia entre procesos de CPU y núcleos de GPU CUDA. Con esta tecnología, podemos optimizar el movimiento de datos dentro y entre nodos.
3. Se utiliza un número reducido de multiprocesadores de transmisión de GPU para devolver más potencia de cálculo a los núcleos del modelo. Las instancias de AWS P4d/P4de están equipadas con GPU NVIDIA A100, cada una de las cuales tiene 108 multiprocesadores de transmisión. Mientras que a NCCL le lleva hasta 24 multiprocesadores de transmisión ejecutar operaciones colectivas, las Colectivas SMDDP solo utilizan hasta nueve multiprocesadores de transmisión. Los multiprocesadores de transmisión ahorrados pueden ser aprovechados por los núcleos de cálculo del modelo para una ejecución más rápida.

Uso

Las colectivas de SMDDP se integran nativamente con PyTorch a través de la abstracción de grupo de procesos en el módulo torch.distributed. Un grupo de procesos define las interfaces para operaciones colectivas comunes como AllGather, ReduceScatter, AllReduce, etc. Los usuarios pueden escribir código distribuido genérico y luego elegir el backend subyacente, que proporciona la implementación de estas operaciones según el dispositivo de cálculo utilizado. Los trabajos de entrenamiento de CPU a menudo utilizan el backend gloo o mpi, mientras que las GPU de NVIDIA utilizan el backend nccl.

La biblioteca SMDDP se registra como un backend personalizado en la abstracción del grupo de procesos. Esto se hace mediante la importación de la declaración que se muestra en los siguientes fragmentos de código. Luego, al seleccionar el backend para su trabajo de entrenamiento distribuido basado en GPU, simplemente reemplace nccl por smddp. El backend smddp cumple con la misma semántica que el backend nccl y admite los mismos escenarios de entrenamiento.

DeepSpeed

import smdistributed.dataparallel.torch.torch_smddpdeepspeed.init_distributed(dist_backend="smddp")  # reemplazando "nccl"

FSDP

import smdistributed.dataparallel.torch.torch_smddpdist.init_process_group(backend="smddp")  # reemplazando "nccl"

Benchmarks

Hemos realizado una comparativa del rendimiento independiente de AllGather, donde la operación colectiva se ejecuta de forma aislada sin ningún entrenamiento de modelos. A continuación se muestra un resultado de muestra en 32 instancias p4d comparando NCCL y SMDDP AllGather. El eje X representa el tamaño de salida de AllGather, y el eje Y representa la tasa de utilización de red de la red EFA de 400 Gbps de p4d. Los 4 subgráficos representan los patrones comunes de grupos de comunicación en los que participan 1, 2, 4 y 8 rangos por instancia p4d en la operación AllGather, respectivamente.

Estas micropruebas muestran que SMDDP supera a NCCL con dos características clave:

1. El rendimiento máximo de SMDDP (aproximadamente 90% de utilización de ancho de banda) es mayor que el de NCCL (aproximadamente 80% de utilización de ancho de banda) en todas las configuraciones.
2. SMDDP alcanza el rendimiento máximo con tamaños de búfer mucho más pequeños que NCCL. Esto mejora especialmente la velocidad de entrenamiento para modelos más pequeños o cuando el usuario establece un tamaño de búfer AllGather pequeño en DeepSpeed (donde el tamaño de AllGather no necesita ser igual al tamaño de la capa).

Benchmarks de entrenamiento de modelos

En trabajos de entrenamiento a gran escala donde la comunicación GPU es un cuello de botella significativo, SMDDP puede aumentar notablemente la velocidad de entrenamiento, medida en TFLOPS/GPU del modelo.

Configuración	Rendimiento
Modelo/Entrenamiento	Cluster	Solución de Paralelismo de Datos en Fragmentos	TFLOPS/GPU del modelo con NCCL	TFLOPS/GPU del modelo con SMDDP	% de aceleración
13B Llama2 Longitud de secuencia: 4096 Tamaño de lote global: 4M tokens	64 nodos p4d.24xlarge (512 GPUs NVIDIA A100)	PyTorch FSDP	97.89	121.85	24.40%
65B GPT-NeoX Longitud de secuencia: 2048 Tamaño de lote global: 4M tokens	64 nodos p4d.24xlarge (512 GPUs NVIDIA A100)	DeepSpeed ZeRO Etapa 3*	99.23	108.66	9.50%

*Se utilizó el repositorio de Megatron-DeepSpeed de EleutherAI. El paralelismo tensor también estaba habilitado con un grado de paralelismo tensorial de ocho.

Nota: TFLOPS/GPU del modelo se basa en el cálculo de Utilización de FLOPS del Modelo definido en el paper aquí y en cifras de referencia en otros lugares se puede citar el TFLOPS/GPU de hardware como métrica de rendimiento. El TFLOPS/GPU de hardware se puede aproximar como 4/3 x TFLOPS/GPU del modelo.

Conclusión

En esta publicación, te mostramos cómo acelerar significativamente los trabajos de entrenamiento paralelo de datos fragmentados en Amazon SageMaker con solo dos líneas de cambio de código. El entrenamiento distribuido a gran escala se está volviendo cada vez más omnipresente con la aparición de LLMs, pero con esta escala también vienen costos elevados. Al reducir el cuello de botella de comunicación entre las GPU, SMDDP te ayuda a entrenar más rápido a gran escala y ahorrar recursos de cómputo. Puedes encontrar más ejemplos de SMDDP con entrenamiento paralelo de datos fragmentados en el repositorio de ejemplos de Amazon SageMaker en GitHub.

We will continue to update Zepes; if you have any questions or suggestions, please contact us!