Performance Optimization and Comparative Analysis of Generative AI Models on Advanced Accelerators

Resumen Técnico: Optimización del Rendimiento y Análisis Comparativo de Modelos de IA Generativa en Aceleradores Avanzados

Planteamiento del Problema
A pesar de las impresionantes capacidades de los modelos de IA Generativa, incluyendo los Grandes Modelos de Lenguaje (LLMs) y los modelos de difusión, su despliegue enfrenta obstáculos significativos. Estos desafíos incluyen requisitos sustanciales de memoria, latencia de inferencia extendida, altas demandas computacionales y elevados costos de hardware. Estos problemas se ven exacerbados al evaluar modelos en plataformas heterogéneas, donde las variaciones en formatos numéricos, anchos de banda de memoria y pilas de software interactúan de forma compleja con las arquitecturas de los modelos y las características de la carga de trabajo. El artículo aborda la necesidad de un estudio sistemático para optimizar el rendimiento y comparar estos modelos a través de diversas tareas descendentes en sistemas modernos de Computación de Alto Rendimiento (HPC) y aceleradores avanzados.

Metodología
Los autores realizaron una evaluación exhaustiva en tres sistemas de HPC distintos: Perlmutter (NERSC, con GPUs NVIDIA A100), Expanse (SDSC, con GPUs NVIDIA V100, A100 y H100) y Voyager (SDSC, con aceleradores Intel Gaudi). El estudio se centró en tres áreas principales:

1. Cuantización Post-Entrenamiento (PTQ) de Precisión Mixta: Los autores introdujeron un novedoso marco de PTQ de precisión mixta sensible a la sensibilidad, diseñado para optimizar LLMs (específicamente TinyLlama-1.1B) a través de aceleradores NVIDIA A100 e Intel Gaudi.

   • Análisis de Sensibilidad: El marco emplea un enfoque de dos fases. Calcula puntuaciones de sensibilidad a nivel de capa de transformador analizando las magnitudes de los pesos (excluyendo sesgos y parámetros de normalización) y aplicando un umbral de poda unificado basado en un objetivo de dispersión (30%).
   • Estrategia de Cuantización: Las capas sensibles reciben una mayor precisión, mientras que las capas robustas reciben una menor precisión. El estudio evaluó la cuantización "falsa" (simulada en FP32) y la cuantización "real" usando FP16/INT8 para A100 y BF16/FP8 (vía Intel Neural Compressor) para Gaudi.
   • Conjuntos de Datos: La evaluación se realizó en WikiText-2 (modelado de lenguaje), HellaSwag (razonamiento de sentido común) y BoolQ (comprensión de lectura).

2. Ajuste Fino (Fine-Tuning) de Rendimiento en Gaudi: El estudio evaluó el ajuste fino de tres LLMs de código abierto (Llama3.1-8B, Gemma2-27B y Llama3.3-70B) en Voyager utilizando la librería Optimum-Habana.

   • Paralelización: Los autores compararon las estrategias de Paralelismo de Datos Distribuido (DDP), DeepSpeed (ZeRO3) y Paralelismo Total Fragmentado (FSDP).
   • Optimización: Se analizó el impacto de Flash Attention (FA) y las diferencias generacionales entre Gaudi1, Gaudi2 y Gaudi3.

3. Entrenamiento de Modelos de Difusión: Un modelo de Super-Resolución, utilizado para mejorar imágenes astrofísicas de polvo galáctico, fue entrenado en nodos Gaudi y comparado contra GPUs NVIDIA (V100, A100, H100) en Expanse. El modelo fue paralelizado usando PyTorch DDP y ejecutado en modo SynapseAI Lazy.

Resultados Clave

• Rendimiento de la Cuantización: La cuantización real en A100 y Gaudi2 logró ratios de compresión entre 2x y 6x con una pérdida de precisión mínima. Por ejemplo, en la tarea WikiText-2, Gaudi2 logró un ratio de compresión de 5.93x con una perplejidad de 12.39 (comparado con 12.02 para FP32), mientras que A100 logró 3.77x de compresión con una perplejidad de 12.02. Las mejoras en el rendimiento (throughput) fueron significativas, mostrando Gaudi2 hasta 34,081 tokens/seg en escenarios de cuantización simulada.
• Eficiencia del Ajuste Fino:
  • Ganancias Generacionales: Las tarjetas Gaudi2 demostraron una reducción de tiempo de ejecución de 2.5x a 3x en comparación con Gaudi1, atribuido al incremento de 8 a 24 Núcleos de Procesamiento de Tensores (TPCs). Gaudi3 ofreció una ganancia de rendimiento adicional de 1.8x a 2x sobre Gaudi2.
  • Flash Attention: Habilitar Flash Attention en Gaudi1 y Gaudi2 resultó en una mejora de rendimiento del 10% al 20% al minimizar el movimiento de datos entre tipos de memoria.
  • Sobrecarga de Fragmentación (Sharding): Para el modelo Gemma2 de 27B parámetros, DDP superó decisivamente a los métodos de fragmentación de modelo (DeepSpeed y FSDP). Los autores encontraron que, si bien la fragmentación reduce la huella de memoria, introduce una sobrecarga significativa que degrada el rendimiento. En consecuencia, la fragmentación solo debe emplearse cuando un modelo no puede caber dentro de la HBM de una sola tarjeta.
• Escalabilidad del Modelo de Difusión:
  • GPUs NVIDIA: Exhibieron un escalamiento casi perfecto con una aceleración de ~2x de V100 a A100 y otra de ~2x de A100 a H100.
  • Intel Gaudi: Mostró una aceleración de 4x de Gaudi1 a Gaudi2. Sin embargo, esta tendencia no continuó linealmente de Gaudi2 a Gaudi3, indicando un estancamiento en la eficiencia de escalamiento para esta carga de trabajo específica.

Significancia y Reivindicaciones
El artículo afirma proporcionar un análisis comparativo sistemático del rendimiento de la IA Generativa a través de entornos de HPC heterogéneos, ofreciendo perspectivas prácticas para el despliegue. Sus principales contribuciones incluyen:

1. Un novedoso marco de cuantización de precisión mixta sensible a la sensibilidad que equilibra la precisión del modelo y la compresión a través de diferentes arquitecturas de hardware.
2. Evidencia empírica que demuestra que, para modelos que caben dentro de la memoria de una sola tarjeta, el DDP estándar es superior a las estrategias de fragmentación complejas como DeepSpeed o FSDP, las cuales introducen una sobrecarga innecesaria.
3. Un benchmarking detallado de la evolución de los aceleradores Intel Gaudi, resaltando ganancias de rendimiento específicas (y limitaciones) en las actualizaciones generacionales y la eficacia de Flash Attention en este ecosistema.
4. La validación de las capacidades de entrenamiento de modelos de difusión en aceleradores Gaudi, proporcionando una base para aplicaciones de computación científica como la mejora de imágenes astrofísicas.

El trabajo enfatiza que, si bien los aceleradores avanzados ofrecen ganancias sustanciales de rendimiento, el despliegue óptimo requiere una consideración cuidadosa del tamaño del modelo, las restricciones de memoria y los trade-offs específicos entre la cuantización, las estrategias de paralelización y las generaciones de hardware.