针对小内存显卡优化大模型推理,实现三倍压缩与三倍加速。
Memory-Efficient Acceleration of Block Low-Rank Foundation Models on Resource Constrained GPUs
- 设计自定义Triton内核,融合计算与内存布局优化
- 在Jetson Orin Nano等设备上实现3.76倍加速和3倍模型压缩
- 支持Llama、GPT2、DiT等多种主流模型,适合边缘部署
基于Transformer的大型基础模型日益普及,但其规模膨胀导致单个GPU难以容纳,计算成本高昂。块低秩(BLR)压缩技术通过学习权重矩阵的紧凑表示来缓解此问题。相比传统低秩方法易造成精度骤降,Monarch与BLAST等BLR方法能更好捕捉结构特征,从而在减少计算与内存占用的同时保持精度。本文通过屋顶分析发现,尽管BLR在单标记推理中具备理论优势与实际提速,多标记推理常因内存瓶颈导致延迟上升,即便使用PyTorch编译优化也难改善。为此,我们为Monarch与BLAST开发了定制Triton内核,结合部分融合与内存布局优化。在如Jetson Orin Nano和A40等内存受限的NVIDIA GPU上,相比使用CUDA后端的PyTorch密集基线,本方法实现最高3.76倍加速与3倍模型压缩,同时支持Llama-7B/1B、GPT2-S、DiT-XL/2和ViT-B等多种模型。代码已开源:https://github.com/pabillam/mem-efficient-blr。
原文摘要 · Abstract (English)
Recent advances in transformer-based foundation models have made them the default choice for many tasks, but their rapidly growing size makes fitting a full model on a single GPU increasingly difficult and their computational cost prohibitive. Block low-rank (BLR) compression techniques address this challenge by learning compact representations of weight matrices. While traditional low-rank (LR) methods often incur sharp accuracy drops, BLR approaches such as Monarch and BLAST can better capture the underlying structure, thus preserving accuracy while reducing computations and memory footprints. In this work, we use roofline analysis to show that, although BLR methods achieve theoretical savings and practical speedups for single-token inference, multi-token inference often becomes memory-bound in practice, increasing latency despite compiler-level optimizations in PyTorch. To address this, we introduce custom Triton kernels with partial fusion and memory layout optimizations for both Monarch and BLAST. On memory-constrained NVIDIA GPUs such as Jetson Orin Nano and A40, our kernels deliver up to $3.76\times$ speedups and $3\times$ model size compression over PyTorch dense baselines using CUDA backend and compiler-level optimizations, while supporting various models including Llama-7/1B, GPT2-S, DiT-XL/2, and ViT-B. Our code is available at https://github.com/pabillam/mem-efficient-blr.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。