過去兩天，我一直在盯著幾個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