馬斯克說到做到開源 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