Relax: Composable Abstractions for End-to-End Dynamic Machine Learning
作者: 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.LG, cs.AI, cs.PL
发布日期: 2023-11-01 (更新: 2025-02-07)
备注: To appear at ASPLOS 2025 (16 pages, 20 figures)
💡 一句话要点
提出Relax以优化动态机器学习工作负载
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 动态机器学习 编译器优化 大型语言模型 跨层次抽象 符号形状注释 性能提升 设备部署 智能应用
📋 核心要点
- 现有动态机器学习模型在不同后端环境中的部署存在性能瓶颈，难以满足多样化需求。
- Relax通过引入跨层次抽象和符号形状注释，优化动态形状计算，提升了模型的灵活性和性能。
- 实验结果显示，Relax在多种GPU上与现有系统性能相当，并支持新兴模型在移动设备等环境中的部署。
📝 摘要（中文）
动态形状计算在现代机器学习工作负载中变得至关重要，尤其是在新兴的大型语言模型中。本文提出了Relax，一个用于优化端到端动态机器学习工作负载的编译器抽象。Relax引入了一种跨层次的抽象，能够将计算图、循环级张量程序和外部库调用封装在一个统一的表示中。此外，Relax还引入了首类符号形状注释，以便在程序中全局跟踪动态形状计算，从而实现动态形状感知的跨层次优化。实验结果表明，Relax在各种GPU上表现出与最先进系统相当的性能，并使新兴模型能够在更广泛的环境中部署，包括手机、嵌入式设备和网页浏览器。
🔬 方法详解
问题定义：本文旨在解决动态机器学习工作负载在不同后端环境中部署时的性能瓶颈，现有方法在处理动态形状计算时缺乏有效的优化手段。
核心思路：Relax的核心思路是通过引入跨层次抽象和符号形状注释，统一表示计算图、张量程序和外部库调用，从而实现动态形状感知的优化。
技术框架：Relax的整体架构包括三个主要模块：计算图表示、符号形状注释和跨层次优化策略。计算图表示用于描述模型结构，符号形状注释用于跟踪动态形状，而优化策略则在不同层次间进行协调。
关键创新：Relax的主要创新在于其跨层次抽象和首类符号形状注释，这使得动态形状计算的优化成为可能，显著提升了模型的灵活性和性能。与现有方法相比，Relax能够在更广泛的环境中实现更高效的部署。
关键设计：Relax在设计中采用了符号形状注释来全局跟踪动态形状计算，并通过优化策略在不同层次间进行动态形状感知的优化，确保了模型在多种后端环境中的高效运行。
🖼️ 关键图片
📊 实验亮点
实验结果表明，Relax在多种GPU上实现了与最先进系统相当的性能，特别是在大型语言模型的动态形状计算中，提升幅度显著，能够支持更广泛的设备部署，增强了模型的适应性和灵活性。
🎯 应用场景
该研究的潜在应用领域包括移动设备、嵌入式系统和网页浏览器等新兴环境。通过优化动态机器学习工作负载，Relax能够提升这些设备上大型语言模型的性能，促进智能应用的普及和发展，具有重要的实际价值和未来影响。
📄 摘要（原文）
Dynamic shape computations have become critical in modern machine learning workloads, especially in emerging large language models. The success of these models has driven the demand for their universal deployment across 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 a cross-level abstraction that encapsulates computational graphs, loop-level tensor programs, and external library calls in a single representation. Relax also introduces first-class symbolic shape annotations to track dynamic shape computations globally across the program, enabling dynamic shape-aware cross-level optimizations. We build an end-to-end compilation framework using the proposed approach to optimize dynamic shape models. Experimental results on LLMs show that Relax delivers performance competitive with state-of-the-art systems across various GPUs and enables deployment of emerging models to a broader set of emerging environments, including mobile phones, embedded devices, and web browsers.