Relax: Abstracciones Componibles para Aprendizaje Automático Dinámico de Extremo a Extremo

Relax: Composable Abstractions for End-to-End Dynamic Machine Learning

November 1, 2023

Autores: Ruihang Lai, Junru Shao, Siyuan Feng, Steven S. Lyubomirsky, Bohan Hou, Wuwei Lin, Zihao Ye, Hongyi Jin, Yuchen Jin, Jiawei Liu, Lesheng Jin, Yaxing Cai, Ziheng Jiang, Yong Wu, Sunghyun Park, Prakalp Srivastava, Jared G. Roesch, Todd C. Mowry, Tianqi Chen

cs.AI

Resumen

Los cálculos de formas dinámicas se han vuelto críticos en las cargas de trabajo modernas de aprendizaje automático, especialmente en los emergentes modelos de lenguaje a gran escala. El éxito de estos modelos ha impulsado la demanda de implementarlos en un conjunto diverso de entornos de backend. En este artículo, presentamos Relax, una abstracción de compilador para optimizar cargas de trabajo de aprendizaje automático dinámico de extremo a extremo. Relax introduce anotaciones simbólicas de formas de primera clase para rastrear globalmente los cálculos de formas dinámicas a lo largo del programa. También introduce una abstracción de nivel cruzado que encapsula gráficos computacionales, programas tensoriales a nivel de bucle y llamadas a bibliotecas en una única representación para habilitar optimizaciones de nivel cruzado. Construimos un marco de compilación de extremo a extremo utilizando el enfoque propuesto para optimizar modelos de formas dinámicas. Los resultados experimentales en modelos de lenguaje a gran escala muestran que Relax ofrece un rendimiento competitivo con los sistemas optimizados manualmente más avanzados en diversas plataformas y permite la implementación de modelos dinámicos emergentes en un conjunto más amplio de entornos, incluyendo teléfonos móviles, dispositivos integrados y navegadores web.

English

Dynamic shape computations have become critical in modern machine learning workloads, especially in emerging large language models. The success of these models has driven demand for deploying them to a diverse set of backend environments. In this paper, we present Relax, a compiler abstraction for optimizing end-to-end dynamic machine learning workloads. Relax introduces first-class symbolic shape annotations to track dynamic shape computations globally across the program. It also introduces a cross-level abstraction that encapsulates computational graphs, loop-level tensor programs, and library calls in a single representation to enable cross-level optimizations. We build an end-to-end compilation framework using the proposed approach to optimize dynamic shape models. Experimental results on large language models show that Relax delivers performance competitive with state-of-the-art hand-optimized systems across platforms and enables deployment of emerging dynamic models to a broader set of environments, including mobile phones, embedded devices, and web browsers.