馬斯克說到做到開源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://huggingface.co/hpcai-tech/grok-1
ModelScope下載鏈接：
https://www.modelscope.cn/models/colossalai/grok-1-pytorch/summary
性能優化
結合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://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在並行加速、量化降低顯存成本等優化，歡迎持續關注。
---[明敏 發自 : 凹非寺*量子位 : 公衆號 QbitAI/來源: 量子位]
Colossal-AI開源地址：https://github.com/hpcaitech/ColossalAI
