Relax: Abstrações Componíveis para Aprendizado de Máquina Dinâmico de Ponta a Ponta

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

Resumo

Cálculos dinâmicos de forma tornaram-se críticos em cargas de trabalho modernas de aprendizado de máquina, especialmente em modelos emergentes de linguagem em grande escala. O sucesso desses modelos tem impulsionado a demanda por sua implantação em um conjunto diversificado de ambientes de backend. Neste artigo, apresentamos o Relax, uma abstração de compilador para otimizar cargas de trabalho dinâmicas de aprendizado de máquina de ponta a ponta. O Relax introduz anotações simbólicas de forma de primeira classe para rastrear cálculos dinâmicos de forma globalmente em todo o programa. Ele também introduz uma abstração de nível cruzado que encapsula grafos computacionais, programas de tensores em nível de loop e chamadas de biblioteca em uma única representação para permitir otimizações de nível cruzado. Construímos uma estrutura de compilação de ponta a ponta usando a abordagem proposta para otimizar modelos de forma dinâmica. Resultados experimentais em modelos de linguagem em grande escala mostram que o Relax oferece desempenho competitivo com sistemas otimizados manualmente de última geração em várias plataformas e permite a implantação de modelos dinâmicos emergentes em um conjunto mais amplo de ambientes, incluindo telefones móveis, dispositivos embarcados e navegadores da 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.