今天小編分享的科學經驗: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 上發布。
性能優化
結合 Colossal-AI 在 AI 大模型系統優化領網域的豐富積累,已迅速支持對 Grok-1 的張量并行。
在單台 8H800 80GB 伺服器上,推理性能相比 JAX、HuggingFace 的 auto device map 等方法,推理時延加速近 4 倍。
使用教程
./run_inference_fast.sh hpcaitech/grok-1
更多詳情可參考 grok-1 使用例:
https://github.com/hpcaitech/ColossalAI/tree/main/examples/language/grok-1
龐然大物 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://github.com/hpcaitech/ColossalAI