BurTorch：通过耦合自动求导、数学优化和系统方法重访深度学习训练

在本工作中，我们引入 BurTorch，一个紧凑高性能框架，旨在通过异常高效的 CPU 基于反向传播实现 (Rumelhart 等，1986；Linnainmaa，1970) 优化单节点工作站上的深度学习 (DL) 训练。虽然现代 DL 框架在内部依赖类编译器优化，但 BurTorch 采取了不同路径。它采用极简设计，展示了在这种情况下，经典编译编程语言在 DL 研究中仍发挥重要作用。通过消除大型框架的开销并作出高效实现选择，BurTorch 在计算 $\nabla f(x)$ 时在 CPU 上实现了数量级的性能和内存效率提升。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 在运行时性能上提升最高可达 $\times 20$，内存消耗降低最高可达 $\times 80$。