Relax: Компонуемые абстракции для сквозного динамического машинного обучения

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

November 1, 2023

Авторы: 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

Аннотация

Динамические вычисления форм стали критически важными в современных задачах машинного обучения, особенно в новых крупных языковых моделях. Успех этих моделей стимулировал спрос на их развертывание в разнообразных средах выполнения. В данной статье мы представляем Relax — абстракцию компилятора для оптимизации сквозных динамических задач машинного обучения. Relax вводит аннотации символических форм первого класса для глобального отслеживания динамических вычислений форм в программе. Также он предлагает кросс-уровневую абстракцию, которая инкапсулирует вычислительные графы, тензорные программы на уровне циклов и вызовы библиотек в единое представление, что позволяет выполнять кросс-уровневую оптимизацию. Мы разработали сквозной фреймворк компиляции с использованием предложенного подхода для оптимизации моделей с динамическими формами. Экспериментальные результаты на крупных языковых моделях показывают, что Relax демонстрирует производительность, сопоставимую с современными системами, оптимизированными вручную, на различных платформах, и позволяет развертывать новые динамические модели в более широком спектре сред, включая мобильные устройства, встроенные системы и веб-браузеры.

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.