Relax : Abstractions composables pour l'apprentissage automatique dynamique de bout en bout

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

Résumé

Les calculs de formes dynamiques sont devenus essentiels dans les charges de travail modernes d'apprentissage automatique, en particulier dans les modèles de langage émergents à grande échelle. Le succès de ces modèles a accru la demande pour leur déploiement dans un ensemble diversifié d'environnements backend. Dans cet article, nous présentons Relax, une abstraction de compilateur pour optimiser les charges de travail d'apprentissage automatique dynamiques de bout en bout. Relax introduit des annotations symboliques de premier ordre pour suivre les calculs de formes dynamiques de manière globale à travers le programme. Il introduit également une abstraction transversale qui encapsule les graphes de calcul, les programmes tensoriels au niveau des boucles et les appels de bibliothèque dans une représentation unique, permettant ainsi des optimisations transversales. Nous construisons un framework de compilation de bout en bout utilisant cette approche pour optimiser les modèles à formes dynamiques. Les résultats expérimentaux sur les modèles de langage à grande échelle montrent que Relax offre des performances compétitives par rapport aux systèmes optimisés manuellement de pointe sur diverses plateformes, et permet le déploiement de modèles dynamiques émergents dans un ensemble plus large d'environnements, incluant les téléphones mobiles, les appareils embarqués et les navigateurs 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.