Show HN: Intel Desk, una mesa geopolítica que califica fuentes y rastrea Hormuz

La serie "Profiling en PyTorch" tiene como objetivo hacer que los lectores se sientan cómodos leyendo trazas de perfil y tablas. En la tercera parte de esta serie, se explora el algoritmo de atención, que es fundamental en la arquitectura Transformer.

El algoritmo de atención funciona con consultas, claves y valores, y su interacción se puede escribir como una secuencia de pasos. Estos incluyen calcular las puntuaciones de atención, aplicar una máscara causal y normalizar las puntuaciones con la función softmax. El resultado es un conjunto de pesos de atención que se utilizan para reponderar los valores.

“El perfil muestra las operaciones discretas, como el kernel de multiplicación de matrices, el kernel de multiplicación, la operación de enmascaramiento y el kernel de softmax”

Se ha implementado un módulo de atención naíf en PyTorch y se ha perfilado. El perfil muestra las operaciones discretas, como el kernel de multiplicación de matrices, el kernel de multiplicación, la operación de enmascaramiento y el kernel de softmax. Al desplegar la pista de GPU, se pueden ver los kernels lanzados, incluyendo la multiplicación de matrices, la multiplicación, la copia de memoria y la función softmax.

La atención es un algoritmo con complejidad cuadrática en el tiempo, pero existen trucos para mitigar este problema y hacerlo más rápido. El objetivo de este artículo es mostrar cómo cada uno de estos trucos se ve diferente bajo el perfilador, en lugar de cubrir todos los detalles. Los scripts utilizados en este artículo están disponibles en línea y se pueden ejecutar en una infraestructura de Hugging Face.

En resumen, el algoritmo de atención es una parte crucial de la arquitectura Transformer, y comprender cómo funciona bajo el perfilador es importante para optimizar su rendimiento. Al analizar el perfil de un módulo de atención naíf, se pueden identificar las operaciones clave y entender cómo se lanzan los kernels en la GPU. Esto puede ayudar a los desarrolladores a optimizar su código y mejorar el rendimiento de sus modelos de aprendizaje automático.