Relax: Astrazioni Componibili per l'Apprendimento Automatico Dinamico End-to-End

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

November 1, 2023

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

Abstract

I calcoli dinamici delle forme sono diventati cruciali nei carichi di lavoro del machine learning moderno, in particolare nei modelli linguistici di grandi dimensioni emergenti. Il successo di questi modelli ha aumentato la richiesta di distribuirli in un insieme diversificato di ambienti backend. In questo articolo, presentiamo Relax, un'astrazione del compilatore per ottimizzare i carichi di lavoro di machine learning dinamici end-to-end. Relax introduce annotazioni simboliche di forma di prima classe per tracciare i calcoli dinamici delle forme in modo globale attraverso il programma. Introduce inoltre un'astrazione cross-level che incapsula grafi computazionali, programmi tensoriali a livello di ciclo e chiamate di libreria in una singola rappresentazione, consentendo ottimizzazioni cross-level. Abbiamo costruito un framework di compilazione end-to-end utilizzando l'approccio proposto per ottimizzare i modelli con forme dinamiche. I risultati sperimentali sui modelli linguistici di grandi dimensioni mostrano che Relax offre prestazioni competitive rispetto ai sistemi ottimizzati manualmente all'avanguardia su diverse piattaforme e consente la distribuzione di modelli dinamici emergenti in un insieme più ampio di ambienti, inclusi telefoni cellulari, dispositivi embedded e browser 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.