大模型优化1：简介

随着深度学习模型越来越大，训练模型需要的GPU越来越多，对于Google、Fackbook来说，不差钱，但是对于很多资金紧张的团队来说，训练大模型的花费变得负担不起。

为了解决这个问题，”模型内存优化“技术开始展露头角，该技术致力于减少模型训练时占据的显存，并极力缩小由此带来的时间开销。

其实在推理领域，也有很多减少模型推理时内存的开源工具，毕竟推理场景中，有更多将模型放入嵌入式系统的需求。本篇只对训练做介绍。

 

现有的模型优化技术，大概分为以下几种类别：

• Memory swaping
• Memory sharing
• Recompute
• Model compression
• Unified memory

1、Memory swaping:

memory swapping技术，主要思想是把暂时用不到的tensor交换到cpu中，用到时再传回gpu中。
对于如何挑选需要swap的Tensor，不同的论文中有不同的挑选原则：
有的论文是小的tensor不swap，因为它占据的内存小，而swap反而会增加数据传输延迟；并且只swap在peak time中会存在的tensor，对于variable的生命周期短的也不swap，因为会增加I/O的带宽。并且它会对swap in的调配过程设计算法，使其能够充分利用IO带宽，即swap out时会按其触发顺序进行，swap in 时会用pefetch策略，并在上一个swap out/in结束后开始。
SuperNeurons只交换卷积层的tensor，而其他层的big tensor不交换，并且为了减少交换带来的通信开销，只有在GPU不够用时才会使用LRU方式交换
TFLMS只要边的长度大于一定的阈值，就会把tensor交换出去，这种频繁的交换会造成很大的communication overhead。(TFLMS)
多数论文中都可以自己设置swap哪些tensor，这就要求用户对模型的理解非常清楚，完美的交换就是把communication overhead 隐藏在计算下。

2、Memory sharing

Memory sharing有buffer reuse和in-place operation两种方式。
In-place opereation就是输出和输入存在同一个物理地址内，输入是可以被输出覆盖的，后面不会用到这个tensor，但这个只有很少的op能采用这种方式，因为大多数op输出的tensor都需要被后面的op用到。
Buffer reuse就是在lifetime上不重叠的tensor可共享同一片内存区域。这类方法的主要改进点就是将memory pool的动态寻找一个合适的block的过程变为静态的，是通过迭代训练一次，获得每个memory block request的信息，然后内存分配就抽象成了如何把一堆小矩形放入一个固定宽度的大容器中，并使高度最小的问题。这个小矩形可以代表一次request的block，也可以代表一个tensor。

如图，通过训练迭代一次获得memory block request信息，然后根据信息将训练时内存分配抽象为两维的矩阵堆放问题，从而后面的训练可采用静态方式分配内存，以此减少寻找block的时间

 

图1：buffer reuse

 

 

3、Recompute

Recompute是把在前向计算中将tensor释放掉，在反向传播时要用时，在重新计算该tensor，占据内存大，并且重计算量小的tensor适合用这种方式。
重计算有很多种方式，speed centric在重新计算出一个tensor后，会保留这个tensor，以备后续使用，memory centric在计算出一个tensor后还会把它释放掉，再次使用时再重新计算。Cost aware再重新计算出一个tensor后有两种选择，当继续保存这个tensor可能导致大于内存峰值时，就释放掉，否则就保留。
可以将swap和recompute结合在一起使用，对于conv的tensor，采用swap的方式，而对于BN的tesnor，采用recompute的方式，这种就是对于特定的op采用不同的方式。此外，还有一种结合方式是预先迭代几次，收集必要的内存和运行时间信息，然后判断哪些tesnor该swap，哪些该recompute

 

4、Unified memory

Unified memory是把cpu和gpu内存放在一起统一管理，在oc-caffe里，用cuMemAllocManaged 代替cuMemAlloc,在内存分配时实现既可以在cpu内存上分配也可以在显存上分配。cudaMemCpy用于GPU内部的数据的拷贝，而在managed buffer上的数据拷贝操作会比在gpu之间的数据拷贝操作性能低，因此oc-caffe用prefetch()预取destination buffer，用Advise(evict())标记source buffer，当需要时会把source buffer传到cpu内存上。

 

5、Model compression

模型压缩是将大模型采用裁剪、权重共享等方式进行处理，以减少模型参数量。通常这种方式，比较少用，因为非常容易使精度降低。

模型压缩方法通常有

• 修剪：可采用对连接、kernel、channel进行裁剪的方式
• 权重共享
• 低秩分解
• 二值化权重，权重由32bits换成8 bits、16bits，混合精度训练
• 知识蒸馏：使用训练好的teacher模型，来指导student模型训练

 

 

 

参考：

memory swapping：SuperNeurons      TFLMS 

 

memory sharing：

Profile-guided memory optimization for deep neural networks

 Efficient Memory Management for GPU-based Deep Learning Systems

 

 Unified Memory:

OC-DNN: Exploiting Advanced Unified Memory Capabilities in CUDA 9 and Volta GPUs for Out-of-Core DNN Training

 

model compression:

Distilling the Knowledge in a Neural Network