Habilite una capacitación más rápida con la biblioteca paralela de datos de Amazon SageMaker

La capacitación en modelos de lenguaje grande (LLM) se ha vuelto cada vez más popular durante el último año con el lanzamiento de varios modelos disponibles públicamente, como Llama2, Falcon y StarCoder. Los clientes ahora están capacitando LLM de un tamaño sin precedentes que van desde mil millones hasta más de 175 mil millones de parámetros. La capacitación de estos LLM requiere importantes recursos informáticos y tiempo, ya que se deben utilizar de cientos a miles de unidades de procesamiento de gráficos (GPU) para manejar los vastos conjuntos de datos de capacitación y tamaños de modelos actuales. Un cuello de botella en la capacitación distribuida puede ser la comunicación de GPU manejada por la Biblioteca de comunicación colectiva de NVIDIA (NCCL). En algunos trabajos de capacitación distribuidos a gran escala, se puede dedicar más tiempo a la comunicación entre GPU que al cálculo real de la GPU. Para aliviar el cuello de botella en la comunicación de la GPU y permitir un entrenamiento más rápido, Amazon SageMaker se complace en anunciar una operación colectiva optimizada de AllGather como parte de la biblioteca paralela de datos distribuidos (SMDDP) de SageMaker. AllGather es la operación colectiva más utilizada en soluciones populares de paralelismo de datos con uso eficiente de la memoria como Optimizador de redundancia cero DeepSpeed ​​(ZeRO) y Paralelismo de datos totalmente fragmentados (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 puede habilitar SMDDP en sus scripts de capacitación de Amazon SageMaker y las mejoras de rendimiento que puede esperar.

Descripción general de la solución

Tradicional entrenamiento paralelo de datos Implica replicar un modelo completo en múltiples GPU, y cada modelo se entrena en diferentes fragmentos de datos del conjunto de datos. Durante el paso hacia atrás, se promedian los gradientes entre los trabajadores de la GPU para que cada réplica del modelo se actualice con los mismos valores de gradiente a pesar de haber sido entrenados con diferentes fragmentos de datos. Esta técnica permite un entrenamiento mucho más rápido en grandes conjuntos de datos al paralelizar el consumo de datos de entrenamiento. Sin embargo, algunos de los modelos grandes actuales (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 sea inutilizable. Para continuar cosechando los beneficios del paralelismo de datos y al mismo tiempo superar la memoria limitada de la GPU, se necesitan soluciones paralelas de datos fragmentados como DeepSpeed ​​ZeRO, PyTorch FSDP y Amazon. Biblioteca de paralelismo de modelos SageMaker han ganado popularidad.

En el paralelismo de datos fragmentados, en lugar de replicar todo el modelo en los trabajadores de GPU, los parámetros del modelo, los gradientes y los estados del optimizador se dividen y distribuyen (es decir, se fragmentan) entre las GPU en el trabajo de entrenamiento. Para realizar cálculos de paso hacia adelante y hacia atrás, los parámetros se recopilan de fragmentos en otros trabajadores de GPU para formar una o más capas de modelo. Una vez realizado el cálculo, estas capas se liberan de la memoria para permitir que se recopile el siguiente conjunto de capas. Tenga en cuenta que existen variantes del paralelismo de datos fragmentados en las que solo se fragmentan los estados y gradientes del optimizador, pero no los parámetros del modelo. AllGather todavía se usa en este tipo de paralelismo de datos fragmentados, pero solo antes del cálculo del paso directo para recopilar parámetros del modelo que han sido actualizados por diferentes fragmentos de estado de optimizador o gradiente de otros trabajadores de GPU. Consulte los diferentes Etapas DeepSpeed ​​ZeRO y el SHARD_GRAD_OP Estrategia de fragmentación FSDP para obtener más detalles.

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

Aunque NCCL se utiliza normalmente para AllGather en capacitación distribuida, su implementación subyacente de bajo nivel no está adaptada a la infraestructura de red de Amazon Elastic Compute Cloud (Amazon EC2) instancias y, por lo tanto, su rendimiento puede ralentizar el entrenamiento de un extremo a otro. La biblioteca SMDDP es una biblioteca de comunicación colectiva para GPU NVIDIA que sirve como reemplazo directo de NCCL y proporciona un mejor rendimiento para trabajos de capacitación distribuidos con PyTorch. Específicamente, SMDDP proporciona una implementación optimizada de AllGather para tipos de instancia p4d/p4de.

Dado que las operaciones colectivas como AllGather bloquean el cálculo de pases 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 un extremo a otro 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 recurriendo a NCCL.

Tutorial

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 Adaptador de tela elástica (EFA) red 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 integral para la comunicación de red entre nodos se adapta mejor a las características de EFA y la infraestructura de red de AWS al requerir menos saltos de paquetes en comparación con el anillo o el NCCL. patrón de comunicación del árbol.
2. RDACopiar para coordinar el tráfico de red local NVLink y EFA. GDRCopy es una biblioteca que proporciona comunicación de baja latencia entre los procesos de la CPU y los núcleos CUDA de la GPU. Con esta tecnología, podemos canalizar el movimiento de datos dentro y entre nodos.
3. Uso reducido de multiprocesadores de transmisión de GPU para devolver más potencia de cálculo a los núcleos de modelo. Las instancias AWS P4d/P4de están equipadas con GPU NVIDIA A100, cada una de las cuales tiene 108 multiprocesadores de transmisión. Mientras que NCCL requiere hasta 24 multiprocesadores de transmisión para ejecutar colectivos, los colectivos SMDDP solo usan hasta nueve multiprocesadores de transmisión. Los multiprocesadores de transmisión guardados pueden ser recogidos por núcleos de cómputo modelo para una ejecución más rápida.

Uso

Los colectivos SMDDP se integran de forma nativa con PyTorch a través del grupo de proceso abstracción en el torch.distributed módulo. 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 subyacente. backend, que proporciona la implementación de estas operaciones según el dispositivo informático utilizado. Los trabajos de entrenamiento de CPU a menudo utilizan el gloo o mpi backend mientras que las GPU NVIDIA usan el nccl back-end.

La biblioteca SMDDP entra en escena registrándose como un backend personalizado en la abstracción del grupo de procesos. Esto se hace mediante la declaración de importació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 con smddp. El smddp El backend sigue la misma semántica que el nccl backend y admite los mismos escenarios de entrenamiento.

Velocidad profunda

import smdistributed.dataparallel.torch.torch_smddp
deepspeed.init_distributed(dist_backend="smddp")  # replacing "nccl"

FSDP

import smdistributed.dataparallel.torch.torch_smddp
dist.init_process_group(backend="smddp")  # replacing "nccl"

Puntos de referencia

Comparamos el rendimiento de AllGather independiente donde la operación colectiva se ejecuta de forma aislada sin ningún entrenamiento modelo. A continuación se muestra un resultado de muestra en 32 instancias de p4d que comparan 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 la red EFA de 400 Gbps de p4d. Los 4 subgráficos representan los patrones comunes del grupo de comunicación donde tenemos 1, 2, 4 y 8 rangos por instancia de p4d que participan en la operación AllGather, respectivamente.

Estos microbenchmarks muestran que SMDDP supera a NCCL con dos características clave:

1. El rendimiento máximo de SMDDP (aproximadamente 90 % de utilización del ancho de banda) es mayor que el de NCCL (aproximadamente 80 % de utilización del ancho de banda) en todas las configuraciones.
2. SMDDP alcanza el máximo rendimiento con tamaños de búfer mucho más pequeños que NCCL. Esto mejora particularmente las velocidades 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 tiene por qué ser igual al tamaño de la capa).

Puntos de referencia de entrenamiento modelo

En trabajos de entrenamiento a gran escala donde la comunicación GPU es un cuello de botella importante, SMDDP puede mejorar notablemente las velocidades de entrenamiento, según lo medido por el modelo TFLOPS/GPU.

*Megatron-DeepSpeed ​​de EleutherAI Se utilizó el repositorio. El paralelismo tensorial también se habilitó con un grado de paralelo tensorial de ocho.

Nota: El modelo TFLOPS/GPU se basa en el cálculo de utilización del modelo FLOPS definido en el documento. aquí y las cifras de referencia en otros lugares pueden citar TFLOPS/GPU de hardware como métrica de rendimiento. Los TFLOPS/GPU de hardware se pueden aproximar a 4/3 x TFLOPS/GPU del modelo.

Conclusión

En esta publicación, le mostramos cómo acelerar significativamente los trabajos de capacitación paralela de datos fragmentados en Amazon SageMaker con solo dos líneas de cambio de código. La formación distribuida a gran escala se está volviendo cada vez más omnipresente con la aparición de los LLM, pero esta escala conlleva altos costos. Al reducir el cuello de botella de comunicación entre GPU, SMDDP le ayuda a entrenar más rápido a escala y ahorrar recursos informáticos. Puede encontrar más ejemplos de SMDDP con entrenamiento paralelo de datos fragmentados en el Repositorio GitHub de ejemplos de Amazon SageMaker.

Sobre los autores

Apoorv Gupta es ingeniero de desarrollo de software en AWS y se centra en la creación de sistemas óptimos de aprendizaje profundo para la infraestructura y el hardware de AWS. Está interesado en la informática distribuida, los sistemas de aprendizaje profundo y los aceleradores de aprendizaje automático. Fuera del trabajo, a Apoorv le gustan los viajes, las caminatas y los videojuegos.

Karan Dhiman es ingeniero de desarrollo de software en AWS, con sede en Toronto, Canadá. Le apasiona el espacio del aprendizaje automático y la creación de soluciones para acelerar las cargas de trabajo informáticas distribuidas.

Ruhan Prasad es un ingeniero de desarrollo de software en AWS que trabaja para hacer que la capacitación en aprendizaje profundo distribuido sea más rápida, económica y fácil de usar en SageMaker. Fuera del trabajo, a Ruhan le gusta jugar tenis, viajar y cocinar.

Zhaoqi Zhu es Ingeniero Senior de Desarrollo de Software en AWS, apasionado por los sistemas distribuidos y las optimizaciones de bajo nivel. Le gusta ver partidos de fútbol mientras bebe refrescos (no dietéticos).