3140参数Grok-1推理加速3.8倍,PyTorch+HuggingFace版来了

人工智能 新闻
Grok-1使用Rust+JAX构建,对于习惯Python+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上发布。

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在并行加速、量化降低显存成本等优化,欢迎持续关注。

Colossal-AI开源地址:https://github.com/hpcaitech/ColossalAI

责任编辑:张燕妮 来源: 量子位
相关推荐

2024-03-18 08:47:34

马斯克GrokOpenAI

2024-06-28 13:42:07

2024-03-18 10:17:00

开源AI

2024-03-18 15:00:48

Grok现已开源Grok-1混合专家模型

2024-01-24 13:11:00

AI模型

2024-03-28 12:53:00

AI数据

2024-03-19 14:00:59

2024-07-19 09:59:31

2023-12-11 15:40:32

PyTorch代码大模型

2019-09-02 14:58:03

深度学习编程人工智能

2023-12-03 08:49:38

微软开源

2023-09-01 15:22:49

人工智能数据

2024-07-08 13:04:01

2023-06-21 13:20:14

系统模型

2023-11-30 18:25:57

数据训练

2024-09-19 13:32:24

2023-05-23 14:06:53

微软研究

2024-04-03 12:32:00

数据训练

2024-03-29 14:04:00

模型训练

2024-09-09 16:22:51

点赞
收藏

51CTO技术栈公众号