3140參數Grok-1推理加速3.8倍，PyTorch+HuggingFace版來了

馬斯克說到做到開源Grok-1，開源社群一片狂喜。

但基於Grok-1做改動or商用，都還有點難題：

Grok-1使用Rust JAX建構，對於習慣Python PyTorch HuggingFace等主流軟體生態的用戶上手門檻高。

△圖註：Grok登上GitHub熱度榜世界第一

Colossal-AI團隊最新成果，解大家燃眉之急，提供方便易用的Python PyTorch HuggingFace Grok-1，能將推理延遲加速近4倍！

現在，模型已在HuggingFace、ModelScope上發布。

HuggingFace下載連結：
https://www.php.cn/link/335396ce0d3f6e808c26132f91916eae

ModelScope下載連結：
https: //www.php.cn/link/7ae7778c9ae86d2ded133e891995dc9e

性能優化

結合Colossal-AI在AI大模型系統優化領域的豐富積累，已迅速支持對Grok-1的張量並行。

在單一8H800 80GB伺服器上，推理效能相較於JAX、HuggingFace的auto device map等方法，推理延遲加速近4倍。

使用教學

下載安裝Colossal-AI後，啟動推理腳本即可。

./run_inference_fast.sh hpcaitech/grok-1

模型權重將會自動下載和​​加載，推理結果也能保持對齊。如下圖Grok-1 greedy search的運行測試。

更多詳情可參考grok-1使用範例：
https://www.php.cn/link/e2575ed7d2c481c414c10e688bcbc4cf

龐然大物Grok-1

這次開源，xAI發布了Grok-1的基本模型權重和網路架構。

具體來說是2023年10月預訓練階段的原始基礎模型，沒有針對任何特定應用（例如對話）進行微調。

結構上，Grok-1採用了混合專家（MoE）架構，包含8個專家，總參數量為314B（3140億），處理Token時，其中的兩個專家會被激活，激活參數量為86B。

單看這活化的參數量，就已經超過了密集模型Llama 2的70B，對於MoE架構來說，這樣的參數量稱之為龐然大物也毫不為過。

更多參數資訊如下：

• 視窗長度為8192tokens，精確度為bf16
• Tokenizer vocab大小為131072（2^17），與GPT-4接近；
• embedding大小為6144（48×128）；
• Transformer層數為64，每層都有一個解碼器層，包含多頭注意力塊和密集塊；
• key value大小為128；
• 多頭注意力區塊中，有48 個頭用於查詢，8 個用於KV，KV 大小為128；
• 密集區塊（密集前饋區塊）擴充因子為8，隱藏層大小為32768

#在GitHub頁面中，官方提示，由於模型規模較大（314B參數），需要有足夠GPU和記憶體的機器才能運作Grok。

這裡MoE層的實作效率並不高，選擇這種實作方式是為了避免驗證模型的正確性時需要自訂核心。

模型的權重檔案則是以磁力連結的形式提供，檔案大小接近300GB。

值得一提的是，Grok-1採用的是Apache 2.0 license，商用友善。

目前Grok-1在GitHub上的標星已達到43.9k Stars。

量子位元了解，Colossal-AI將在近期進一步推出對Grok-1在並行加速、量化降低顯存成本等優化，歡迎持續關注。

Colossal-AI開源位址：#https://www.php.cn/link/b9531e7d2a8f38fe8dcc73f58cae9530

#################################################################################################################

以上是3140參數Grok-1推理加速3.8倍，PyTorch+HuggingFace版來了的詳細內容。更多資訊請關注PHP中文網其他相關文章！