给定一个模型架构、数据类型、输入形状和优化器,你能否计算出前向传播和反向传播所需的GPU内存量?要回答这个问题,我们需要将流程分解为基本组件,并从底层理解内存需求。以下实验(可以在Google Colab上运行)将帮助你理解核心概念。
预留与分配
PyTorch预留了更多内存,但只分配所需的内存。这样做是为了在需要更多内存时能够快速分配,而不是进行昂贵的预留操作。我们只关心内存分配,而不关心预留。
当删除变量x或当x超出作用域时,x的内存被释放,但仍然为将来使用而预留。只有在调用
时,才会释放预留的内存。
这里的
将返回PyTorch在此进程上分配的内存。如果有另一个进程正在使用一些GPU内存,将返回0。为了获取真实的GPU内存使用情况,可以使用以下函数。
数据类型
需要4字节的内存,
需要2字节,我们可以绘制一些数据类型所需的内存图。
图1:不同数据类型的内存分配
内存块
内存以512字节的块分配。当创建一个张量时,它被分配到下一个可用的块中。对于形状为(800,)的float32张量,不是分配800 * 4 = 3200字节,而是分配3584(512 * 7)字节。
图2:不同张量大小的内存分配。
可训练参数(单个线性层前向传播)
接下来我们将看一个单一的线性层。进行前向传播,并计算所需的内存。
有一个形状为(256, 250)的float32
矩阵,占用(256 * 250 * 4) = 256,000字节,这正好是内存块大小512的倍数(512 * 500 = 256,000)。但是
有250个float32需要占用(250 * 4) = 1000字节。而512的更高倍数是2,(512 * 2) = 1024字节。
和
是形状为(256,)的张量,所以它们各占用1024字节。总内存 =
当我们将所有内容加起来时,应该得到259,072字节(256,000 + 1024 + 1024 + 1024)。但是实际观察到的大小是8,778,752字节。这额外的8,519,680字节来自分配cuBLAS工作空间。
这是为快速矩阵乘法操作预留的内存空间。对于某些matmul操作,会分配一个新的8,519,680字节的块。这个大小可能会根据GPU和Python环境而变化。当调用
时,cublas内存不会消失。它需要
来实际清除它。也可以设置环境变量
来禁用cublas工作空间。但这可能是一种以牺牲执行速度为代价来优化内存的方法,所以我们使用默认就好。
梯度(单个线性层反向传播)
使用相同的模型,运行
。为简单起见假设损失为
。
由于每个具有
的模型参数都会有一个
成员来存储底层张量的梯度,所以模型的大小会翻倍。
这次分配了2个cublas工作空间内存块,假设一个用于前向传播,一个用于反向传播。此时cublas何时确切地分配新块还不确定。
中间张量(多层前馈网络)
当模型在推理模式下运行时,没有自动求导图,不需要存储中间张量。所以内存量只是简单地将每一层的内存相加。
在需要跟踪计算图的训练模式下情况会有所不同。当有多个串行应用的操作时,比如在前馈网络或任何深度网络中,自动求导图需要记住这些操作的中间张量。存储需求取决于它们的偏导数操作的性质。这些中间张量在反向传播过程中从内存中清除。我们看一些例子:
是输入,
是需要梯度的参数(
)。
-
x @ w
不需要额外的存储。偏导数x
已经存储。但是当x
是某个输出,如x = u * w1
时,x
也需要被存储。 -
x + w
也不需要存储,因为对w
的偏导数是0。 -
(x * 2) @ w
将需要存储操作数x * 2
,因为它将用于找到梯度。 -
(((x + 2) @ w1) + 3) * w2
是一个有趣的案例,模仿了2层。 - 对于关于w1
的偏导数,我们需要存储x + 2
- 对于关于w2
的偏导数,我们需要存储((x + 2) @ w1) + 3
让我们看看更深网络的实现:
在像BatchNorm1d、LayerNorm、RMSNorm这样的归一化层中,在与参数
相乘之前,有一个对输入
的操作,如
。操作数
是需要存储的中间输出。并且可能还有其他状态,如running_mean、running_std或
方法中的中间张量需要考虑。其中一些中间张量我们无法访问,所以我们无法确定发生了什么。当包含批量大小时,这变得更加复杂。
反向传播非常相似,但有一些变化:
- 模型大小因梯度存储而翻倍。
- 所有中间张量在最后都被清除。
- 分配了一个新的cublas工作空间。
优化器(单个线性层反向传播)
我们观察一些优化步骤的内存分配。
图3:使用SGD优化器在训练的各个阶段的内存分配
图4:使用Adam优化器在训练的各个阶段的内存分配
直到backward_1,我们看到内存分配如预期。当
结束时,在这个特定的代码中删除了
,所以该内存被释放。在底层优化器会获取额外的内存(等于可训练参数的大小)来更新它们,并在更新后释放该内存。这在图中没有显示。更详细的时间图可以在下图5中看到。
对于Adam对每个可训练参数都有一阶矩和二阶矩。所以它总是在内存中保留2倍的模型大小。这是这段代码中训练最耗费内存的部分。
图5:按毫秒计的内存分配时间图。
现在让我们尝试手动计算这些内存需求:
图6:使用SGD优化器在训练的不同阶段的内存使用分段
图7:使用Adam优化器在训练的不同阶段的内存使用分段
在手动计算内存分配后,我们的计算与观察结果相匹配。这次实际上可以看到内存分配到各种张量的分段。例如,Adam的状态占用了两倍的模型大小。梯度(红色)的不同变化。如果向继续测试,还可以尝试向这个模型添加更多层,添加中间张量并在适当的时候删除它们。这应该在这些条形图中创建另一个代表中间张量的分段。
总结
结合上面的每个概念我们可以回答主要问题:
- 可训练参数:固定的模型大小
- 内存块:它只以512字节的块出现
- Cublas内存:前向传播一个块,反向传播一个块
- 梯度:与模型大小相同
- 中间张量:最麻烦的部分,取决于代码如何编写
- 优化器:至少分配一倍的模型大小
最后一个问题就是,我们只处理了前馈层,那么CNN、Transformers、RNN等呢?首先CNN是类似前馈层的操作,所以我们可以根据他的计算规则进行计算,而Transformers、RNN都基础操作的组合,我们计算了一个前馈层可以根据他们的架构进行组合计算。我们已经掌握了计算前馈层内存需求的方法,所以我们可以自己解决这些问题!
作者:Akhilez