AI算力的阿喀琉斯之踵：內(nèi)存墻

作者 | Amir Gholami

翻譯與編輯 | OneFlow社區(qū)

這篇文章是我（Amir Gholami）， Zhewei Yao，Sehoon Kim，Michael W. Mahoney 和 Kurt Keutzer 等人共同協(xié)作完成的。本文中用到的數(shù)據(jù)可以參考鏈接https://github.com/amirgholami/ai_and_memory_wall 

https://github.com/amirgholami/ai_and_memory_wall/blob/main/imgs/pdfs/ai_and_compute.pdf

如圖表 1 所示，最近幾年，計(jì)算機(jī)視覺(CV)，自然語言處理(NLP)和語音識(shí)別領(lǐng)域最新模型的訓(xùn)練運(yùn)算量，以大約每兩年翻15倍數(shù)的速度在增長。而 Transformer 類的模型運(yùn)算量的增長則更為夸張，約為每兩年翻 750 倍。這種接近指數(shù)增長的趨勢(shì)驅(qū)動(dòng)了 AI 硬件的研發(fā)，這些 AI 硬件更專注于提高硬件的峰值算力，但是通常以簡化或者刪除其他部分（例如內(nèi)存的分層架構(gòu)）為代價(jià)。

然而，在應(yīng)付最新 AI 模型的訓(xùn)練時(shí)，這些設(shè)計(jì)上的趨勢(shì)已經(jīng)顯得捉襟見肘，特別是對(duì)于 NLP 和 推薦系統(tǒng)相關(guān)的模型：有通信帶寬瓶頸。事實(shí)上，芯片內(nèi)部、芯片間還有 AI 硬件之間的通信，都已成為不少 AI 應(yīng)用的瓶頸。特別是最近大火的 Transformer 類模型，模型大小平均每兩年翻240倍（如圖表2所示）。類似的，大規(guī)模的推薦系統(tǒng)模型，模型大小已經(jīng)達(dá)到了 O(10) TB 的級(jí)別了。與之相比，AI 硬件上的內(nèi)存大小僅僅是以每兩年翻2倍的速率在增長。

值得注意的是，訓(xùn)練 AI 模型時(shí)候所需要的內(nèi)存一般比模型參數(shù)量還要多幾倍。這是因?yàn)橛?xùn)練時(shí)候需要保存中間層的輸出激活值，通常需要增加3到4倍的內(nèi)存占用。圖表3中展示了最新的 AI 模型訓(xùn)練時(shí)候，內(nèi)存占用大小逐年的增長變化趨勢(shì)。從中能清楚地看到，神經(jīng)網(wǎng)絡(luò)模型的設(shè)計(jì)是如何受 AI 硬件內(nèi)存大小影響的。

這些挑戰(zhàn)也就是通常所說的 “內(nèi)存墻” 問題。內(nèi)存墻問題不僅與內(nèi)存容量大小相關(guān)，也包括內(nèi)存的傳輸帶寬。這涉及到多個(gè)級(jí)別的內(nèi)存數(shù)據(jù)傳輸。例如，在計(jì)算邏輯單元和片上內(nèi)存之間，或在計(jì)算邏輯單元和主存之間，或跨不同插槽上的不同處理器之間的數(shù)據(jù)傳輸。上述所有情況中，容量和數(shù)據(jù)傳輸?shù)乃俣榷即蟠舐浜笥谟布挠?jì)算能力。

大家可能會(huì)想到，是否可以嘗試采用分布式的策略將訓(xùn)練擴(kuò)展到多個(gè) AI 硬件（GPU）上，從而突破于單個(gè)硬件內(nèi)存容量和帶寬的限制。然而這么做也會(huì)遇到內(nèi)存墻的問題：AI 硬件之間會(huì)遇到通信瓶頸，甚至比片上數(shù)據(jù)搬運(yùn)更慢、效率更低。和單設(shè)備的內(nèi)存墻問題類似，擴(kuò)展 AI 硬件之間的網(wǎng)絡(luò)帶寬的技術(shù)難題同樣還未被攻破。如圖表4所示，其中展示了在過去20年中，硬件的峰值計(jì)算能力增加了90,000倍，但是內(nèi)存/硬件互連帶寬卻只是提高了30倍。而要增加內(nèi)存和硬件互連帶寬[1]，需要克服非常大的困難。因此，分布式策略的橫向擴(kuò)展僅在通信量和數(shù)據(jù)傳輸量很少的情況下，才適合解決計(jì)算密集型問題。

 https://github.com/amirgholami/ai_and_memory_wall/blob/main/imgs/pdfs/hw_scaling.pdf

有希望打破內(nèi)存墻的解決方案

“指數(shù)增長不可持續(xù)”，即使對(duì)于巨頭公司來說，以每2年240倍的速度持續(xù)指數(shù)增長也是不可能的。再加上算力和帶寬能力之間的差距越來越大，訓(xùn)練更大的模型的成本將以指數(shù)級(jí)增長，訓(xùn)練更大模型將更具有挑戰(zhàn)性。

為了繼續(xù)創(chuàng)新和 “打破內(nèi)存墻”，我們需要重新思考人工智能模型的設(shè)計(jì)。這里有幾個(gè)要點(diǎn)：

首先，當(dāng)前人工智能模型的設(shè)計(jì)方法大多是臨時(shí)的，或者僅依賴非常簡單的放大規(guī)則。例如，最近的大型 Transformer 模型大多是原始 BERT 模型[22]的縮放版本，二者基本架構(gòu)幾乎一樣。

其次，我們需要設(shè)計(jì)更有效的數(shù)據(jù)方法來訓(xùn)練 AI 模型。目前的網(wǎng)絡(luò)訓(xùn)練非常低效，需要大量的訓(xùn)練數(shù)據(jù)和數(shù)十萬次的迭代。有些人也指出，這種訓(xùn)練方式，不同于人類大腦的學(xué)習(xí)方式，人類學(xué)習(xí)某個(gè)概念或分類，往往只需要很少的學(xué)習(xí)例子。

第三，現(xiàn)有的優(yōu)化和訓(xùn)練方法需要大量的超參調(diào)整(如學(xué)習(xí)率、動(dòng)量等) ，在設(shè)置好參數(shù)從而訓(xùn)練成功前，往往需要數(shù)以百計(jì)次的試錯(cuò)。這樣看來，圖1中只是展示了訓(xùn)練成本的下限，實(shí)際成本通常要高得多。

第四，SOTA 類網(wǎng)絡(luò)規(guī)模巨大，使得光部署它們就極具挑戰(zhàn)。這不僅限于 GPT-3 等模型。事實(shí)上，部署大型推薦系統(tǒng)(類似于 Transformers ，但 embedding 更大且后接的 MLP 層更少)是巨頭公司所面臨的主要挑戰(zhàn)。最后，AI 硬件的設(shè)計(jì)主要集中在提高算力上，而較少關(guān)注改善內(nèi)存。這讓訓(xùn)練大模型、探索新模型都變得困難。例如圖神經(jīng)網(wǎng)絡(luò)（GNN）就常常受限于帶寬，不能有效地利用當(dāng)前硬件（的算力）。

以上幾點(diǎn)都是機(jī)器學(xué)習(xí)中的重要基礎(chǔ)問題。在這里，我們簡要討論最近針對(duì)后三點(diǎn)（包括我們自己）的研究。

高效的訓(xùn)練算法

訓(xùn)練模型時(shí)的一大困難是需要用暴力探索的方法調(diào)整超參。尋找學(xué)習(xí)率以及其配套的退火策略，模型收斂所需的迭代次數(shù)等等，這給訓(xùn)練 SOTA 模型帶來了不少額外開銷（overhead）。這些問題大多是由于訓(xùn)練中使用的是一階 SGD 優(yōu)化方法。雖然 SGD 超參容易實(shí)現(xiàn)，卻沒有穩(wěn)健的方法去調(diào)試超參，特別是對(duì)于那些還沒得到正確超參集合的新模型，調(diào)參就更加困難了。一個(gè)可能的解決方法是使用二階 SGD 優(yōu)化方法，如我們最近發(fā)表的 ADAHESSIAN 方法[4]。這類方法在超參調(diào)優(yōu)時(shí)往往更加穩(wěn)健，從而達(dá)到可以達(dá)到 SOTA。但是，這種方法也有亟待解決的問題：目前占用的內(nèi)存是原來的3-4倍。微軟關(guān)于 Zero 論文中介紹了一個(gè)很有前景的工作：可以通過刪除/切分冗余優(yōu)化器狀態(tài)參數(shù)[21, 3]，在保持內(nèi)存消耗量不變的前提下，訓(xùn)練8倍大的模型。如果這些高階方法的引入的 overhead 問題可以得到解決，那么可以顯著降低訓(xùn)練大型模型的總成本。

另一種很有前景的方法是提高優(yōu)化算法的數(shù)據(jù)本地性(data locality)并減少內(nèi)存占用，但是這會(huì)增加計(jì)算量。一個(gè)簡單的例子是，在前向 forward 期間，不保存所有的激活參數(shù)(activations)，而只保存它的子集，這樣可以減少圖3所示的用于特征映射內(nèi)存占用。未保存的激活參數(shù)可以在需要的時(shí)候進(jìn)行重計(jì)算，盡管這個(gè)方法會(huì)增加計(jì)算量，但只增加 20% 的計(jì)算量，可以減少高達(dá)5倍 [2]的內(nèi)存占用。

還有另一個(gè)重要的解決方案是設(shè)計(jì)足夠穩(wěn)健的、適用于低精度訓(xùn)練的優(yōu)化算法。事實(shí)上，AI 硬件的主要突破之一是支持了半精度(FP16)運(yùn)算，用以替代單精度運(yùn)算[5,6]。這使得算力提高了10倍以上。接下來的挑戰(zhàn)是，如何在保證準(zhǔn)確度不降低的前提下，進(jìn)一步將精度從半精度降低到 INT8。

高效部署

當(dāng)部署最新的 SOTA 模型，如 GPT-3 或大型推薦系統(tǒng)模型時(shí)，為了推理，常常需要做分布式部署，因此相當(dāng)具有挑戰(zhàn)性。可能的解決方案是，通過降低精度（如量化）或移除冗余參數(shù)（如剪枝）來壓縮推理模型。量化方法，既可以用于訓(xùn)練，也可以用于推理。雖然量化用于推理是可能做到超低精度級(jí)別的，但是用于訓(xùn)練時(shí)，想要將精度做到遠(yuǎn)低于 FP16 的級(jí)別是非常困難的。目前，在最小限度影響準(zhǔn)確率的前提下，已經(jīng)可以相對(duì)容易地將推理精度量化至 INT4 級(jí)別，這使得模型所占空間及延時(shí)，減少至原有的 1/8。然而，如何將精度量化至低于 INT4 級(jí)別，是一個(gè)頗具挑戰(zhàn)的問題，也是當(dāng)前研究的熱門領(lǐng)域。

除量化外，剪枝掉模型中冗余的參數(shù)也是高效部署的一種辦法。在最小限度影響準(zhǔn)確率的前提下，目前已經(jīng)可以使用基于 structured sparsity 的方法剪枝掉高達(dá) 30% 的神經(jīng)元，使用基于 non-structured sparsity 的方法可以剪枝掉高達(dá) 80% 的神經(jīng)元。然而，如果要進(jìn)一步提高剪枝的比率，則非常困難，常常會(huì)導(dǎo)致準(zhǔn)確度下降非常多，這該如何解決，還是一個(gè)開放問題。

AI 硬件設(shè)計(jì)的再思考

如何同時(shí)提高硬件帶寬和算力是一個(gè)極具挑戰(zhàn)的基本問題，不過，通過犧牲算力來謀求更好的“算力/帶寬”平衡點(diǎn)是可行的。事實(shí)上，CPU 架構(gòu)包含了充分優(yōu)化后的緩存架構(gòu)，因此在內(nèi)存帶寬受限類問題（如大型推薦系統(tǒng)）上，CPU 的性能表現(xiàn)要明顯優(yōu)于 GPU。然而，當(dāng)前 CPU 的主要問題是，它的計(jì)算能力（FLOPS）與 GPU 和 TPU 這類 AI 芯片相比，要弱一個(gè)數(shù)量級(jí)。個(gè)中原因之一，就是 AI 芯片為追求算力最大化，往往在設(shè)計(jì)時(shí)，就考慮移除了一些組件（如緩存層級(jí)）來增加更多的計(jì)算單元。我們有理由想象，可以有一種架構(gòu)，處于以上兩種極端架構(gòu)之間：它將使用更高效的緩存，更重要的是，使用更高容量的 DRAM（設(shè)計(jì) DRAM 層次結(jié)構(gòu)，不同層次擁有不同帶寬）。后者對(duì)于解決分布式內(nèi)存通信瓶頸將非常有幫助。

結(jié)論

目前 NLP 中的 SOTA Transformer 類模型的算力需求，以每兩年750倍的速率增長，模型參數(shù)數(shù)量則以每兩年240倍的速率增長。相比之下，硬件算力峰值的增長速率為每兩年3.1倍。DRAM 還有硬件互連帶寬增長速率則都為每兩年1.4倍，已經(jīng)逐漸被需求甩在身后。深入思考這些數(shù)字，過去20年內(nèi)硬件算力峰值增長了90000倍，但是DRAM/硬件互連帶寬只增長了30倍。在這個(gè)趨勢(shì)下，數(shù)據(jù)傳輸，特別是芯片內(nèi)或者芯片間的數(shù)據(jù)傳輸會(huì)迅速成為訓(xùn)練大規(guī)模 AI 模型的瓶頸。所以我們需要重新思考 AI 模型的訓(xùn)練，部署以及模型本身，還要思考，如何在這個(gè)越來越有挑戰(zhàn)性的內(nèi)存墻下去設(shè)計(jì)人工智能硬件。

感謝 Suresh Krishna 跟 Aniruddha Nrusimha 給出的非常有價(jià)值的回答。

[1] Patterson DA. Latency lags bandwidth. Communications of the ACM. 2004 Oct 1;47(10):71–5.

[2] Jain P, Jain A, Nrusimha A, Gholami A, Abbeel P, Keutzer K, Stoica I, Gonzalez JE. Checkmate: Breaking the memory wall with optimal tensor rematerialization. arXiv preprint arXiv:1910.02653. 2019 Oct 7.

[3] Rajbhandari S, Rasley J, Ruwase O, He Y. Zero: Memory optimizations toward training trillion parameter models. InSC20: International Conference for High Performance Computing, Networking, Storage and Analysis 2020 Nov 9 (pp. 1–16). IEEE.

[4] Yao Z, Gholami A, Shen S, Keutzer K, Mahoney MW. ADAHESSIAN: An adaptive second order optimizer for machine learning. arXiv preprint arXiv:2006.00719. 2020 Jun 1.

[5] Ginsburg B, Nikolaev S, Kiswani A, Wu H, Gholaminejad A, Kierat S, Houston M, Fit-Florea A, inventors; Nvidia Corp, assignee. Tensor processing using low precision format. United States patent application US 15/624,577. 2017 Dec 28.

[6] Micikevicius P, Narang S, Alben J, Diamos G, Elsen E, Garcia D, Ginsburg B, Houston M, Kuchaiev O, Venkatesh G, Wu H. Mixed precision training. arXiv preprint arXiv:1710.03740. 2017 Oct 10.

[7] Yao Z, Dong Z, Zheng Z, Gholami A, Yu J, Tan E, Wang L, Huang Q, Wang Y, Mahoney MW, Keutzer K. HAWQV3: Dyadic Neural Network Quantization. arXiv preprint arXiv:2011.10680. 2020 Nov 20.

[8] Gholami A, Kim S, Yao Z, Dong Z, Mahoney M, Keutzer K, A Survey of Quantization Methods for Efficient Neural Network Inference, arxiv preprint, arxiv:arXiv:2103.13630, 2021.

[9] Gale T, Elsen E, Hooker S. The state of sparsity in deep neural networks. arXiv preprint arXiv:1902.09574. 2019 Feb 25.

[10] Hoefler T, Alistarh D, Ben-Nun T, Dryden N, Peste A. Sparsity in Deep Learning: Pruning and growth for efficient inference and training in neural networks. arXiv preprint arXiv:2102.00554. 2021 Jan 31.

[11] Iandola FN, Han S, Moskewicz MW, Ashraf K, Dally WJ, Keutzer K. SqueezeNet: AlexNet-level accuracy with 50x fewer parameters and< 0.5 MB model size. arXiv preprint arXiv:1602.07360. 2016 Feb 24.

[12] Gholami A, Kwon K, Wu B, Tai Z, Yue X, Jin P, Zhao S, Keutzer K. Squeezenext: Hardware-aware neural network design. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops 2018 (pp. 1638–1647).

[13] Wu B, Iandola F, Jin PH, Keutzer K. Squeezedet: Unified, small, low power fully convolutional neural networks for real-time object detection for autonomous driving. InProceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops 2017 (pp. 129–137).

[14] Shaw A, Hunter D, Landola F, Sidhu S. SqueezeNAS: Fast neural architecture search for faster semantic segmentation. InProceedings of the IEEE/CVF International Conference on Computer Vision Workshops 2019 (pp. 0–0).

[15] Wu B, Wan A, Yue X, Keutzer K. Squeezeseg: Convolutional neural nets with recurrent crf for real-time road-object segmentation from 3d lidar point cloud. In2018 IEEE International Conference on Robotics and Automation (ICRA) 2018 May 21 (pp. 1887–1893). IEEE.

[16] Iandola FN, Shaw AE, Krishna R, Keutzer KW. SqueezeBERT: What can computer vision teach NLP about efficient neural networks?. arXiv preprint arXiv:2006.11316. 2020 Jun 19.

[17] Howard AG, Zhu M, Chen B, Kalenichenko D, Wang W, Weyand T, Andreetto M, Adam H. Mobilenets: Efficient convolutional neural networks for mobile vision applications. arXiv preprint arXiv:1704.04861. 2017 Apr 17.

[18] Krishna S, Krishna R. Accelerating Recommender Systems via Hardware” scale-in”. arXiv preprint arXiv:2009.05230. 2020 Sep 11.

[19] Kim S, Gholami A, Yao Z, Mahoney MW, Keutzer K. I-BERT: Integer-only BERT Quantization. arXiv preprint arXiv:2101.01321. 2021 Jan.

[20] Patrick Judd, Senior Deep Learning Architect, Integer Quantization for DNN Acceleration, Nvidia, GTC 2020.

[21] Bottou L, Curtis FE, Nocedal J. Optimization methods for large-scale machine learning. Siam Review. 2018;60(2):223–311.

[22] Devlin J, Chang MW, Lee K, Toutanova K. BERT: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805. 2018 Oct 11.

[23] Naumov M, Mudigere D, Shi HJ, Huang J, Sundaraman N, Park J, Wang X, Gupta U, Wu CJ, Azzolini AG, Dzhulgakov D. Deep learning recommendation model for personalization and recommendation systems. arXiv preprint arXiv:1906.00091. 2019 May 31.

[24] LeCun Y, Bottou L, Bengio Y, Haffner P. Gradient-based learning applied to document recognition. Proceedings of the IEEE. 1998 Nov;86(11):2278–324.

本文由oneflow社區(qū)團(tuán)隊(duì)翻譯編輯，并已獲得作者授權(quán)。

原文鏈接：https://medium.com/riselab/ai-and-memory-wall-2cb4265cb0b8

作者簡介：Amir Gholami is a senior research fellow atICSI and Berkeley AI Research (BAIR). He received his PhD from UT Austin,working on large scale 3D image segmentation, a research topic which receivedUT Austin’s best doctoral dissertation award in 2018. He is a Melosh Medalfinalist, the recipient of best student paper award in SC'17, Gold Medal in theACM Student Research Competition, as well as best student paper finalist inSC’14. He was also part of the Nvidia team that for the first time made lowprecision neural network training possible (FP16), enabling more than 10xincrease in compute power through tensor cores. That technology has been widelyadopted in GPUs today. Amir's current research focuses on efficient AI at theEdge, and scalable training of Neural Network models.