3140參數(shù)Grok-1推理加速3.8倍，PyTorch+HuggingFace版來(lái)了
馬斯克說(shuō)到做到開(kāi)源Grok-1，開(kāi)源社區(qū)一片狂喜。
但基于Grok-1做改動(dòng)or商用，都還有點(diǎn)難題：
Grok-1使用Rust+JAX構(gòu)建，對(duì)于習(xí)慣Python+PyTorch+HuggingFace等主流軟件生態(tài)的用戶上手門檻高。
△圖注：Grok登上GitHub熱度榜世界第一
Colossal-AI團(tuán)隊(duì)最新成果，解大家燃眉之急，提供方便易用的Python+PyTorch+HuggingFace Grok-1，能將推理時(shí)延加速近4倍！
現(xiàn)在，模型已在HuggingFace、ModelScope上發(fā)布。
HuggingFace下載鏈接：
https://huggingface.co/hpcai-tech/grok-1
ModelScope下載鏈接：
https://www.modelscope.cn/models/colossalai/grok-1-pytorch/summary
性能優(yōu)化
結(jié)合Colossal-AI在AI大模型系統(tǒng)優(yōu)化領(lǐng)域的豐富積累，已迅速支持對(duì)Grok-1的張量并行。
在單臺(tái)8H800 80GB服務(wù)器上，推理性能相比JAX、HuggingFace的auto device map等方法，推理時(shí)延加速近4倍。
使用教程
下載安裝Colossal-AI后，啟動(dòng)推理腳本即可。
./run_inference_fast.sh hpcaitech/grok-1模型權(quán)重將會(huì)被自動(dòng)下載和加載，推理結(jié)果也能保持對(duì)齊。如下圖中Grok-1 greedy search的運(yùn)行測(cè)試。
更多詳情可參考grok-1使用例：
https://github.com/hpcaitech/ColossalAI/tree/main/examples/language/grok-1
龐然大物Grok-1
此次開(kāi)源，xAI發(fā)布了Grok-1的基本模型權(quán)重和網(wǎng)絡(luò)架構(gòu)。
具體來(lái)說(shuō)是2023年10月預(yù)訓(xùn)練階段的原始基礎(chǔ)模型，沒(méi)有針對(duì)任何特定應(yīng)用（例如對(duì)話）進(jìn)行微調(diào)。
結(jié)構(gòu)上，Grok-1采用了混合專家（MoE）架構(gòu)，包含8個(gè)專家，總參數(shù)量為314B（3140億），處理Token時(shí)，其中的兩個(gè)專家會(huì)被激活，激活參數(shù)量為86B。
單看這激活的參數(shù)量，就已經(jīng)超過(guò)了密集模型Llama 2的70B，對(duì)于MoE架構(gòu)來(lái)說(shuō)，這樣的參數(shù)量稱之為龐然大物也毫不為過(guò)。
更多參數(shù)信息如下：
- 窗口長(zhǎng)度為8192tokens，精度為bf16
- Tokenizer vocab大小為131072（2^17），與GPT-4接近；
- embedding大小為6144（48×128）；
- Transformer層數(shù)為64，每層都有一個(gè)解碼器層，包含多頭注意力塊和密集塊；
- key value大小為128；
- 多頭注意力塊中，有48 個(gè)頭用于查詢，8 個(gè)用于KV，KV 大小為 128；
- 密集塊（密集前饋塊）擴(kuò)展因子為8，隱藏層大小為32768
在GitHub頁(yè)面中，官方提示，由于模型規(guī)模較大（314B參數(shù)），需要有足夠GPU和內(nèi)存的機(jī)器才能運(yùn)行Grok。
這里MoE層的實(shí)現(xiàn)效率并不高，選擇這種實(shí)現(xiàn)方式是為了避免驗(yàn)證模型的正確性時(shí)需要自定義內(nèi)核。
模型的權(quán)重文件則是以磁力鏈接的形式提供，文件大小接近300GB。
值得一提的是，Grok-1采用的是Apache 2.0 license，商用友好。
目前Grok-1在GitHub上的標(biāo)星已達(dá)到43.9k Stars。
量子位了解，Colossal-AI將在近期進(jìn)一步推出對(duì)Grok-1在并行加速、量化降低顯存成本等優(yōu)化，歡迎持續(xù)關(guān)注。
Colossal-AI開(kāi)源地址：https://github.com/hpcaitech/ColossalAI
