Axonn: Avance de la capacitación de modelos de lenguaje grande a través de la computación paralela híbrida de cuatro dimensiones

La capacitación de la red neuronal profunda (DNN) ha experimentado un crecimiento sin precedentes con el aumento de los grandes modelos de idiomas (LLM) y la IA generativa. La efectividad de estos modelos se correlaciona directamente con el aumento de su tamaño, un desarrollo hecho posible por los avances en la tecnología y marcos de GPU como Pytorch y TensorFlow. Sin embargo, la capacitación de redes neuronales con miles de millones de parámetros presenta desafíos técnicos significativos a medida que los modelos exceden la capacidad de las GPU individuales. Esto requiere distribuir el modelo en múltiples GPU y operaciones de multiplicación de matriz de paralelización. Varios factores afectan la eficiencia de la capacitación, incluido el rendimiento computacional sostenido, las operaciones de comunicación colectiva sobre los subcomunicadores y la superposición de cálculo con colectivos que no son bloqueos.

Los recientes esfuerzos para entrenar LLM han superado los límites de la utilización del clúster basada en GPU, aunque la eficiencia sigue siendo un desafío. Meta entrenado Llama 2 usando 2,000 GPU A100 NVIDIA, mientras que el paralelismo de la tubería de Megatron-LM alcanzó el 52% del rendimiento máximo al evaluar un modelo de parámetros de 1000B en 3,072 GPU. La combinación Megatron-LM y DeepSpeed ​​alcanzó el 36% del rendimiento máximo al entrenar un modelo de parámetros de 530B en 4,480 GPU A100. Megascale logró el 55.2% del rendimiento máximo para un modelo de parámetros de 175B en 12,288 GPU A100. En los sistemas AMD, el entrenamiento forjado alcanzó el 28% del rendimiento máximo en 2,048 GPU MI250X, mientras que otros estudios lograron el 31.96% del pico al evaluar un modelo de parámetro 1T en 1.024 GPU MI250X.

Investigadores de la Universidad de Maryland, College Park, EE. UU.; Instituto Max Planck de Sistemas Inteligentes, Tübingen, Alemania; y la Universidad de California, Berkeley, EE. UU. Han propuesto Axonn, un nuevo algoritmo paralelo híbrido de cuatro dimensiones implementado en un marco de código abierto altamente escalable, portátil. Los investigadores introdujeron varias optimizaciones de rendimiento en AXONN para mejorar el rendimiento del núcleo de multiplicación de la matriz, superponerse efectivamente a los colectivos sin bloqueo con cálculo y emplear el modelado de rendimiento para identificar configuraciones óptimas. Más allá del rendimiento, también abordaron las preocupaciones críticas de privacidad y derechos de autor que surgen de la memorización de datos de capacitación en LLM al investigar la “memorización catastrófica”. Un parámetro de 405 mil millones de LLM se ajusta fino usando Axonn en la frontera.

AxoNN is evaluated on three leading supercomputing platforms: Perlmutter at NERSC/LBL with NVIDIA A100 GPUs (40GB DRAM each), Frontier at OLCF/ORNL featuring AMD Instinct MI250X GPUs (128GB DRAM each, divided into two independently managed 64GB Graphic Compute Dies), and Alps at CSCS equipped with GH200 Superchips (96 GB DRAM por GPU H100). Todos los sistemas utilizan cuatro NIC de HPE Slingshot 11 por nodo, cada uno entregando velocidades de enlace bidireccionales de 25 GB/s. Las mediciones de rendimiento siguen a una metodología rigurosa, ejecutando diez iteraciones y promediando los últimos ocho para tener en cuenta la variabilidad del calentamiento. La evaluación comparativa se realiza para obtener resultados contra los valores de rendimiento máximo teórico, informando el porcentaje de pico alcanzado y total flop/s sostenido BF16 sostenido.

Axonn muestra un rendimiento de escala débil excepcional en las tres supercomputadoras con transformadores de estilo GPT. La escala casi ideal se logra hasta 4,096 GPU/GCD en todas las plataformas, cubriendo el rango de hardware típico para el entrenamiento LLM a gran escala. Mientras se ejecuta el modelo 60B en 6,144 GPU de H100 de Alpes, muestra una ligera reducción de eficiencia a 76.5% en comparación con el rendimiento de 1.024 GPU, la extensa disponibilidad de GPU de Frontier permite pruebas de escala sin precedentes. Axonn mantiene una escala débil casi perfecta de hasta 8,192 GCD en la frontera con un 88.3% de eficiencia en relación con el rendimiento de 512 GCD. En Perlmutter, Axonn logra constantemente el 50% o más del pico de 312 tflop/s anunciado por GPU. La escala de rendimiento lineal se evidencia por un aumento de casi 8 veces en las operaciones de punto flotante sostenidos, desde 80.8 PFLOP/s en 512 GPU hasta un impresionante 620.1 PFLOP/s en 4,096 GPU.

En conclusión, los investigadores han introducido Axonn, cuyas contribuciones al aprendizaje automático se extienden más allá de las métricas de rendimiento al proporcionar acceso escalable, fácil de usar y portátil al paralelismo modelo. Permite que la capacitación y el ajuste de modelos más grandes bajo restricciones de computación de productos básicos, lo que permite que las bases de código de entrenamiento LLM secuenciales utilicen recursos distribuidos de manera eficiente. Además, al democratizar la capacidad de ajustar grandes modelos en datos específicos del dominio, Axonn amplía las capacidades de los profesionales en varios campos. Por lo tanto, existe una urgencia para comprender y abordar los riesgos de memorización, ya que más investigadores ahora pueden trabajar con modelos de escala y complejidad sin precedentes que pueden capturar inadvertidamente información confidencial de los datos de capacitación.


Verificar el Papel. Todo el crédito por esta investigación va a los investigadores de este proyecto. Además, siéntete libre de seguirnos Gorjeo Y no olvides unirte a nuestro Subreddit de 80k+ ml.

🚨 Lectura de lectura recomendada Liberaciones de investigación de IA: un sistema avanzado que integra el sistema de IA del agente y los estándares de cumplimiento de datos para abordar las preocupaciones legales en los conjuntos de datos de IA


Sajjad Ansari es un pregrado de último año de IIT Kharagpur. Como entusiasta de la tecnología, profundiza en las aplicaciones prácticas de la IA con un enfoque en comprender el impacto de las tecnologías de IA y sus implicaciones del mundo real. Su objetivo es articular conceptos complejos de IA de manera clara y accesible.