Press "Enter" to skip to content

这篇AI论文介绍了Relax:一种用于优化端到端动态机器学习工作负载的编译器抽象

通过使用动态形状优化机器学习模型,可以实现更好的性能和灵活性。动态形状是指模型在运行时处理具有不同尺寸输入数据的能力。用户可以使用支持动态计算图的框架,例如TensorFlow的即时执行或PyTorch来构建能够在运行时适应可变输入尺寸的模型。

优化具有动态形状的机器学习模型面临许多挑战,因为许多传统优化依赖于静态形状分析。动态尺寸的缺失信息可能会对跨运算符和函数执行的优化产生重大影响。具有动态形状的模型需要处理不同的批处理大小。在生产环境中,为不同的批处理大小进行优化比为固定批处理大小进行优化更具挑战性。

当前的机器学习(ML)编译器通常将程序转换为传统的单次转换流程,一次应用一种优化,通常将程序重写为较低级别的表示。这种方法经常导致在抽象层之间丢失形状和附加信息,使跨边界进行增量优化变得更加困难。

研究人员提出了名为Relax的编译器抽象。它具有一流的符号形状注释,可以全局追踪程序中的动态形状计算。它还具有跨级别抽象,将计算图、循环级张量程序和库调用封装在单个表示中,以实现跨级别的优化。它是一个优化动态形状模型的端到端编译框架。

研究人员采用正向推理方法,根据其输入组件推断表达式的注释。正向推理简单且局部化,编译器可以在各个步骤中为临时变量获取注释。此外,当无法自动推断形状时,正向推理可以使用用户插入的匹配转换的结果来继续推断后续注释。

研究人员表示,Relax中的所有优化都是可组合的动态形状感知转换。它通过不同的方法增量优化或部分降低计算的部分。它考虑来自其他级别的分析并结合了假设动态形状关系的进一步优化。

实验结果表明,Relax将新兴的LLM编译并优化到不同的硬件后端,性能竞争力与经过大量优化的平台特定解决方案相当。此外,Relax通过WebAssembly和WebGPU支持广泛的设备和环境中的LLM,包括手机、嵌入式设备和Web浏览器。

Leave a Reply

Your email address will not be published. Required fields are marked *