En la formación de modelos de lenguaje, la comunicación de Expert Parallel (EP) para modelos de mezcla de expertos (MoE) es un reto. Esta comunicación es esencialmente un todo para todos, pero su dinámica y escasez dificultan la implementación y optimización.
Este artículo presenta una solución eficiente de comunicación EP para MoE, denominada Hybrid-EP, y su aplicación en la familia de frameworks NVIDIA Megatron, utilizando plataformas NVIDIA Quantum InfiniBand y NVIDIA Spectrum-X Ethernet. También se analiza la efectividad de Hybrid-EP en la formación de modelos en escenarios reales.
Desafíos de eficiencia en la formación de modelos MoE a gran escala
DeepSeek-V3 representa una nueva generación de modelos MoE a gran escala. Estos modelos equilibran el rendimiento y la sobrecarga computacional mediante la activación esparcida de tamaño de hiperparámetro. Sin embargo, presentan serios desafíos para los frameworks de formación de modelos existentes.
Cuellos de botella en la eficiencia de comunicación: El modelo MoE depende de expertos paralelos y requiere comunicación frecuente todo para todos. A medida que aumenta el número de expertos, también lo hace la carga de comunicación EP. En DeepSeek-V3, el tiempo de comunicación puede representar más del 50% del tiempo total de formación sin optimización.
- Desequilibrio de carga: Mecanismos de enrutamiento dinámico hacen que algunos “expertos calientes” reciban más tokens que el promedio, mientras que los “expertos fríos” están infrautilizados, resultando en una carga computacional desigual.
- Retos de adaptabilidad del framework: Los modelos MoE actuales imponen requisitos más complejos para las estrategias paralelas y la programación de recursos dinámica, además de necesitar optimización para maximizar el potencial del hardware de próxima generación.
Optimización del framework de formación MoE y solución de comunicación
Megatron Core, una biblioteca de formación de modelos a gran escala y de código abierto, es fundamental para la formación de modelos MoE a gran escala. Sus beneficios principales incluyen:
- Estrategias de paralelismo multidimensional que admiten paralelismo tensorial, secuencial, en pipeline y paralelismo experto MoE, entre otros, que se pueden combinar para adaptarse a cargas de trabajo complejas.
- Optimización de recursos y eficiencia que integra entrenamiento de precisión mixta FP8, descarga de valores de activación y funciones de recalculo para reducir el consumo de memoria GPU.
- Adaptación específica a MoE que proporciona soporte completo para modelos MoE como DeepSeek, Mixtral y Qwen, permitiendo una formación escalable y eficiente.
Hybrid-EP como solución de optimización de comunicación
Hybrid-EP es una biblioteca de comunicación MoE EP recientemente diseñada. Utiliza avances en hardware y software en la plataforma NVIDIA para lograr un ancho de banda de comunicación cercano a los límites del hardware y minimizar el uso de recursos de hardware GPU.
Implementa dos operadores clave en la comunicación EP de MoE: dispatch, que dirige los tokens generados por el operador de atención a los expertos correspondientes, y combine, que envía los tokens de vuelta al operador de atención.
Los objetivos de diseño de Hybrid-EP incluyen aprovechar las últimas tecnologías de comunicación en la plataforma NVIDIA, maximizando el ancho de banda en la comunicación híbrida y optimizando el uso de recursos GPU.
Hybrid-EP logra un ancho de banda de comunicación óptimo utilizando menos multiprocesadores de transmisión (SM), dejando más recursos disponibles para el cómputo. Soporta operaciones dispatch y combine en baja precisión.
Cada bloque CUDA se diseña como un canal de datos independiente que ocupa un SM para ejecutar un pipeline de datos completo. Diferentes grupos de warp manejan distintas etapas del pipeline sin necesidad de sincronización.
Las cajas punteadas en la Figura 2 representan las etapas del pipeline utilizadas en la comunicación de red RDMA. Grupos de warp RDMA son responsables de transmitir tráfico de red a través de las tarjetas de interfaz de red (NIC) utilizando tecnología IBGDA.
Durante este proceso, los tokens se envían según el mapa de enrutamiento para evitar transmitir datos no deseados. Cada bloque CUDA utiliza este pipeline para procesar los datos.
El operador combine realiza operaciones de acumulación de alta precisión, que se llevan a cabo jerárquicamente dentro del SM. En el caso de múltiples nodos, se completan los trabajos acumulativos de forma parcial antes de ser enviados al GPU correspondiente.
Hybrid-EP se ha probado en múltiples plataformas de hardware con condiciones específicas, logrando un rendimiento notable en la comunicación y procesamiento de datos.
Se realizaron pruebas en configuraciones de GPU NVIDIA DGX Hopper, analizando el rendimiento de Hybrid-EP en condiciones específicas de ancho de banda y comunicación.
Finalmente, se testó el rendimiento de Hybrid-EP en grandes redes NVLink utilizando 36 GPUs, logrando un uso eficiente del ancho de banda.
Casos prácticos: Verificación de modelo combinado y hardware
Hybrid-EP se basa en plantillas y implementaciones CUDA C, requiriendo trabajo adicional para su uso en el framework Megatron Core basado en PyTorch. Está disponible en la rama DeepEP/Hybrid-EP, facilitando la integración y prueba para los usuarios.
La gestión de buffers es crucial, ya que Hybrid-EP requiere un mecanismo razonable de gestión de memoria. Los buffers se dividen en dos categorías: buffers registrados y normales.
- Buffer registrado: Memoria GPU especialmente registrada accesible por kernels en otros rangos.
- Buffer normal: Memoria GPU asignada con cudaMalloc, gestionada por el asignador de PyTorch.
La aplicación y registro de buffers deben completarse preferiblemente durante la fase de inicialización de Hybrid-EP. Sin embargo, dado que el modelo MoE es dinámico, se requiere una estrategia de preasignación de buffers para controlar el uso de memoria GPU.
El flujo de trabajo de Hybrid-EP en el entorno de PyTorch requiere sincronización para determinar tamaños de tensores, lo que puede evitarse si se predefinen tamaños de buffers adecuados.
Prácticas de optimización en Grace Blackwell
Megatron Core ha integrado Hybrid-EP en la plataforma Grace Blackwell, optimizándose para diferentes tipos de modelos MoE.
| Modelo | Precisión | Despachador | TFLOPS/GPU | Mejora |
| DeepSeek-V3 | MXFP8 | DeepEP | 829 | 1x |
| MXFP8 | Hybrid-EP | 943 | 1.14x | |
| DeepSeek-V3- FSDP | MXFP8 | A2A | 597 | 1x |
| MXFP8 | Hybrid-EP | 645 | 1.08x | |
| Qwen 3 235B | BF16 | A2A | 665 | 1x |
| BF18 | Hybrid-EP | 698 | 1.05x | |
| MXFP8 | A2A | 728 | 1x | |
| MXFP8 | Hybrid-EP | 800 | 1.10x |
Los resultados indican:
- DeepSeek-V3, 256 expertos, topk-8, muestra una mejora del 14% con Hybrid-EP.
- En el escenario Megatron-FSDP, Hybrid-EP logra un 8% de mejora.
- Para Qwen 3 235B, se observa una mejora del 5.5% en BF16 y de aproximadamente 9.9% en MXFP8.
Para más información, descubre cómo NVIDIA está habilitando 10x de rendimiento y 1/10 del costo para implementar modelos MoE.
