大模型LLM训练显存消耗详解

最新推荐文章于 2024-06-29 09:44:56 发布

ZYM66

最新推荐文章于 2024-06-29 09:44:56 发布

阅读量1.4k

点赞数 14

文章标签：人工智能 llama 语言模型自然语言处理 pytorch

本文链接：https://blog.csdn.net/ZYM66/article/details/136158868

版权

文章深入解析了大模型训练中的显存消耗，重点关注了ModelStates（包括OptimizerStates、GradientsStates和ParametersStates）和ResidualMemoryConsumption（涉及Activations、Temporarybuffers和MemoryFragmentation）。通过比较Float32和Half精度对内存的影响，展示了如何通过优化来减少显存占用。

摘要由CSDN通过智能技术生成

参考论文：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，内存碎片的产生会导致内存空间的利用效率低下，即使有空余空间但是不足以分配给一个新的内存请求。