Relax: Komponierbare Abstraktionen für End-to-End dynamisches maschinelles Lernen

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

November 1, 2023

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

Zusammenfassung

Dynamische Formberechnungen sind in modernen Machine-Learning-Workloads, insbesondere bei aufkommenden großen Sprachmodellen, von entscheidender Bedeutung. Der Erfolg dieser Modelle hat die Nachfrage nach ihrer Bereitstellung in einer Vielzahl von Backend-Umgebungen gesteigert. In diesem Artikel stellen wir Relax vor, eine Compiler-Abstraktion zur Optimierung von end-to-end dynamischen Machine-Learning-Workloads. Relax führt symbolische Formannotationen erster Klasse ein, um dynamische Formberechnungen global über das Programm hinweg zu verfolgen. Es führt außerdem eine übergreifende Abstraktion ein, die Berechnungsgraphen, Tensorprogramme auf Schleifenebene und Bibliotheksaufrufe in einer einzigen Repräsentation zusammenfasst, um übergreifende Optimierungen zu ermöglichen. Wir entwickeln ein end-to-end Compiler-Framework, das den vorgeschlagenen Ansatz nutzt, um Modelle mit dynamischen Formen zu optimieren. Experimentelle Ergebnisse an großen Sprachmodellen zeigen, dass Relax eine Leistung erzielt, die mit state-of-the-art handoptimierten Systemen über verschiedene Plattformen hinweg konkurrieren kann, und die Bereitstellung von aufkommenden dynamischen Modellen in einer breiteren Palette von Umgebungen, einschließlich Mobiltelefonen, eingebetteten Geräten und Webbrowsern, ermöglicht.

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.