arXiv:2501.05408cs.DCcs.AI2025-01被引 3

让动态深度学习既有灵活执行,又能编译优化。

Tempo: Compiled Dynamic Deep Learning with Symbolic Dependence Graphs

  • 用带时间维度的张量和符号表达式定义动态依赖
  • 实现7倍于JAX的Llama解码速度,强化学习快54倍
  • 适合需要动态时序计算的模型开发与性能优化

深度学习算法常涉及时间依赖:某一时刻的张量可能依赖于前或后时刻的张量。这种动态依赖(及对应动态张量形状)难以表达与优化:虽然急切执行系统支持动态性,但无法应用编译优化;图执行系统要求静态张量形状,迫使用户填充张量或拆分程序为多个静态图。本文提出Tempo,一种结合急切执行灵活性与图编译全程序优化的新系统。Tempo通过包含显式时间维度的递归张量实现声明式编程,利用符号表达式索引时间维度以表达对过去和未来张量的动态依赖。基于此,Tempo构建符号依赖图,精炼编码算子间的动态依赖关系,并应用代数简化、向量化、分块与融合等全程序优化。通过将动态依赖分块为静态大小单元,Tempo可复用现有静态代码生成器。随后采用多面体模型求解可行执行调度,包括内存管理操作。实验表明,Tempo在Llama-3.2-3B解码上比JAX快7倍;在强化学习算法中,速度提升54倍,峰值内存降低16倍。

原文摘要 · Abstract (English)

Deep learning (DL) algorithms are often defined in terms of temporal relationships: a tensor at one timestep may depend on tensors from earlier or later timesteps. Such dynamic dependencies (and corresponding dynamic tensor shapes) are difficult to express and optimize: while eager DL systems support such dynamism, they cannot apply compiler-based optimizations; graph-based systems require static tensor shapes, which forces users to pad tensors or break-up programs into multiple static graphs. We describe Tempo, a new DL system that combines the dynamism of eager execution with the whole-program optimizations of graph-based compilation. Tempo achieves this through a declarative programming model with recurrent tensors, which include explicit temporal dimensions. Temporal dimensions can be indexed using symbolic expressions to express dynamic dependencies on past and future tensors. Based on this, Tempo constructs a symbolic dependence graph, which concisely encodes dynamic dependencies between operators, and applies whole-program optimizations, such as algebraic simplifications, vectorization, tiling, and fusion. By tiling dynamic dependencies into static-size blocks, Tempo can also reuse existing static code-generators. It then uses a polyhedral model to find a feasible execution schedule, which includes memory management operations. We show that Tempo achieves a 7$\times$ speedup over JAX for Llama-3.2-3B decoding; for reinforcement learning algorithms, Tempo achieves a 54$\times$ speedup, with 16$\times$ lower peak memory usage.

动态深度学习编译优化符号依赖图张量计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。