Mirage: un superoptimizador de álgebra tensorial multinivel que automatiza la generación del kernel de GPU para aplicaciones PyTorch

Con el creciente crecimiento de la inteligencia artificial (introducción de grandes modelos de lenguaje (LLM) e IA generativa), ha habido una demanda creciente de unidades de procesamiento de gráficos (GPU) más eficientes. Las GPU son hardware especializado que se utiliza ampliamente para tareas informáticas exigentes y son capaces de ejecutar cálculos en paralelo. Escribir kernels de GPU adecuados es importante para utilizar las GPU en todo su potencial. Esta tarea requiere bastante tiempo y es compleja, y requiere una profunda experiencia en la arquitectura de GPU y algunos lenguajes de programación como C++, CUDA, etc.

Los compiladores de Machine Learning ML como TVM, Triton y Mojo proporcionan cierta automatización, pero aún necesitan el manejo manual de los núcleos de la GPU para obtener el resultado óptimo. Para lograr resultados óptimos y evitar tareas manuales, investigadores de La Universidad Carnegie Mellon ha desarrollado Mirageuna herramienta innovadora diseñada para automatizar la generación de núcleos de GPU de alto rendimiento buscándolos y generándolos. Los núcleos generados por Mirage se pueden usar directamente en tensores de PyTorch y llamarse en programas de PyTorch. Los usuarios necesitan escribir unas pocas líneas de código en Mirage en comparación con el script tradicional, que utiliza muchas líneas.

Mirage puede verse como un cambiador del futuro, que logrará una alta productividad, un mejor rendimiento y una mayor corrección en las aplicaciones de IA. Escribir códigos manuales requiere una importante experiencia en ingeniería debido a la naturaleza compleja de la arquitectura de la GPU, pero Mirage simplifica el proceso generando núcleos automáticamente, lo que facilita y simplifica las tareas de los ingenieros.

Los núcleos de GPU escritos manualmente pueden tener algunos errores, lo que dificulta lograr los resultados requeridos, pero la investigación sobre Mirage ha demostrado que los núcleos generados por Mirage son entre 1,2 y 2,5 veces más rápidos que el mejor código escrito por humanos. Además, la integración de Mirage en PyTorch reduce la latencia general entre un 15 y un 20 %.

      # Use Mirage to generate GPU kernels for attention
      import mirage as mi
      graph = mi.new_kernel_graph()
      Q = graph.new_input(dims=(64, 1, 128), dtype=mi.float16)
      K = graph.new_input(dims=(64, 128, 4096), dtype=mi.float16)
      V = graph.new_input(dims=(64, 4096, 128), dtype=mi.float16)
      A = graph.matmul(Q, K)
      S = graph.softmax(A)
      O = graph.matmul(S, V)
      optimized_graph = graph.superoptimize()

El código en Mirage ocupa pocas líneas en comparación con el método tradicional con muchas líneas

Todos los cálculos en las GPU se centran en núcleos, que son funciones que se ejecutan en paralelo alrededor de múltiples multiprocesadores de transmisión (SM) en forma de un solo programa y múltiples datos (SPMD). Los núcleos organizan el cálculo en una cuadrícula de bloques de subprocesos, y cada bloque de subprocesos se ejecuta en un único SM. Cada bloque tiene además múltiples subprocesos para realizar cálculos en elementos de datos individuales.

La GPU sigue una jerarquía de memoria particular con:

• Registre el archivo para acceder rápidamente a los datos
• Memoria compartida: compartida por todos los subprocesos de un bloque para un intercambio de datos eficiente.
• Memoria del dispositivo: accesible para todos los subprocesos de un kernel

La arquitectura se representa con la ayuda de la representación uGraph, que contiene gráficos en múltiples niveles: nivel de kernel, nivel de bloque de subprocesos y nivel de subproceso con cálculo encapsulado a nivel de kernel en toda la GPU, nivel de bloque de subprocesos que aborda el cálculo en un multiprocesador de transmisión individual ( SM) y gráfico de subprocesos que aborda el cálculo en el nivel CUDA o núcleo tensor. uGraph proporciona una forma estructurada de representar los cálculos de GPU.

Cuatro categorías de optimización de GPU:

1. Normalización + Lineal

Los LLM generalmente utilizan técnicas LayernNorm, RMSNorm, GroupNorm y BatchNorm, que los compiladores de ML a menudo tratan por separado. Esta separación se debe a que las técnicas de normalización requieren operaciones tanto de reducción como de difusión. Estas capas de normalización se pueden fusionar con capas lineales mediante multiplicación de matrices.

2. LoRA + Lineal

Fusiona la adaptación de rango bajo (LoRA), una técnica para adaptar modelos previamente entrenados a nuevas tareas o conjuntos de datos al tiempo que reduce los requisitos computacionales con capas lineales. Es 1,6 veces más rápido que los sistemas existentes.

3. MLP cerrado

Combina dos MatMuls, activación SiLU y multiplicación de elementos. MLP cerrado reduce la sobrecarga de lanzamiento del kernel y el acceso a la memoria del dispositivo a 1,3 veces más rápido que la mejor línea de base.

4. Variantes de atención

a. Normalización de claves de consulta

Chameleon, ViT-22B y el artículo reciente de Google introdujeron la normalización de claves de consulta y fusionaron LayerNorm en el núcleo de atención. Este kernel personalizado también realiza optimizaciones de GPU existentes adaptadas a la atención con una mejora de rendimiento de 1,7 a 2,5 veces.

b. Atención latente de múltiples cabezas

Optimiza el uso de la memoria al comprimir la caché de atención de valor-clave tradicional en un vector latente más compacto. Este cambio introduce dos capas lineales ante la atención. Mirage genera un kernel personalizado que integra las capas lineales con el mecanismo de atención en un solo kernel. Esto evita almacenar vectores clave-valor intermedios en la memoria del dispositivo GPU.

En conclusión, Mirage aborda el desafío crítico de lidiar con núcleos de GPU con alto contenido de energía en problemas avanzados de inteligencia artificial. Elimina los problemas de una importante inversión de tiempo, una alta experiencia en codificación y la generación de errores al proporcionar los mejores núcleos de GPU óptimos que funcionan en un entorno basado en PyTorch. También aborda las lagunas que la informática manual podría pasar por alto, acelerando la implementación de LLM y otras tecnologías de inteligencia artificial en aplicaciones del mundo real.

Mira el página de GitHub y Detalles. Todo el crédito por esta investigación va a los investigadores de este proyecto. Además, no olvides seguirnos en Gorjeo y únete a nuestro Canal de telegramas y LinkedIn Grarriba. Si te gusta nuestro trabajo, te encantará nuestro hoja informativa..

No olvides unirte a nuestro SubReddit de más de 50.000 ml

¿Está interesado en promocionar su empresa, producto, servicio o evento ante más de 1 millón de desarrolladores e investigadores de IA? ¡Colaboremos!

Nazmi Syed es pasante de consultoría en MarktechPost y está cursando una licenciatura en ciencias en el Instituto Indio de Tecnología (IIT) Kharagpur. Tiene una profunda pasión por la ciencia de datos y explora activamente las amplias aplicaciones de la inteligencia artificial en diversas industrias. Fascinada por los avances tecnológicos, Nazmi está comprometida a comprender e implementar innovaciones de vanguardia en contextos del mundo real.