参考论文：ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
大模型的显存消耗一直都是面试常见的问题，这次我就彻彻底底的根据论文ZeRO中的调研和分析做一次分析
显存消耗的两个部分：Model States（跟模型的参数量和优化器相关）、 Residual Memory Consumption（跟训练时的batchsize，序列长度有关）
接下来，我就从这两个部分详细分析：
Model States
一个模型在显存消耗上，分为三个部分
- Optimizer States
- Gradients States
- Parameters States
更加具体的说，对于一个模型参数（Parameters）我们需要维护维护三个不同方面的参数
我们假设：模型的参数量大小为ModelSize
Parameters States
故名思义就是模型本身的权重参数，对于一个使用Float32存储的参数，我们需要32/8=4byte进行存储。
Gradients States
记录参数的梯度，对于一个使用Float32存储的参数，我们同样需要一个相同大小的梯度（4byte）保存它的梯度。
Optimizer States
对于最常用的Adam优化器以及其变体，对于一个使用Float32存储的参数需要维护两个额外的参数momentum和variance，也就是需要2*4=8byte进行保存
总的来说，对于Float32保存的模型来说，我们显存消耗是16（4+4+8）* ModelSize byte
但是对于半精度保存的模型（Float16），每个参数Parameters States 和 Gradients States的显存消耗都是2byte。在训练时，我们仍然需要保存其Float32的Parameters States用以加速运算，同时Adam优化器的两个参数momentum和variance同样也是Float32形式保存的，每个参数消耗的即为4+4+4=12 byte。所以半精度保存的模型，计算时的显存消耗仍然为16（2+2+12）* ModelSize byte
Residual Memory Consumption
剩下的显存消耗跟我们训练时的配置有关
主要有三个部分
- Activations
- Temporary buffers
- Memory Fragmentation
Activations
对于一个transformer based的模型来说，Activations的显存消耗和如下公式是成比例的：
number of transformer layers × hidden dimensions × sequence length × batch size
对于GPT2来说，这个比例大约为12
Temporary buffers 和 Memory Fragmentation
这两个参数不容易具体量化，Temporary buffers是多卡训练过程中为了提升梯度计算的效率，通常会执行一些类似于gradient all-reduce，gradient norm computation等操作，把数据集合到一个临时的缓存区中，这个临时区也会占用相当数量的显存
Memory Fragmentation，内存碎片的产生会导致内存空间的利用效率低下，即使有空余空间但是不足以分配给一个新的内存请求。
