Relax: Samenvoegbare Abstracties voor End-to-End Dynamisch Machine Learning

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

November 1, 2023

Auteurs: 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

Samenvatting

Dynamische vormberekeningen zijn cruciaal geworden in moderne machine learning-workloads, vooral in opkomende grote taalmodellen. Het succes van deze modellen heeft de vraag aangewakkerd om ze te implementeren in een diverse set van backend-omgevingen. In dit artikel presenteren we Relax, een compilerabstractie voor het optimaliseren van end-to-end dynamische machine learning-workloads. Relax introduceert eersteklas symbolische vormannotaties om dynamische vormberekeningen globaal in het programma te volgen. Het introduceert ook een abstractie op meerdere niveaus die computergrafieken, tensorprogramma's op loopniveau en bibliotheekaanroepen in één enkele representatie omvat, om optimalisaties op meerdere niveaus mogelijk te maken. We bouwen een end-to-end compilatiekader met behulp van de voorgestelde aanpak om dynamische vormmodellen te optimaliseren. Experimentele resultaten op grote taalmodellen laten zien dat Relax prestaties levert die competitief zijn met state-of-the-art handmatig geoptimaliseerde systemen op verschillende platforms, en het mogelijk maakt om opkomende dynamische modellen te implementeren in een bredere set van omgevingen, waaronder mobiele telefoons, embedded apparaten en webbrowsers.

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.