當前位置: 華文世界 > 數碼

開源AirLLM,單卡4GB視訊記憶體跑70B大模型

2024-04-28數碼

目標:在單個4GB顯卡上執行LLama 3 70B模型 - 使用airllm和分層推理技術

分層推理基本上是「分而治之」的方法。

並且這是在不使用量化、裁剪或其他模型壓縮技術的情況下實作的

大型語言模型之所以體積大並占用大量記憶體,主要是因為它們的結構包含許多「層」。

一個LLM(Large Language Model)從嵌入投影層開始,然後是許多相同的變換器層。

一個70B的模型可以有多達80層。但在推理過程中,每一層都是獨立的,只依賴於前一層的輸出。

因此,執行完一層後,其記憶體可以被釋放,只保留該層的輸出。基於這個概念,我們用AirLLM實作了分層推理。

怎麽實作呢 ❓

在基於Transformer的LLM中進行推理時,層會依序執行。前一層的輸出為下一層的輸入。每次只執行一層。

因此,完全沒有必要把所有的層都保留在GPU記憶體中。在執行某一層時,我們可以把需要的層從磁盤載入進來,完成所有的計算,然後完全釋放記憶體。

這樣一來,每層所需的GPU記憶體只有一個變換器層的參數大小,約為整個模型的1/80,大概1.6GB。

接著使用flash attention深度最佳化CUDA記憶體存取以實作多倍的速度提升

按層分割模型檔。

使用HuggingFace Accelerate提供的meta device特性。當你透過meta device載入模型時,模型數據實際上並未讀入,只載入了程式碼。所需記憶體為0。

提供了一個`compression`參數,用於做量化。

`compression`:支持的選項:4bit, 8bit,進行4位元或8位元的塊級量化。