前情提要#
當我確定要去投一個會議的稿件的時候,距離截稿日只有一週的時間了。但是原先取得的數據沒有辦法直接用,只能從頭開始設計程序獲取數據,值得慶幸的是好在核心的計算模塊可以復用。我在剛開始的第一天還是很樂觀的,預估用兩天時間獲取實驗數據,三天寫文章,剩下的時間裡修改,趕在截稿前就能提交。
但是實際情況並不盡人意。由於過往的經驗,我錯誤的估計了實驗的運行時間。我原先的實驗是在每輪計算 4 個大規模複數矩陣,但是這次實驗為了得到更詳細的結果,每輪需要計算的矩陣暴漲到 42 個。更糟糕的是,我原先預計跑 28 輪計算,而每一輪的每一個矩陣的規模都會是上一輪的兩倍。第 28 輪時每個矩陣的數據量會達到大約 1GiB,而為了對這個矩陣做處理還需要額外存儲三個同等規模的矩陣(為了方便說明,讓我們把之前的那個矩陣稱為核心矩陣,這三個矩陣稱為輔助矩陣)!也就是說光是處理一個核心矩陣就需要 4GiB 的空間,更別提中間產生的一些中間量數據,而這樣的計算我需要做 42 次(這還只是第 28 輪單輪計算過程!)......
我沒有辦法把這樣的計算量壓在一顆 CPU 核心上,除非我想算到天荒地老...... 我最早的方案是用多核並行計算。因為每一個核心矩陣的計算是彼此獨立的,非常適合並行計算的方案,這在我最早的實驗裡是對的,因為我每輪涉及到的只有 4 個核心矩陣,哪怕加上中間量,也不會超過我有的 64GB 內存,我不需要考慮其他的方案,只要在 4 個 CPU 核心同時啟動,只需要大約 2 個小時就能給我結果。
但當每輪的核心矩陣數達到 42 個,這樣的方案就行不通了!首先為了縮短計算時間,我需要更大的並行數,而當我把並行數提高到 16 的時候,第 28 輪的總占用內存量達到了 100GB 以上!我不得已為此額外添加了內存,但是事實上,哪怕我添加了足夠的內存,這樣的計算量和數據交換量對於一顆普通的家用 CPU 而言還是太大了,第 20 輪以後的每輪的計算時間都是用小時作為單位,而且還會呈指數級增長...... 更不用說,如果在途中出現了意外退出或者輸出結果不如人意,我免不了回頭重新審視程序設計,更有可能需要從頭再跑幾輪......
我真的很希望時間能倒退幾個月!這樣我可以好好設計我的實驗方案,再不濟我也有足夠的時間去等待整個計算過程。但是現在哪怕我能兩天寫完文章,我也只有四到五天時間了!我沒有辦法去做一些很複雜的優化設計,因為我沒有時間去反復調試程序,現在擺在我面前的只有一條路:去找一種簡單粗暴的方案,能夠復用大部分程序,不需要對整體框架做出大範圍修改,同時要能夠想方設法把指數級增長的計算時間壓下去。
讓我們看一下這個問題的核心思路:每一輪核心矩陣的規模都是翻倍的,為了能夠讓計算時間保持線性增長,我們需要讓每個矩陣的計算時間也保持線性,所以最簡單的方法,如果矩陣規模翻倍,那就讓計算模塊也翻倍。
很遺憾我的 CPU 沒有辦法做到這一點,他只有 8 個核心!但是我們也不是沒有其他的手段,想一想那個著名的視頻,在流言終結者裡有一集,他們用彩彈槍形象地再現了 CPU 和 GPU 的主要區別。如果我們沒有辦法讓我們的 CPU 的核心滿足我們的指數增長的需求,那麼把矩陣相關的計算交給並行計算能力更強的 GPU 就是一個很明智的選擇,GPU 簡直就是為此而生的!
異構計算方案#
我們把結合不同計算單元進行計算的方法稱作異構計算。最典型的例子就是 CPU 和 GPU 的結合。是的,雖然這個詞聽上去非常的高級,但是它廣泛應用在許多領域,其中最有親切感的可能就是遊戲。CPU 和 GPU 的結合是個熱門領域,正因為這一方案的廣泛性,我們有很多簡單好用的工具來支持我們將單一 CPU 架構的程序轉換為結合 GPU 的程序,這些工具裡最有代表性的是 NVIDIA 公司針對自己 GPU 產品推出的 CUDA,還有對多個廠商的 GPU 都提供了支持的 OpenCL。這些年隨著對 GPU 計算的重視,AMD 也推出了他們家的 ROCm 方案,但是這個工具的社區和生態還待提高。我儘管很願意去嘗試新的技術,但是現階段我得先完成我的實驗!
我家裡的電腦和研究室的電腦使用的都是 NVIDIA 的 GPU,所以我最早考慮的就是使用 CUDA 的方案。CUDA 是一個非常成熟的工具,也有人針對它開發了許多好用的庫。尤其是近些年隨著 GPU 在 AI 領域的使用,湧現了很多傻瓜式調用 CUDA 的庫,甚至有不需要對原有代碼進行修改,就能直接將計算模塊下放到 GPU 的手段,遺憾的是我目前的程序並不適用。我最後用的是一個叫做 CuPy 的 Python 庫,因為他和 NumPy 有很高的相似性,你也可以認為他是 NumPy 的 GPU 版本的簡單實現。這一點對我來說很重要,因為涉及到矩陣計算,為了提高計算速度,我在程序裡大量的使用了 NumPy 的計算方法和數據類型,我並不想在將計算模塊下放(我指 offload 這個詞)給 GPU 時還需要對數據進行過多的整形,這不僅僅只是因為我需要大量修改我的代碼,而且這也會產生許多額外的計算開銷。
CuPy 的用法很簡單,我用了大概一小時左右就成功修改並跑通了程序,在這裡我無意過多地講述代碼實現的細節,如果有時間的話我可能會單獨寫一篇文章來聊一聊這些(當然最好是在我對 CUDA 有更深的理解以後)。但是在這個時間點我卻並沒有得到我想要的結果,現實總是比人們想象到的更為複雜。
首當其衝的依然是數據傳輸問題,儘管現代 GPU 採用的都是高速存儲,但是如果反復在 CPU 和 GPU 之間交換數據,那麼大量的時間將會被消耗在這上面,計算部分必須一直等待數據就位,而集齊了數據以後又算不了多久,馬上又要把結果傳輸出去了。這並不是使用 pipeline 結構或者將數據分塊傳輸就能解決的問題。因為在最初的設計裡,GPU 只需要負責對每個矩陣要素單獨進行的幾個簡單計算,你可以簡單理解為我們只是將每個要素都單獨乘以 2,所以 pipeline 一類的計算結構並不十分適用,同時,從一開始我們的計算就不需要等待整個矩陣的傳輸,計算過程也不會造成巨大的延遲,一切的問題根源就在於這裡,我們的計算部分足夠快,所以性能瓶頸就落在了數據傳輸上。
在不更改原有設計的基礎上這一點很難改變,我不得已對計算結構進行了調整,將後續很多放在 CPU 上的數據處理部分也下放給了 GPU。這個修改是很有效的,我們加大了 GPU 的計算負載,同時盡可能減少 CPU-GPU 之間頻繁的數據傳輸,這使得我的 GPU 的最高計算負載到了 80%。
然而在這裡出現了另外一個問題,現在我的 GPU 不再僅僅對每個矩陣元素進行處理,而是需要保存處理後的矩陣,再對其整體進行後續的處理,就像最開始提到的一樣,我們會需要非常巨大的矩陣,這種時候我們的 GPU 內存就開始不足了,很遺憾我的 GPU 的內存只有 4GB,儘管通過 SWAP 我們可以共享部分的系統內存,但是這部分就會遭遇頻繁數據傳輸的問題... 我當時能想到的手段,是通過分割矩陣,只給 GPU 部分矩陣來獲取部分的結果,再將其傳到 CPU 後進行合併,這比原先會好一些,但在 24 輪之後依然無可避免的遭遇了數據傳輸帶來的瓶頸,硬件帶來的限制比想象的更嚴重一些。
由於時間上的限制,我迫不得已放棄了原先運行 28 輪的方案,而是在盡可能保證實驗結果的可驗證性和完整性的基礎上,降低了計算次數,選擇了運行 25 輪的方案,到這一步我成功在 8 小時內得到了所有我需要的結果,並趕在截止前完成了我的文章(感謝我的老師!)
我依然在繼續進行程序的優化,接下來的方向有兩個,一個是通過 MPI 等方案分散計算壓力到多個機器,但預想這依然有數據傳輸瓶頸等著我解決,另一個是使用 FPGA 做數據計算後傳遞給主機,資源夠多的情況下這種方案比 GPU 在數據傳輸和延遲方面更有優勢,但是開發成本也有很大的提升,我還在猶豫這種方案,除非某一天我發現我確實需要定制化硬件來解決一些傳統方案所無法解決的問題。