大模型LLM压缩量化和GPU显存管理paper阅读笔记

最新推荐文章于 2024-08-21 19:12:43 发布

不归牛顿管的熊猫

最新推荐文章于 2024-08-21 19:12:43 发布

阅读量443

点赞数

文章标签：笔记论文阅读人工智能性能优化持续部署

本文链接：https://blog.csdn.net/weixin_43568400/article/details/133792662

版权

前言

欢迎大家关注我的公众号“AI不止算法”，不定时分享C++以及AI高性能优化部署的技术体会

不知道大家有没有思考过，GPT如此之大的模型是怎么做到能在这么快的速度下吐出一个个word/token来回复我们的？其实这离不开AI部署的各种手段

算法方向的paper和落地工程方面的paper我个人在学校和工作中都分别读过，我个人觉得落地工程的paper能更给人带来一种参与“工业革命”的快感，贡献独有的工业价值，让我更加入迷。

这里一共两篇paper notes，一篇关于大模型压缩量化，这在LLM工业部署上必不可少，对减少显存占用，提升推理性能非常可观。一篇关于GPU显存利用的优化，可以给你省钱！

大模型量化的研究

第一篇，今年才出的大模型量化技术，title叫做《SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models》

动机
某些LLM的某些channel或某些维度的activation outlier值很多且很大(ep. GLM-130B有30%），导致量化的有效位变少，比如int8本来是-128到127，没有outlier的时候，映射到-128到127的数据分布均匀，占满了8bit位范围，精度损失很低，但是有了outlier之后，多数正常值的分布区间可能在[-20,20]或者[-10,10]，8bit位范围只利用到了5bit，甚至4bit，由此导致精度损失
observation
发现论文里面经常用到observation这个字眼，自我感觉，可能一些创新的想法，不太好找到驱动的源头，就用observation来说明。主要有两个发现，如下图：

- 在这里插入图片描述

不同token的异常分布一致
红色橙色是outlier，outlier一般集中存在于某几个channel或axis
challenge
1.activation比weight更难量化，后者数据分布一般比较均匀
2.outlier导致non-outlier在per tensor级别的量化误差大
3.activation的量化方式不适宜用per channel，weight用per tensor/channel都可以，这是因为activation用perchannel的话很难dequantize
4.很多人也用per token量化，但只是一个workaround，这没有解决act比weight更难量化的问题

提出方法

这里我尽量简单讲

key points：weight比activation更难量化，所以可以把activation量化难度迁移到weight上来。

怎么迁移呢？

把activation里面不均匀的分布用weight中和一下，具体来讲，主要是在fp32阶段做的，保证一个等式的等价，X为输入act，W为weight，s为因子，通过s来中和

在这里插入图片描述

直接上图：左边为smooth前，右边为smooth后，可以明显看到X乘以s^（-1）之后数据分布明显均匀了，把难度匀了一点给weight
在这里插入图片描述
效果
smoothquant和几个baseline的比较，在OPT-175B上，平均精度明显比另外四个int8方法要好

GPU显存优化

第二篇，比较老了，20年的一篇省显存的paper，如今已经有了很多在他之上的工作，title:《SuperNeurons: Dynamic GPU Memory Management for Training Deep Neural Networks》

动机
GPU上训练深层大神经网络通常显存不够，尤其对于现在千变万化的non-linear的网络拓扑模型，主流DL fwk（文中只提到了caffe和mxnet和torch）只有静态的memory reduction tech，由于静态的本质，不能很好的handle non-linear的网络拓扑，因此本文基于以上动机，提出一种dynamic memory solution来handle non-linear的网络拓扑结构

挑战
1.GPU resident mem有限，DNN的特点是高内存需求和高计算密度，因此有了利用多GPU的模型并行训练，把子网络分到不同的GPU里，但是这样造成inter/intra-network 通信的overhead，由此造成perf issue，而data parallism，每个卡计算sub batch的sub gradients，最后聚合来更新整网的gradients ，虽然没有通信overhead，但是也需要model可以fit into GPU DRAM

2.non-linear网络数据依赖类型变种很多，导致依赖关系和执行流复杂，所以需要一个dynamic的memory schedule解决方法来handle这些千变万化的拓扑关系的内存分配，以fit into GPU DRAM，non-linear说白了好像就是join和fan两种形式，前者是两层layer喂到一层，后者是一层分叉输出到多个branch

方法
以下为本文提出的三个内存优化技术，peak memory usage的baseline为 , 第一项为前向过程的各层中间tensor占用的总内存，第二项为后向过程的各层中间tensor占用的总内存。

liveness analysis
先认识一下，训练过程中包含的前向和后向计算图的样子以及包含join和fan两种non-linear形式的计算图：

liveness analysis的核心思想是分析各中间tensor的存活时间段，即lifetime，比如tensor1在倒数第二层的前向layer生成，那么在正数第二层的后向layer消费，那么tensor1的lifetime为这两个layer之间的这个时间段；本文在每一层记录了该层运行前后的存活tensor，如下图，in表示该层运行前的存活tensor，out表示运行后的，拿step7举例，FC7在运行后，t2和t5这俩tensor不在被需要，就被free了
在这里插入图片描述

根据这样的原理，peak memory usage可以降低到，对应于上图中softmax6的位置，因为随着bwd的进行，比如进行到POOL8，那么FC4,softmax5，softmax6，FC7这四层layer的tensor都不被需要了，都被free了，所以memory usage随着backward的进行在下降，在一步步被free。

文中还提出了一种改进办法：因为liveness analysis的作用下，我们会经常的alloc内存给tensor，以及free tensor，如果此时还是用cudaMalloc和cudaFree API来做，在成千上万个迭代次数的训练过程中，那自然性能会很差，所以提出了memory pool来做alloc和free，简单来讲，mp的内存空间被划分为了以1kb block为单位的空间，mp数据结构里面包含了各个链表node，每个node包含三个数据：GPU address，node ID，occupied blocks。应对alloc请求的时候，遍历链表上的node，直到该node有足够大的block供使用，返回起始地址。

unified tensor pool
网络层数增加到1000时，imagenet training仍然会消耗100G显存，此时liveness analysis不够，作者提出了把某些layer占用的显存offload到CPU或者其他GPU，要用到的时候再prefetch会来。整个过程分为offload和prefetch阶段。核心原理是，通过CUDA async API，GPU上的计算时间掩盖GPU-CPU的memcpy时间，因此要求这些layer为计算密集型，不能为访存密集型，不然不能掩盖memcpy时间，会造成GPU的等待，所以文中只提出了offload convolution layer，prefetch与以上相反。

文中还提出了一种改进办法：因为conv在某些网络中比较有限，导致优化机会比较少，现代GPU和CPU之间的memcpy都是通过PCIe来传输数据，GPU和remote GPU之间通过DMA来传输，纵然传输速度可以达到几GB/s，每个iter都去传输GB数据量的数据会引入较大的overhead，所以需要优化一下，提出了一种办法，必要的时候才传，具体是引入了一个LRU cache在GPU DRAM上cache tensor来最大化利用GPU DRAM，减小数据传输开销。之所以选择LRU，主要是深度学习前向计算图是head to tail的顺序，而后向图是tail to head返过来的顺序，所以比较匹配LRU的特点，也是头部为最近使用，尾部为最远使用，主要分为三个函数：

1.in：将tensor插入到LRU的链表头

2.out：把一些tensor offload到CPU，空出显存

3.check：查询tensor是否在 LRU内，如果在，切到头部，不在则在头部创建一个node，如果此时显存空间不足，则调用out，把一些tensor offload到CPU，空出显存来存放该tensor，然后创建node
在这里插入图片描述
cost-aware recomputation

一共有三种策略：一种是面向速度的，一种是面向显存的，一种是两者之间的。第一种是只需要在前向过程中把所有中间tensor一次并存下来，以后要是还需要用到该tensor，则直接从显存里拿，reuse。第二种每次用到某个前向layer的tensor每次都去重算一遍。第三种则是一种新的策略，见c图，如果某几层的内存占用之和小于单层的最大内存占用，则用第一种策略，反之用第二种策略
在这里插入图片描述