Grass: 利用结构化稀疏梯度实现计算高效的低内存大语言模型训练
机器学习
2024-06-26 v1
摘要
大语言模型(LLM)的训练和微调常常受限于有限的GPU内存。现有的基于投影的优化方法通过将梯度投影到低维子空间来减少优化器状态内存，但它们通常依赖稠密投影矩阵，这可能会引入计算和内存开销。在这项工作中，我们提出了Grass(GRAdient Structured Sparsification，梯度结构化稀疏化)，一种利用稀疏投影将梯度转化为结构化稀疏更新的新颖方法。这种设计不仅显著减少了优化器状态的内存使用，还最小化了梯度内存占用、计算和通信成本，从而带来了显著的吞吐量提升。在预训练和微调任务上的大量实验表明，Grass达到了与全秩训练和现有基于投影的方法相竞争的性能。值得注意的是，Grass能够在单个40GB A100 GPU上对13B参数的LLaMA模型进行半精度预训练——这是先前方法无法实现的——并在8-GPU系统上实现了高达2倍的吞吐量提升。代码可在https://github.com/aashiqmuhamed/GRASS 获取。
引用
@article{arxiv.2406.17660,
title = {Grass: Compute Efficient Low-Memory LLM Training with Structured Sparse Gradients},
author = {Aashiq Muhamed and Oscar Li and David Woodruff and Mona Diab and Virginia Smith},
journal= {arXiv preprint arXiv:2406.17660},
year = {2024}
}