un poco de optimización es muy útil. Modelos como GPT4 cuestan más de $ 100 millones para entrenar, lo que hace que una ganancia de eficiencia del 1% valga más de un millón de dólares. Una forma poderosa de optimizar la eficiencia de los modelos de aprendizaje automático es escribir algunos de sus componentes directamente en la GPU. Ahora, si eres como yo, la simple mención de los núcleos Cuda es suficiente para enviar escalofríos por la columna vertebral, ya que son notoriamente complejos para escribir y depurar.
Afortunadamente, Operai lanzó Triton en 2021, un nuevo lenguaje y un compilador que abstrae gran parte de la complejidad de Cuda y permite a los profesionales menos experimentados escribir núcleos performantes. Un ejemplo notable es Uncera, un servicio de entrenamiento LLM que promete una capacitación 30x más rápida con un 60% menos de uso de memoria, todo gracias a reemplazar las capas escritas en Pytorch con núcleos Triton.
¡En esta serie tutorial, aprenderemos los conceptos básicos de la arquitectura de GPU y cómo implementar los núcleos Triton de alto rendimiento! Todo el código presentado en esta serie estará disponible en https://github.com/rpegoud/triton-kernels.
Conceptos básicos de arquitectura de GPU
En esta sección, pasaremos por los conceptos básicos de las GPU (NVIDIA) para comenzar y escribir nuestro primer núcleo Triton al final de este artículo.
A partir de la unidad de software más pequeña, podemos describir la jerarquía de las unidades de ejecución de la siguiente manera:
Threads: la unidad de trabajo más pequeña, ejecutan el código del núcleo definido por el usuario. Warps: la unidad de programación más pequeña, siempre están compuestas por 32 hilos paralelos, cada uno con su propio contador de direcciones de instrucciones y estado de registro. Los hilos en una urdimbre comienzan juntos, pero son libres de ramificarse y ejecutar de forma independiente. Bloques de hilos: grupo de urdimbres, donde todos los hilos pueden cooperar a través de la memoria compartida y las barreras de sincronización. Se requiere que los bloques de subprocesos puedan ejecutarse de forma independiente y en cualquier orden, en paralelo o secuencialmente. Esta independencia permite que los bloques de subprocesos se programen en cualquier orden en cualquier número de núcleos, de modo que los programas de GPU escalen de manera eficiente con el número de núcleos. Podemos sincronizar los hilos dentro de un bloque en puntos específicos del núcleo si es necesario, por ejemplo, para sincronizar el acceso a la memoria. Multiprocesador de transmisión (SM): una unidad a cargo de ejecutar muchas deformaciones en paralelo, posee la memoria compartida y un caché L1 (posee las líneas de memoria global más recientes a las que ha accedido el SM). Un SM tiene un planificador de urdimbre dedicado que extrae las deformaciones de los bloques de hilo que están listos para funcionar.
En el lado del hardware, la unidad de trabajo más pequeña es un núcleo CUDA, la unidad de lógica aritmética física (ALU) que realiza operaciones aritméticas para un hilo (o partes).
Para resumir esta sección con una analogía, pudimos ver los núcleos de Cuda como trabajadores individuales, mientras que una urdimbre es un escuadrón de 32 trabajadores que reciben la misma instrucción a la vez. Pueden o no ejecutar esta tarea de la misma manera (ramificación) y potencialmente pueden completarla en un punto diferente en el tiempo (independencia). Un bloque de subprocesos se compone de varios escuadrones que comparten un espacio de trabajo común (es decir, tienen memoria compartida), los trabajadores de todos los escuadrones en el espacio de trabajo pueden esperar unos a otros para almorzar al mismo tiempo. Un multiprocesador de transmisión es un piso de fábrica con muchos escuadrones trabajando juntos y compartiendo herramientas y almacenamiento. Finalmente, la GPU es una planta completa, con muchos pisos.
Conceptos básicos de optimización
Al optimizar modelos de aprendizaje profundo, estamos haciendo malabarismos con tres componentes principales:
Calcule: Tiempo dedicado a las operaciones de punto flotante de computación de GPU (FLOPS). Memoria: tiempo dedicado a transferir tensores dentro de una GPU. Overhead: todas las demás operaciones (intérprete de Python, Pytorch Dispatch, …).
Tener en cuenta esos componentes ayuda a descubrir la forma correcta de resolver un cuello de botella. Por ejemplo, aumentar el cómputo (por ejemplo, usar una GPU más potente) no ayuda si la mayor parte del tiempo se dedica a hacer transferencias de memoria. Idealmente, sin embargo, la mayor parte del tiempo debe dedicarse a la computa, más precisamente en las multiplicaciones de matriz, las GPU de operación precisa están optimizadas.
Esto implica minimizar el costo pagado para mover datos, ya sea desde la CPU hasta la GPU (“costo de transferencia de datos”), de un nodo a otro (“costo de red”) o desde la memoria global de CUDA (DRAM, barato pero lento) a la memoria compartida CUDA (SRAM, memoria más costosa pero más rápida en el dispositivo). El último se llama costos de ancho de banda y será nuestro enfoque principal por ahora. Las estrategias comunes para reducir los costos de ancho de banda incluyen:
Reutilizando datos cargados en la memoria compartida para múltiples pasos. Un excelente ejemplo de esto es la multiplicación de matriz de azulejos, que cubriremos en una publicación futura. Fusionando múltiples operaciones en un solo núcleo (ya que cada lanzamiento del núcleo implica datos móvil de DRAM a SRAM), por ejemplo, podemos fusionar una multiplicación de matriz con una función de activación. En general, el operador de fusión puede proporcionar un aumento masivo de rendimiento, ya que evita muchas lecturas/escrituras de memoria global y cualquiera de los dos operadores presentan una oportunidad para la fusión.
En este ejemplo, realizamos una multiplicación de matriz x@w y almacenamos el resultado en una variable intermedia a. Luego aplicamos un RELU a A y almacenamos el resultado en una variable Y. Esto requiere que la GPU lea de X y W en la memoria global, escriba el resultado en A, lea de A nuevamente y finalmente escriba en y. En su lugar, el operador Fusion nos permitiría reducir a la mitad la cantidad de lecturas y escrituras a la memoria global realizando la multiplicación de matriz y aplicando el Relu en un solo núcleo.
Tritón
Ahora escribiremos nuestro primer núcleo Triton, una simple adición de vectores. Primero, caminemos por cómo esta operación se desglosa y ejecutada en una GPU.
Considere querer sumar las entradas de dos vectores x e y, cada uno con 7 elementos (n_elements = 7).
Instruiremos a la GPU que aborde este problema en trozos de 3 elementos a la vez (block_size = 3). Por lo tanto, para cubrir los 7 elementos de los vectores de entrada, la GPU lanzará 3 “programas” paralelos, instancia independiente de nuestro núcleo, cada uno con una identificación de programa única, PID:
El programa 0 se asigna elementos 0, 1, 2. El programa 1 se asigna elementos 3, 4, 5. El programa 2 está asignado Elemento 6.
Luego, estos programas volverán a escribir los resultados en un vector Z almacenado en la memoria global.
Un detalle importante es que un núcleo no recibe un vector X entero, sino que recibe un puntero a la dirección de memoria del primer elemento, x[0]. Para acceder a los valores reales de X, necesitamos cargarlos de la memoria global manualmente.
Podemos acceder a los datos para cada bloque utilizando la ID del programa: block_start = pid * block_size. A partir de ahí, podemos obtener las direcciones de elementos restantes para ese bloque calculando los compensaciones = block_start + rango (0, block_size) y cargarlas en la memoria.
Sin embargo, recuerde que el programa 2 solo se asigna el elemento 6, pero sus compensaciones son [6, 7, 8]. Para evitar cualquier error de indexación, Triton nos permite definir una máscara para identificar elementos objetivo válidos, aquí Mask = Offsets Ahora podemos cargar de forma segura X e Y y agregarlos antes de escribir el resultado a una variable de salida Z en la memoria global de manera similar. Echemos un vistazo más de cerca al código, aquí está el núcleo Triton: import triton import triton.language as tl @triton.jit def add_kernel( x_ptr, # pointer to the first memory entry of x y_ptr, # pointer to the first memory entry of y output_ptr, # pointer to the first memory entry of the output n_elements, # dimension of x and y BLOCK_SIZE: tl.constexpr, # size of a single block ): # — Compute offsets and mask — pid = tl.program_id(axis=0) # block index block_start = pid * BLOCK_SIZE # start index for current block offsets = block_start + tl.arange(0, BLOCK_SIZE) # index range mask = offsets < n_elements # mask out-of-bound elements # --- Load variables from global memory --- x = tl.load(x_ptr + offsets, mask=mask) y = tl.load (y_ptr + offsets, máscara = máscara) # --- operación --- output = x + y # --- Guardar resultados en la memoria global --- tl.store (pointer = output_ptr + offsets, value = output, mask = Mask)
Desglosemos parte de la sintaxis específica de Triton: Primero, un núcleo Triton siempre está decorado por @triton.jit. En segundo lugar, algunos argumentos deben declararse como estáticos, lo que significa que son conocidos en el tiempo de cálculo. Esto se requiere para block_size y se logra agregar la anotación de tipo tl.constexpr. También tenga en cuenta que no anotamos otras variables, ya que no son variables de pitón adecuadas. Utilizamos TL.Program_ID para acceder a la ID del bloque actual, TL.arange se comporta de manera similar a NP.Arange de Numpy. Las variables de carga y almacenamiento se logra llamando a TL.Load y Tl.Store con matrices de punteros. Observe que no hay una declaración de devolución, este papel se delega a TL.Store. Para usar nuestro núcleo, ahora necesitamos escribir un envoltorio de nivel de pytorch que proporcione punteros de memoria y define una cuadrícula del núcleo. En general, la cuadrícula del núcleo es una tupla 1D, 2D o 3D que contiene el número de bloques de rosca asignados al núcleo a lo largo de cada eje. En nuestro ejemplo anterior, utilizamos una cuadrícula 1D de 3 bloques de hilo: cuadrícula = (3,). Para manejar variables tamaños de matriz, predeterminamos a Grid = (techo (n_elements / block_size),). Def add (x: tortch.tensor, y: torc.tensor) -> torc.tensor: “” “” pytorch wrapper para `add_kernel`.” “” output = torc.zeros_like (x) # Memoria de alivio para la salida n_elements = output.numel () # dimensión de x e # cdiv = ceil div, calcula el número de bloqueo de la salida para usar. Meta: (Triton.cdiv (n_elements, meta[“BLOCK_SIZE”]),) # Llamar al kernel almacenará automáticamente `block_size` en` meta` # y actualizar `salida ‘add_kernel[grid](X, y, salida, n_elements, block_size = 1024) Salida de retorno Aquí hay dos notas finales sobre el envoltorio: Es posible que haya notado que la red se define como una función lambda. Esto le permite a Triton calcular el número de bloques de subprocesos para iniciarse en la hora de lanzamiento. Por lo tanto, calculamos el tamaño de la cuadrícula en función del tamaño del bloque que se almacena en Meta, un diccionario de constantes de tiempo de compilación que están expuestos al núcleo. Al llamar al núcleo, el valor de salida se modificará en el lugar, por lo que no necesitamos reasignar salida = add_kernel[…]. x, y = atorch.randn ((2, 2048), dispositivo = “cuda”) print (add (x, y)) >> tensor ([ 1.8022, 0.6780, 2.8261, …, 1.5445, 0.2563, -0.1846]dispositivo = ‘Cuda: 0’) ABS_DiFference = Torch.abs ((x + y) – ADD (X, Y)) Impresión (F “MAX Diferencia absoluta: {Torch.max (ABS_DIFEFERENT)}”) >> Diferencia absoluta máxima: 0.0 Eso es todo para esta introducción, en las siguientes publicaciones aprenderemos a implementar núcleos más interesantes, como la multiplicación de matriz de mosaico y ver cómo integrar los núcleos Triton en los modelos Pytorch usando Autograd. ¡Hasta la próxima! 👋
Podemos concluir este tutorial verificando que nuestro kernel funciona correctamente:Referencias y recursos útiles