Los modelos de lenguaje grande (LLM) basados en la mezcla del marco de expertos (MOE) han ganado tracción por su capacidad de escala de manera eficiente activando solo un subconjunto de parámetros por token. Esta dispersión dinámica permite a los modelos MOE retener una alta capacidad de representación al tiempo que limita el cálculo por token. Sin embargo, con su creciente complejidad y tamaño del modelo que se acerca a billones de parámetros, capacitarlos de manera eficiente requiere una innovación algorítmica y una optimización de software de hardware estrechamente integrada. Estos desafíos son especialmente relevantes al implementar modelos en aceleradores de IA no estándar como ASCEND NPUS, que requieren una alineación arquitectónica específica para ofrecer un rendimiento óptimo.
Un desafío técnico importante radica en la utilización ineficiente de los recursos de hardware mientras capacita a los LLMS dispersos. Dado que solo una parte de los parámetros está activa para cada token, las cargas de trabajo entre los dispositivos se vuelven desequilibrados, lo que lleva a retrasos de sincronización y potencia de procesamiento infrautilizada. Este desequilibrio también afecta la utilización de la memoria ya que diferentes expertos procesan diferentes números de tokens, a veces excediendo la capacidad. Estas ineficiencias se agravan a gran escala, como entre miles de chips de IA, donde los cuellos de botella de la gestión de la comunicación y la memoria obstaculizan significativamente el rendimiento. La incapacidad de aprovechar completamente la promesa computacional de escasez en la práctica restringe la implementación de tales modelos en sistemas de hardware como ASCEND NPUS.
Se han propuesto varias estrategias para abordar estos desafíos. Estos incluyen pérdidas auxiliares para equilibrar la distribución del token entre expertos y estrategias de caída y caída que limitan la sobrecarga de expertos al descartar los tokens que exceden la capacidad. Sin embargo, estas técnicas reducen el rendimiento del modelo o introducen ineficiencias en la memoria y el cálculo. Otros esfuerzos incluyen la colocación de expertos heurísticos y los patrones de comunicación tradicionales como el envío total, pero estos a menudo no pueden escalar bien o mantener un alto rendimiento. Además, las técnicas estándar de ahorro de memoria como la recomputación suelen ser de grano grueso, dirigidos a capas enteras en lugar de operaciones específicas, lo que lleva a un mayor tiempo de ejecución sin ahorros de memoria proporcionales.
Investigadores del equipo de Pangu en Huawei Cloud introdujeron un enfoque de entrenamiento altamente estructurado y optimizado para grandes modelos MOE adaptados a NPUS. Desarrollaron Pangu Ultra Moe, un SPARSE LLM con 718 mil millones de parámetros, centrándose en alinear la arquitectura del modelo y el diseño del sistema con las capacidades del hardware Ascend. Su enfoque comienza con un proceso de configuración del modelo basado en simulación que evalúa miles de variantes de arquitectura utilizando métricas basadas en el comportamiento real de hardware. Estas simulaciones informan las decisiones de diseño antes de que se realice cualquier entrenamiento físico, ahorrando así recursos computacionales sustanciales y permitiendo el ajuste informado de los hiperparámetros del modelo.
El método de simulación analiza combinaciones de parámetros como el número de capas, el tamaño oculto y el recuento de expertos utilizando una estrategia de paralelismo de cinco dimensiones que incluye el paralelismo de la tubería, el paralelismo tensor, el paralelismo experto, el paralelismo de datos y el paralelismo de contexto. La configuración del modelo final adoptada por Huawei incluyó 256 expertos, un tamaño oculto 7680 y 61 capas del transformador. Para optimizar aún más el rendimiento, los investigadores integraron un mecanismo de superposición de tuberías adaptativas para enmascarar los costos de comunicación y utilizaron la comunicación jerárquica total para reducir la transferencia de datos entre nodos. Emplearon la recomputación de grano fino, como recomputar solo vectores de valor clave en los módulos de atención, e introdujo el intercambio de tensor para la memoria de activación de descarga a los dispositivos de host de forma dinámica.
Pangu Ultra Moe logró una utilización de Flops de Modelo (MFU) de 30.0% y tokens procesados a una tasa de 1.46 millones por segundo usando 6,000 NPUS ascend. La MFU de línea de base fue del 18.9% con 0.61 millones de tokens por segundo en 4.000 NPU. Los investigadores también introdujeron estrategias dinámicas de colocación de expertos, mejorando el equilibrio de carga a nivel de dispositivo y logrando una mejora relativa del 10% de MFU. El modelo se desempeñó competitivamente en evaluaciones de referencia, que alcanzó el 81.3% en AIME2024, 97.4% en Math500, 94.8% en CLUEWSC y 91.5% en MMLU. En el dominio de la salud, superó a Deepseek R1 al calificar 87.1% en MEDQA y 80.8% en MEDMCQA, confirmando su fuerza en aplicaciones específicas de dominio.
Este estudio ilustra cómo el equipo de Pangu en Huawei abordó efectivamente las dificultades centrales de entrenar modelos MOE masivos en hardware especializado. Su búsqueda sistemática de arquitectura, técnicas de comunicación eficientes y optimizaciones de memoria a medida representan un marco fuerte para la capacitación de IA escalable. El trabajo demuestra formas prácticas de desbloquear el potencial de rendimiento de los modelos dispersos y establece una dirección para el diseño de IA del sistema futuro.
Verificar Papel aquí. 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 95k+ ml de subreddit.
Aquí hay una breve descripción de lo que estamos construyendo en MarkTechPost:
Nikhil es consultor interno en MarktechPost. Está buscando un doble grado integrado en materiales en el Instituto Indio de Tecnología, Kharagpur. Nikhil es un entusiasta de AI/ML que siempre está investigando aplicaciones en campos como biomateriales y ciencias biomédicas. Con una sólida experiencia en la ciencia material, está explorando nuevos avances y creando oportunidades para contribuir.