今天小编分享的科学经验：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
