AI News HubLIVE
站內改寫1 分鐘閱讀

MegaSlide-DiT:以內存為核心的適配與可變形局部注意力實現高效視頻擴散

MegaSlide-DiT提出一種系統級方案,通過將模型狀態駐留於主機內存並流式傳輸至GPU,結合3D可變形滑動注意力機制,在單塊H200 GPU上適配105B參數的DiT模型,解決了參數內存和激活內存兩大瓶頸。

來源arXiv Computer Vision作者: Jiacheng Liu, Jason Liu

基於擴散變換器(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模型部署的門檻提供了重要參考,尤其對資金有限的研究團隊和中小企業具有重大意義。