機器之心報道

編輯：杜偉

想要了解自己的 PyTorch 項目在哪些地方分配 GPU 内存以及為什麼用完嗎？不妨試試這個可視化工具。

近日，PyTorch 核心開發者和 FAIR 研究者 Zachary DeVito 創建了一個新工具（添加實驗性 API），通過生成和可視化内存快照（memory snapshot）來可視化 GPU 内存的分配狀态。這些内存快照記錄了内存分配的堆棧跟蹤以及内存在緩存分配器狀态中的位置。

接下來，通過将這些内存快照可視化為火焰圖（flamegraphs），内存的使用位置也就能一目了然地看到了。

圖靈獎得主 Yann Lecun 也轉推了這個工具。

Twitter@Zachary DeVito

下面我們來看這個工具的實現原理（以第一人稱「我們」描述）。

生成快照

首先，我們必須記錄每次分配的堆棧幀信息。

記錄這些堆棧跟蹤的速度非常快（每次分配約 1 us，正常的 PyTorch 内核調用需要至少 8 us），但我們默認将其關閉。而啟用之後，我們可以分配一些内存并拍攝快照。

快照記錄了整個分配器的狀态，如下所示。

快照是具有以下結構的 Segment 字典列表。

Segments 是直接從 cudaMalloc 請求并由分配器緩存的内存。因此，我們可以隻使用這些 segments 中的一部分，緩存分配器将它們分為一個或更多個 Block。所有的塊始終保持相同的分配狀态。同時，使用_record_memory_history，每個塊還将記錄一個 History 對象，該對象會記住塊中最後一次分配的位置，包括作為 Frames 列表的堆棧跟蹤。

對于 active_allocated 塊，它其中存在的内容和當前分配的内容将有一個曆史記錄。對于 inactive 塊，可能會有多個條目來記錄塊内存中最後存在的内容。可能不止一個條目的原因在于分配器在空閑時會合并分割塊，并記錄下兩次拆分的曆史。為了避免出現大量的曆史記錄，我們隻保留不與任何更新塊重合的塊的曆史記錄。

保存快照

快照因自身設計而可以之後離線查看。

文件_memory_viz.py 可以直接用作交互式命令來處理保存的快照。

可視化快照

_memory_viz.py 工具也可以生成内存的可視化火焰圖。

可視化圖将分配器中所有的字節（byte）按不同的類來分割成段，如下圖所示（原文為可交互視圖）。

火焰圖可視化是一種将資源（如内存）使用劃分為不同類的方法，然後可以進一步細分為更細粒度的類别。

memory 視圖很好地展現了内存的使用方式。但對于具體地調試分配器問題，首先将内存分類為不同的 Segment 對象是有用的，而這些對象是分配軌迹的單個 cudaMalloc 段。

比較快照

該可視化器還可以生成顯示在兩個快照之間添加和删除的段的可視化。例如，我們可以使用更大的輸入重新運行模型，并查看分配器如何為更大的臨時對象請求更多内存。

比較視圖僅顯示新段，這有助于找出哪些代碼路徑提示分配更多内存。

$ python _memory_viz.py compare snapshot.pickle snapshot2.pickle -o compare.svg only_before = [] only_after = [140636932014080, 140636827156480, 140634912456704, 140634839056384, 140634843250688, 140634841153536, 140634866319360, 140634811793408, 140634845347840, $ 140636806184960, 140636778921984, 140634878902272]

原文鍊接：https://zdevito.github.io/2022/08/16/memory-snapshots.html

,