近日，PyTorch核心开发者兼FAIR研究员ZacharyDeVito创建了一个新工具（添加实验性API），通过生成和可视化内存快照（memorysnapshot）来可视化分配GPU内存的状态。这些内存快照记录了内存分配的堆栈跟踪以及内存处于缓存分配器状态的位置。接下来，通过将这些内存快照可视化为火焰图，您可以一目了然地看到哪些地方正在使用内存。图灵奖获得者YannLeCun也转发了这个工具。Twitter@ZacharyDeVito让我们来看看这个工具是如何工作的（以第一人称“我们”描述）。生成快照首先，我们必须记录每次分配的栈帧信息。记录这些堆栈跟踪非常快（每个分配大约1us，正常的PyTorch内核调用至少需要8us），但我们默认关闭它。启用后，我们可以分配一些内存并拍摄快照。快照记录了整个分配器的状态，如下所示。快照是具有以下结构的段字典列表。段是直接从cudaMalloc请求并由分配器缓存的内存。因此，我们只能使用这些段中的一部分，缓存分配器将它们分成一个或多个Block。所有块始终保持相同的分配状态。此外，使用_record_memory_history，每个块还将记录一个历史对象，该对象记住最后一次分配在块中的位置，包括作为帧列表的堆栈跟踪。对于一个active_allocated块，将有一个历史记录，说明其中存在什么以及当前分配了什么。对于非活动块，可能有多个条目来记录块内存中最后的内容。可能有多个条目的原因是分配器在空闲时合并拆分并保留两个拆分的历史记录。为了避免大量的历史记录，我们只保留与任何更新的区块不一致的区块的历史记录。保存快照按照设计，快照可以在以后离线查看。文件_memory_viz.py可以直接用作交互式命令来处理保存的快照。可视化snapshot_memory_viz.py工具也可以生成内存的可视化火焰图。可视化图将分配器中的所有字节按照不同的类划分成段，如下图所示（原文为交互视图）。火焰图可视化是一种将资源（如内存）使用划分为不同类别的方法，然后可以进一步细分为更细粒度的类别。内存视图可以很好地显示内存的使用情况。但是为了专门调试分配器问题，首先将内存分类到不同的Segment对象中是有用的，这些对象是分配跟踪的各个cudaMalloc段。比较快照可视化工具还可以生成可视化效果，显示在两个快照之间添加和删除的段。例如，我们可以使用更大的输入重新运行模型，并查看分配器如何为更大的临时对象请求更多内存。比较视图仅显示新段，这有助于找出哪些代码路径暗示分配更多内存。$python_memory_viz.pycomparesnapshot.picklesnapshot2.pickle-ocompare.svgonly_before=[]only_after=[140636932014080,140636827156480,140634912456704,140634839056384,140634843250688,140634841153536,140634866319360,140634811793408,140634845347840,$140636806184960,140636778921984,140634878902272]
