arXiv:2509.02197cs.LGcs.PF2025-09被引 1

无需修改代码,实现机器学习与科学计算的高效自动微分。

DaCe AD: Unifying High-Performance Automatic Differentiation for Machine Learning and Scientific Computing

  • 基于整数线性规划优化存储与重算的权衡,提升性能。
  • 在NPBench测试中平均比JAX快92倍以上,且无需代码改动。
  • 适合需要高性能微分的科研人员和高性能计算开发者。

自动微分(AD)是一类系统应用链式法则计算函数梯度的技术,无需人工干预。尽管其基础理论已确立多年,但近年来因在机器学习反向传播中的关键作用而迎来复兴。它对科学计算领域同样重要,尤其在机器学习模型嵌入科学模拟的新技术中。现有AD框架存在四大局限:编程语言支持有限、需修改代码以兼容AD、科学计算代码性能差、前向数据采用全存储策略导致内存开销大。这些限制迫使领域科学家手动推导梯度。本文提出DaCe AD,一种无需代码修改的通用高效自动微分引擎。其采用新型基于整数线性规划(ILP)的算法,在给定内存约束下优化存储与重算的权衡,实现最佳性能。我们在包含多样化科学计算模式的高性能计算基准套件NPBench上验证了方法的通用性,结果表明,相比具有领先通用性能力的Python框架JAX,DaCe AD在无需任何代码更改的情况下,平均提速超过92倍。

原文摘要 · Abstract (English)

Automatic differentiation (AD) is a set of techniques that systematically applies the chain rule to compute the gradients of functions without requiring human intervention. Although the fundamentals of this technology were established decades ago, it is experiencing a renaissance as it plays a key role in efficiently computing gradients for backpropagation in machine learning algorithms. AD is also crucial for many applications in scientific computing domains, particularly emerging techniques that integrate machine learning models within scientific simulations and schemes. Existing AD frameworks have four main limitations: limited support of programming languages, requiring code modifications for AD compatibility, limited performance on scientific computing codes, and a naive store-all solution for forward-pass data required for gradient calculations. These limitations force domain scientists to manually compute the gradients for large problems. This work presents DaCe AD, a general, efficient automatic differentiation engine that requires no code modifications. DaCe AD uses a novel ILP-based algorithm to optimize the trade-off between storing and recomputing to achieve maximum performance within a given memory constraint. We showcase the generality of our method by applying it to NPBench, a suite of HPC benchmarks with diverse scientific computing patterns, where we outperform JAX, a Python framework with state-of-the-art general AD capabilities, by more than 92 times on average without requiring any code changes.

自动微分高性能计算科学计算代码无侵入

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