記憶體優化系列文章(1)AI效能加速1.3x倍就靠一個環境變數!

愛上一隻檸檬的pig_head發表於2020-10-22

記憶體優化系列文章(1)AI效能加速1.3x倍就靠一個環境變數!

虛擬記憶體、實體記憶體、記憶體分配、記憶體洩漏、TLB 不命中、缺頁異常這些看似熟悉的計算名詞卻似懂非懂,沒關係!本系列文章通過理論和實戰帶你一網打盡!
自從工作中接觸到深度學習記憶體優化知識以來,白丁決定重新系統研究一下計算機背後的記憶體故事。通讀本文你將獲得以下乾貨知識:
為什麼記憶體優化可以大幅度提升深度學習演算法的效能?如何通過工具定量分析記憶體優化帶來的效能好處?

作業系統為什麼採用虛擬記憶體管理而不直接使用實體記憶體?虛擬記憶體的核心思想是什麼?虛擬記憶體管理對應用程式的效能有什麼影響?應用程式如何從作業系統申請和釋放記憶體?

常見記憶體分配演算法 PTMalloc、TCMalloc 和 Jemalloc 詳解

什麼是零拷貝?

程式中的記憶體管理和共享記憶體?

如何利用虛擬記憶體進行 I/O 優化?

BERT 模型記憶體優化

一個環境變數效能直接飆升1.3x, 記憶體優化這麼牛?## 標題
在實際的優化工作中,我們經常利用 TCMalloc 或者 JEMalloc 來替換 glibc 中預設使用的 PTMalloc 來獲得更好的模型效能。 以下圖中BERT 模型在 Pytorch 中的推理為例, 在採用了TCMalloc之後,BERT 的吞吐量直接從60.8 sentence/s 飆升到79.2 sentence/s,效能直接提升了1.3x。

所謂知其然並知其所以然。替換 malloc 演算法其實就是一個export 一個 LD_PRELOAD 環境變數的事(下一篇就給你教程),為什麼會發生這麼大的效能影響?這會影響作業系統裡的哪些指標?如何通過現有的工具分析?
###缺頁異常大幅度降低
簡單來說,現代作業系統採用虛擬記憶體(後續詳解)的管理機制,系統在訪問某一記憶體頁時需要首先利用頁表機制將虛擬地址對映到實體地址,如果在查詢的過程中發現頁表中沒有對應的對映條目,就會發成缺頁異常。因為要執行完異常處理程式後,程式才能繼續執行,這難免會拖累效能。在 Linux 系統中 ps 命令可以用來檢視程式的缺頁狀況。在本實驗中,我們統計在實際推理過程中(不包括模型初始化的過程)的 minflt 的情況如下圖所示。相對而言,PTMalloc 的缺頁次數明顯多於兩種 malloc 演算法。這也在一定程度上解釋了為什麼推理的效能會有如此大的差距。對於記憶體機制而言,除了缺頁可能導致程式變慢之外?還有什麼可能影響程式效能呢?

ps -o majflt, minflt  -C program

TLB 不命中率降低

上述分析中提到虛擬記憶體管理機制需要通過查詢頁表才能獲取實際的實體地址。頁表通常儲存在實體記憶體中,類似於快取機制,常用的頁表項會被載入到 TLB (Translation Look-aside Buffer)中以加快頁表的訪問。記憶體對映單元(MMU)會優先訪問 TLB 獲取實體地址對映,但是如果 TLB 不命中將導致 MMU 需要從記憶體中載入頁表項,因此 TLB 不命中勢必會影響程式的執行效率。perf 是一個強大的效能分析工具,它可以觀測到諸如快取不命中、缺頁異常、TLB 不命中等系統級的效能事件。下圖的分析結果表明,相對於其他兩種 malloc 演算法,PTMalloc 也具有明顯的 TLB 不命中比例,這也在一定程度上拖累了程式效能。試想,缺頁異常表示的是頁表項中沒有該虛擬頁到實體記憶體的訪問。那也意味著這個虛擬地址很可能是第一次訪問,在訪問該虛擬地址之前沒有建立對映關係,因此,缺頁異常必然導致 TLB 不命中。

perf stat -e dTLB-loads,dTLB-load-misses,iTLB-loads,iTLB-load-misses your_program 

上文我們介紹了 malloc 演算法對程式效能的影響以及背後的原因分析。下一篇文章我們將提供 TCMalloc 和 JeMalloc 快速入門指南以及使用時候的注意事項。更多記憶體優化的知識請持續關注 "程式設計師的大廠之路" 微信公眾號

作者簡介:小編是某知名企業的一名深度學習軟體工程師,主要從事深度學習效能優化方面的工作包括但不限於深度學習框架優化、模型量化、分散式訓練等。歡迎留言一起討論計算機基礎,深度學習效能優化等方面的知識。

相關文章