用极简代码实现超快反向传播,让单机CPU高效训练深度模型。
BurTorch: Revisiting Training from First Principles by Coupling Autodiff, Math Optimization, and Systems
- 用传统编译语言重构反向传播,去除大框架冗余开销。
- 小模型上速度比主流框架快2000倍,内存降低3500倍。
- 适合想深入理解训练机制或资源受限环境的开发者。
本文提出BurTorch,一个轻量级高性能框架,通过高效的纯CPU反向传播实现深度学习训练优化。相比现代框架依赖复杂编译优化,BurTorch采用极简设计,摒弃大框架的高开销,利用经典编译语言优势,在小型计算图上实现反向传播性能提升数个数量级。在多个基准测试中,其运行时间较最佳实践方案最高快2000倍,内存消耗降低至原系统的1/3500。针对微型GPT-3模型,相较于PyTorch,BurTorch实现最高20倍加速与80倍内存节省。框架兼顾脚本式编程体验与极低运行开销,凸显传统编程语言在深度学习研究中的价值。
原文摘要 · Abstract (English)
In this work, we introduce BurTorch, a compact high-performance framework designed to optimize Deep Learning (DL) training on single-node workstations through an exceptionally efficient CPU-based backpropagation (Rumelhart et al., 1986; Linnainmaa, 1970) implementation. Although modern DL frameworks rely on compilerlike optimizations internally, BurTorch takes a different path. It adopts a minimalist design and demonstrates that, in these circumstances, classical compiled programming languages can play a significant role in DL research. By eliminating the overhead of large frameworks and making efficient implementation choices, BurTorch achieves orders-of-magnitude improvements in performance and memory efficiency when computing $\nabla f(x)$ on a CPU. BurTorch features a compact codebase designed to achieve two key goals simultaneously. First, it provides a user experience similar to script-based programming environments. Second, it dramatically minimizes runtime overheads. In large DL frameworks, the primary source of memory overhead for relatively small computation graphs $f(x)$ is due to feature-heavy implementations. We benchmarked BurTorch against widely used DL frameworks in their execution modes: JAX (Bradbury et al., 2018), PyTorch (Paszke et al., 2019), TensorFlow (Abadi et al., 2016); and several standalone libraries: Autograd (Maclaurin et al., 2015), Micrograd (Karpathy, 2020), Apple MLX (Hannun et al., 2023). For small compute graphs, BurTorch outperforms best-practice solutions by up to $\times 2000$ in runtime and reduces memory consumption by up to $\times 3500$. For a miniaturized GPT-3 model (Brown et al., 2020), BurTorch achieves up to a $\times 20$ speedup and reduces memory up to $\times 80$ compared to PyTorch.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。