过去两天,我一直在盯着几个GPU内存使用情况图表看。关于如何解读这些图表的信息并不多,所以我写了一个脚本,只需一行代码就能完成分析,然后写了一篇博客文章。 那么,这些奇怪的形状到底代表什么呢? 快速教程🧵
首先,要分析任何代码块的性能,只需将此 Python 文件放入您的项目中,然后使用如下代码获取性能数据: https://t.co/iYTzEFksVR
你会看到类似下图的可视化效果。你可以立即看到三个训练步骤的三个重复模式。这是一个好的开始。这些条纹表示逐层发生的内存操作。下方的常量带代表模型参数和优化器状态。 但那些东西到底是什么呢?
我们来看第一个训练步骤,并进行分解。它以半三角形开始,代表前向传播。最后生成logits,也就是巨大的红色平行四边形。开头的尖峰是softmax函数的临时内存占用。logits之所以这么大,是因为词汇量有15.1万。
接下来我们调用 `zero_grad()` 函数,它会清除旧的 `.grad` 分配。在 `zero_grad()` 函数执行完毕后,会出现一个小尖峰,然后是一个巨大的断崖。尖峰是由于 `.backward()` 函数在 softmax 函数处开始执行,而断崖则是因为不再需要 logits 张量。这会产生 LM 头梯度(绿色带状区域)。
在上图中,请注意右侧的波浪状图案。这是因为我启用了激活检查点。这会导致每一层重新计算激活值,从而产生一些小的尖峰,同时也会随着 `.backward()` 逐层执行而释放其输入张量。
最后,我最喜欢的部分:optimizer.step()。它简直美得令人惊叹,就像一座矗立着旗杆的庙宇。它建立在两个基座之上。橙色条纹实际上是.backward()结束时生成的嵌入层梯度。
上图中,绿色条纹是存储 Adam 算法所需空间,用于存放梯度平方的平方根。融合 Adam 算法逐层计算更新,从而形成条纹。顶部区域则一次性应用所有更新。尖峰是应用更新前进行除法运算所需的临时缓冲区。
关于这些精美的可视化图表,还有很多可以说的。如果您发现任何错误或其他值得了解的有趣之处,请告诉我。 更多内容请访问我的博客: 如何获取和解读GPU内存分析数据 https://t.co/3uPt0S6RIp





