用随机线性代数技术,让深度模型更省内存、更快运行
Panther: Faster and Cheaper Computations with Randomized Numerical Linear Algebra
- 将随机线性代数算法整合为可直接替换PyTorch组件的高性能库
- 在BERT上替换线性层后内存减少75%且损失基本不变
- 支持CPU/GPU,仅需少量代码即可接入现有模型
现代深度学习模型训练正受限于GPU内存与计算资源。尽管随机数值线性代数(RandNLA)提供了有效的模型压缩方法,但缺乏统一的生产级库阻碍了其广泛应用。本文提出Panther,一个兼容PyTorch的库,将成熟的RandNLA算法整合为高性能框架。Panther提供高效、即插即用的替代组件,包括草图化线性层、2D卷积、多头注意力及随机矩阵分解(如选主元的CholeskyQR)。通过自研的C++/CUDA后端(pawX),Panther实现跨CPU/GPU优化。实验表明,仅需几行代码替换标准PyTorch线性层,即可在BERT上实现最高75%的内存节省,同时保持相近损失。源码已开源(MIT许可),附演示视频链接。
原文摘要 · Abstract (English)
Training modern deep learning models is increasingly constrained by GPU memory and compute limits. While Randomized Numerical Linear Algebra (RandNLA) offers proven techniques to compress these models, the lack of a unified, production-grade library prevents widely adopting these methods. We present Panther, a PyTorch-compatible library that consolidates established RandNLA algorithms into a single high-performance framework. Panther engineers efficient, drop-in replacements for standard components including sketched linear layers, 2D convolution, multi-head attention, and randomized matrix decompositions (such as pivoted CholeskyQR). By implementing a custom C++/CUDA backend (pawX), Panther provides an optimized implementation that can run on both CPUs and GPUs. We demonstrate the effectiveness of RandNLA techniques and Panther's ease of adoption. By replacing standard PyTorch linear layers with Panther layers (requiring only a few lines of code) we achieve significant memory savings (up to 75%) on BERT while maintaining comparable loss. Source code is available (MIT License) at https://github.com/FahdSeddik/panther, along with demonstration video at https://youtu.be/7M3RQb4KWxs.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。