MegaSlide-DiT:以記憶體為核心的適配與可變形區域性注意力實現高效影片擴散
MegaSlide-DiT提出一種系統級方案,透過將模型狀態駐留於主機記憶體並流式傳輸至GPU,結合3D可變形滑動注意力機制,在單塊H200 GPU上適配105B引數的DiT模型,解決了引數記憶體和啟用記憶體兩大瓶頸。
基於擴散變換器(DiT)的高解析度影片擴散模型能夠生成極其逼真的影像和影片,但其巨大的引數量和對記憶體的需求遠遠超出了單個工作站的承載能力。例如,一個超過1000億引數的DiT模型需要超過1TB的持久化狀態來儲存權重、最佳化器狀態等,而傳統時空自注意力機制的計算複雜度隨序列長度呈平方增長,導致啟用記憶體也迅速膨脹。這“兩堵牆”——引數記憶體牆和啟用記憶體牆——使得研究人員在沒有大規模GPU叢集的情況下,幾乎無法適配這些龐大的生成模型。
來自研究者Jiacheng Liu和Jason Liu的最新論文MegaSlide-DiT從一個新穎的系統級視角重新審視了這一問題。他們提出了一套原型方案,證明可以在配備1.5 TB主機記憶體的單塊NVIDIA H200 GPU上,成功適配一個預訓練的105B引數DiT模型。其核心洞察在於:GPU無需擁有整個模型狀態。具體來說,所有持久化的權重、主權重和最佳化器狀態都保留在主機記憶體中,只有臨時計算所需的分片被按需流式傳輸到GPU。這種方法極大地降低了對GPU視訊記憶體的要求。
與此同時,研究者用創新的3D可變形滑動注意力(3D-DSA)取代了傳統的二次複雜度全域性注意力。3D-DSA是一種運動自適應的區域性注意力運算元,它透過關注影片幀中隨時間變化的區域性區域,將注意力的記憶體和計算複雜度都降低到序列長度的線性級別。這一設計不僅大幅節省了視訊記憶體,還顯著加快了計算速度。
論文提供了詳細的記憶體佔用統計、執行軌跡和評估結果來支撐其設計。需要注意的是,MegaSlide-DiT並非聲稱能在單GPU上從頭訓練105B模型,也沒有神奇地解決頻寬限制問題——它本質上提供了一條在高階工作站上對大規模影片擴散模型進行全引數適配的實用路徑。這項研究為降低大規模AI模型部署的門檻提供了重要參考,尤其對資金有限的研究團隊和中小企業具有重大意義。