OpenAI lanza Triton, lenguaje de programación basado en Python para la optimización de cargas de trabajo de IA

OpenAI lanzó su idioma más nuevo, Triton . Este lenguaje de programación de código abierto que permite a los investigadores escribir código GPU altamente eficiente para cargas de trabajo de IA es compatible con Python y permite a los nuevos usuarios lograr resultados de calidad experta en solo 25 líneas de código.

Triton usa Python como base. El desarrollador escribe código en Python utilizando las bibliotecas de Triton, que luego se compilan con JIT para ejecutarse en la GPU. Esto permite la integración con el resto del ecosistema de Python, actualmente el mayor destino para el desarrollo de soluciones de aprendizaje automático.

Las bibliotecas de Triton , que recuerdan a NumPy , proporcionan una variedad de operaciones matriciales, por ejemplo, o funciones que realizan reducciones en matrices de acuerdo con algún criterio. El usuario combina estas primitivas en su propio código, agregando el decorador @ triton.jit compilado para ejecutarse en la GPU. En este sentido, Triton también se parece a Numba , el proyecto que permite que el código Python numéricamente intensivo sea compilado con JIT en un ensamblaje nativo de la máquina para mayor velocidad.

Ejemplos simples de Triton en funcionamiento incluyen un kernel de adición de vectores y una operación softmax fusionada . El último ejemplo puede ejecutarse muchas veces más rápido que el softmax fusionado de PyTorch nativo para operaciones que se pueden realizar completamente en la memoria de la GPU.

Según OpenAI, Triton simplifica el desarrollo de núcleos especializados que pueden ser mucho más rápidos que los de las bibliotecas de uso general. Su compilador simplifica el código y lo optimiza y paraleliza automáticamente, convirtiéndolo en código para su ejecución en GPU Nvidia recientes.

La arquitectura de las GPU modernas se puede dividir en tres componentes principales: DRAM , SRAM y ALU . Cada uno debe tenerse en cuenta al optimizar el código CUDA; no se pueden pasar por alto los desafíos que conlleva la programación de GPU, como asegurarse de que las transferencias de memoria de DRAM se fusionen para aprovechar anchos de bus más grandes en las interfaces de memoria actuales. Los datos deben almacenarse manualmente en SRAMS antes de volver a utilizarlos para que no entren en conflicto con otros bancos de memoria compartida al recuperarlos.

OpenAI afirma que esto hace posible alcanzar el máximo rendimiento del hardware sin mucho esfuerzo, lo que hace que la creación de flujos de trabajo más complejos sea más fácil que nunca. Siempre que los datos en forma de caché son demasiado grandes para el proceso, Triton transfiere casi el doble de la cantidad de memoria necesaria, lo que la hace más rápida que las demás. Codificar en Triton es más fácil y confiable que el kernel CUDA de PyTorch y fácil de entender y mantener el proceso.

Triton tiene la intención de convertirse en un proyecto impulsado por la comunidad, y el repositorio para bifurcar está disponible en GitHub .

Fuente: https://www.infoq.com/news/2021/08/openAI-triton/