今天小編分享的科學經驗：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
