CUDA 有 unified memory 還需要記憶體優化嗎？

CUDA 有 unified memory 還需要記憶體優化嗎？

先說結論，不管有沒有 unified memory，只要有效能要求的軟體，都必需自己做優化。

因為 unified memory 的主要工作，並不是去掉了 Host Memory 與 Device Memory 之間的拷貝，而是將原有的需要人工手寫的這部分程式碼，交給了執行時，使得這些工作對程式設計師透明，但是 Host2Device、Device2Host 的記憶體拷貝所造成的 overhead 依然存在。

其次，“記憶體優化”其實是一個非常寬泛的話題，他通常可以指記憶體生命週期的自動管理（如各種解釋型語言使用的垃圾回收機制）；也包括減少動態記憶體重複申請、釋放頻率提高效率；還包括用更少的記憶體佔用量完成一樣的功能（當年阿波羅登月飛船上的計算機記憶體只有4KB）等。

以深度學習框架 OneFlow 為例，研發團隊在不依賴 unified memory 的前提下，就在框架中自己實現了對開發者透明的 Host/Device 記憶體管理。

在 OneFlow 中，開發者只需要使用簡單的表示式：user_op::HobDeviceTag() == DeviceType::kGPU 或者 user_op::HobDeviceTag() == DeviceType::kCPU 就可以告之框架相關功能是使用 GPU 記憶體還是 CPU 記憶體，程式設計過程中拿到的記憶體地址就是對應記憶體空間的地址。那些煩人的記憶體分配、裝置之間的記憶體拷貝、動態記憶體的生命週期管理統統不用操心。

除了程式設計易用性外，怎樣使用更少的記憶體完成更多的工作量也是記憶體優化重點。很多人會認為這塊應該交給硬體廠商做，同時也常認為：如果硬體廠商做得足夠好，就沒有軟體開發者什麼事情了，遺憾的是，這是一個無法達到的理想境界。因為極致的記憶體優化與業務邏輯強相關，而硬體廠商要兼顧通用。

舉個業務邏輯相關的例子，比如做深度學習框架，假設我們有 N 個引數，那麼在反向過程中，需要為這 N 個引數分配空間存放梯度。一個新手在一次訓練迭代中，可能會這樣管理記憶體生命週期：

• 獲取引數資訊 N
• 根據引數資訊 N ，分配空間 R
• 反向傳播，更新梯度、更新引數
• 釋放空間 R

當第二次訓練迭代過程中，會重複以上過程。然而，只要稍加分析，就可以發現，在一次迭代結束後，不必要立即釋放空間 R，而可以留給下一次迭代複用。這是一個很簡單的例子，但是硬體廠商無法為此”擅作主張“，需要結合業務邏輯分析。

在實際的實現中，業務邏輯對程式碼複用的影響就更加複雜，而優化結果的回報自然更加豐厚。

可以有多豐厚呢？以深度學習 OneFlow 為例，在一流的演算法工程師和框架工程師的配合下，同樣的任務，OneFlow 和英偉達的親兒子 HugeCTR 比較，OneFlow 記憶體節省驚人。

對於不熟悉 HugeCTR 的同學，簡單介紹下，它是英偉達為解決大模型訓練問題開發的框架，在 OneFlow 釋出之前，HugeCTR 效能一直是世界第一。

以下是跑 Wide&Deep 不同 batch size 的的對比結果，可以看到，從 batch szie 一萬到百萬數量級，OneFlow 的記憶體佔用率都只有一半不到，最少時只佔 24%，四捨五入就是，用 OneFlow 買一張顯示卡可以完成的工作，換其他人得買4張顯示卡。當 batch size 達到 2097152 之後，就無法比較了，因為 OneFlow 還可以頂一下，HugeCTR 已經記憶體溢位（OOM Out of Memory）。

綜上所述， unified memory 簡化了程式設計的過程，但是我們不能期待它可以幫助我們一勞永逸地解決記憶體優化的問題。反過來說，恰恰是因為特定的領域（包括但不限於深度學習框架開發）還存在著許多頗具挑戰、回報豐厚、懸而未決的問題，才讓這些領域充滿魅力。

最後，安利一下 OneFlow 的測試倉庫 DLPerf。