过去两天，我一直在盯着几个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