1 puntos por GN⁺ 2024-05-13 | 1 comentarios | Compartir por WhatsApp

Características de la GPU H100

  • Ofrece 80 GB de memoria HBM3 y 3 TB/s de ancho de banda (en la práctica es un poco menor)
  • Ofrece 50 MB de caché L2 y 12 TB/s de ancho de banda. Está dividida en dos secciones de 25 MB dentro de la GPU y conectadas por un crossbar (el crossbar es un factor de degradación del rendimiento)
  • Está compuesta por 132 Streaming Multiprocessors (SM), y cada SM tiene la siguiente configuración:
    • Ofrece hasta 227 KB de memoria compartida dentro de una caché L1 de 256 KB (en conjunto, alrededor de 33 TB/s de ancho de banda)
    • Incluye Tensor Memory Accelerator (TMA), que permite generación de direcciones asíncrona y fetching de memoria. También ofrece funciones como soporte para la red de memoria on-chip, pero no se tratan en este post
    • Está dividido en 4 cuadrantes, y cada cuadrante está compuesto por un warp scheduler, 512 vector registers (cada uno con 32 palabras de 4 bytes), un tensor core para multiplicación de matrices, e instrucciones integradas que soportan operaciones paralelas como suma/multiplicación

Consejos para optimizar el rendimiento de la GPU H100

  • Es importante aprovechar al máximo los Tensor Cores. La H100 tiene un rendimiento de 989 TFLOPs en multiplicación de matrices FP16, por lo que el nivel de uso de los Tensor Cores influye fuertemente en la utilización total de la GPU
  • Sin embargo, para aprovechar al máximo los Tensor Cores, hay que considerar lo siguiente:
    • Es indispensable usar la instrucción Warp Group Matrix Multiply Accumulate (WGMMA), pero su uso es complicado
    • La memoria compartida no es tan rápida como podría parecer, y requiere mucho cuidado al usarla
    • La generación de direcciones tiene un costo alto, por lo que debe optimizarse usando Tensor Memory Accelerator (TMA), entre otros recursos
    • Aumentar la occupancy sigue siendo útil, y los registers son el recurso principal
  • Estas características se aplican en cierta medida no solo a la H100 sino también a otras GPU, pero en la H100 el aprovechamiento de los Tensor Cores es particularmente importante y difícil

ThunderKittens: un DSL embebido de CUDA optimizado para la H100

  • DSL embebido basado en CUDA desarrollado para extraer al máximo el rendimiento de GPU modernas como la NVIDIA H100
  • Ofrece los siguientes 4 tipos de plantillas basadas en tiles:
    • Register Tile (tensor 2D almacenado en registros)
    • Register Vector (tensor 1D almacenado en registros)
    • Shared Tile (tensor 2D almacenado en memoria compartida)
    • Shared Vector (tensor 1D almacenado en memoria compartida)
  • También ofrece varios operadores para manipulación de tiles (exp, mul, sum, etc.) a nivel de warp o grupo de warps
  • Al implementar los kernels de Flash Attention y Flash Attention 2 con ThunderKittens, el código se simplifica mucho y se obtiene hasta 30% más rendimiento en la H100
  • También se implementó el kernel de Based Linear Attention con ThunderKittens, logrando 215 TFLOPs de rendimiento (por las características del algoritmo, si se incluye recompute supera los 300 TFLOPs)

Reflexión desde una perspectiva filosófica

  • ThunderKittens funciona bien porque no intenta soportarlo todo, sino que ofrece una abstracción simple basada en tiles que encaja muy bien con la arquitectura de la GPU
  • Usar vector registers de 1024 bits en lugar de las palabras tradicionales de 32 bits también es un avance, pero hace falta un cambio de paradigma para ver los tiles de 16x16 como la unidad del registro
  • Dado que las cargas de trabajo de IA al final se centran en multiplicación de matrices, reduction y reshape, el enfoque basado en tiles tiene sentido, y desde la perspectiva del hardware también es probable que evolucione en una dirección que soporte pequeñas multiplicaciones de matrices además de los systolic arrays
  • Más allá de eso, hace falta cambiar la forma de pensar hacia el diseño de algoritmos de IA en formas optimizadas para el hardware. Por ejemplo, limitar el tamaño del estado de una RNN a lo que quepa dentro de un SM, y ajustar también la densidad de cómputo al nivel que exige el hardware

Opinión de GN⁺

  • ThunderKittens puede ser una opción atractiva para desarrolladores familiarizados con CUDA, ya que permite obtener mejoras de rendimiento fácilmente sin modificar demasiado el código existente de los kernels
  • Aun así, para principiantes la barrera de entrada puede seguir siendo alta. Harían falta más ejemplos de código y materiales de aprendizaje en el futuro
  • Resulta interesante el plan de ampliar el soporte de ThunderKittens no solo para la H100 sino también para otras GPU como AMD. Podría contribuir a reducir la dependencia de proveedor
  • En última instancia, un punto muy importante es diseñar los propios modelos/algoritmos de IA en formas optimizadas para el hardware. Para ello, hace falta primero una comprensión profunda de las características del hardware, y ThunderKittens puede ayudar a los desarrolladores a obtener ese tipo de insights
  • Se espera que la investigación y desarrollo continuos de Hazy Research, así como sus contribuciones open source, ayuden mucho a dinamizar el ecosistema de CUDA. Aun así, a largo plazo también parece necesario que surjan frameworks con un nivel de abstracción más alto

1 comentarios

 
GN⁺ 2024-05-13
Comentarios en Hacker News
  • Los requisitos del hardware para IA se están volviendo cada vez más claros. Las GPU fueron diseñadas originalmente para otros propósitos, pero se usan para IA porque tienen buen hardware para multiplicación de matrices. Una "GPU para IA" podría excluir algunas funciones de una GPU real, y la tendencia apunta a números más cortos (coma flotante de 16 bits, 8 bits, 2 bits y 1 bit). Este artículo sugiere que el hardware que prefiere mosaicos de 16x16 tiene muchas ventajas.

  • Se necesitan coprocesadores dedicados para IA (NPU), especialmente en sistemas de escritorio prosumer para desarrolladores, especialistas y gamers. Las GPU funcionan en entornos empresariales, pero son incómodas para usar en IA desde la perspectiva de la computación personal. En particular, el límite de VRAM y la ausencia de una API estándar abierta fuera de Vulkan son un problema.

  • Para avanzar en la investigación de IA, hay que estudiar mejor la neurociencia y la psicología. Además, también podrían estar relacionadas cosas vinculadas con la topología de grafos de las redes neuronales.

  • ¿Es esto una versión de CUTLASS más amigable para el usuario?

  • La mascota de ThunderKittens tiene una vibra de gato / Sony Aibo. Parece estar bien generada por IA.

  • Si la computación básica universal (UBC) se considerara como sustituto del ingreso básico universal, sería un futuro muy distópico. Imaginen que una sola empresa como Nvidia fabrica toda la computación.