BurTorch: Revisiting Training from First Principles by Coupling Autodiff, Math Optimization, and Systems

在本工作中，我们介绍了 BurTorch，这是一个紧凑且高性能的框架，旨在通过高效的基于 CPU 的反向传播（Rumelhart 等，1986；Linnainmaa，1970）实现来优化单节点工作站上的深度学习（DL）训练。尽管现代 DL 框架内部依赖于类似编译器的优化技术，BurTorch 走了一条不同的路径。它采用了极简主义设计，并证明在这种情况下，经典的编译型编程语言可以在深度学习研究中发挥重要作用。通过消除大型框架的开销并做出高效的设计选择，BurTorch 在使用 CPU 计算 $\nabla f(x)$ 时实现了性能和内存效率的数量级提升。 BurTorch 的代码库设计紧凑，旨在同时实现两个关键目标。首先，它提供了类似于脚本编程环境的用户体验。其次，它极大地减少了运行时开销。在大型 DL 框架中，对于相对较小的计算图 $f(x)$，内存开销的主要来源是功能丰富的实现方式。我们对 BurTorch 进行了基准测试，对比了广泛使用的 DL 框架及其执行模式：JAX（Bradbury 等，2018）、PyTorch（Paszke 等，2019）、TensorFlow（Abadi 等，2016），以及几个独立的库：Autograd（Maclaurin 等，2015）、Micrograd（Karpathy，2020）、Apple MLX（Hannun 等，2023）。对于小型计算图，BurTorch 在运行时间上比最佳实践解决方案快至 $\times 2000$ 倍，并将内存消耗减少至多 $\times 3500$ 倍。对于一个微型化的 GPT-3 模型（Brown 等，2020），BurTorch 相较于 PyTorch 实现了高达 $\times 20$ 的加速，并将内存使用减少至多 $\times 80$ 倍。

该论文试图解决在单节点工作站上优化深度学习训练性能的问题，特别是在计算较小的计算图时。这是一个长期存在的问题，但论文专注于通过高效的CPU实现来减少运行时间和内存消耗，这为现有框架提供了一种替代方案。

BurTorch采用极简设计，使用经典的编译型编程语言实现高效的反向传播算法，从而显著减少运行时开销和内存占用。相比现代深度学习框架依赖内部编译器优化的方式，BurTorch通过消除大框架的冗余功能和优化实现选择，实现了数量级上的性能提升。

实验表明，BurTorch在小型计算图上比最佳实践解决方案快2000倍，并将内存消耗降低3500倍。对于迷你版GPT-3模型，它比PyTorch快20倍，内存使用减少80倍。论文与JAX、PyTorch、TensorFlow等主流框架进行了基准测试对比，验证了其优越性。代码是否开源未明确提及，但其设计理念值得进一步研究，尤其是在资源受限环境下的应用。

相关工作包括：1) JAX（Bradbury et al., 2018），支持自动微分和高效计算；2) PyTorch（Paszke et al., 2019），以动态计算图为特色；3) TensorFlow（Abadi et al., 2016），广泛应用于工业界；4) Autograd（Maclaurin et al., 2015），早期的自动微分库；5) Micrograd（Karpathy, 2020），一个简单的教育性微分框架；6) Apple MLX（Hannun et al., 2023），苹果推出的机器学习工具包。这些框架主要依赖GPU加速或复杂的优化技术，而BurTorch则专注于CPU优化。