grok-1_jax由xAI从头开始训练的3140亿个参数混合专家&对话问答模型

Grok1

Grok-1是由xAI从头开始训练的3140亿个参数混合专家模型。

论文

暂无

模型结构

Grok-1 是一个8个专家（2个活跃）的混合专家（MoE）模型。

算法原理

Grok-1目前的设计规格如下:

• 参数: 314B
• 体系结构: Mixture of 8 Experts (MoE)
• 专家利用: 2 experts used per token
• 层数: 64
• 多头注意力: 48 个 head 用于查询，8 个用于键 / 值（KV）
• 嵌入大小: 6,144
• Tokenization: tokenizer 词汇大小为 131,072
• 附加特性:
  • Rotary embeddings (RoPE)
  • Supports activation sharding and 8-bit 量化
• 最大序列长度(上下文): 8,192 tokens

环境配置

Docker（方法一）

-v 路径、docker_name和imageID根据实际情况修改

docker pull image.sourcefind.cn:5000/dcu/admin/base/jax:0.4.23-ubuntu20.04-dtk24.04-py310
docker run -it -v /path/your_code_data/:/path/your_code_data/ -v /opt/hyhal:/opt/hyhal:ro --shm-size=200G --privileged=true --device=/dev/kfd --device=/dev/dri/ --group-add video --name docker_name imageID bash

cd /your_code_path/grok-1_jax
pip install -r requirements.txt

Dockerfile（方法二）

cd ./docker
docker build --no-cache -t grok1:latest .

docker run -it -v /path/your_code_data/:/path/your_code_data/ -v /opt/hyhal:/opt/hyhal:ro --shm-size=200G --privileged=true --device=/dev/kfd --device=/dev/dri/ --group-add video --name docker_name imageID bash

cd /your_code_path/grok-1_jax

Anaconda（方法三）

关于本项目DCU显卡所需的特殊深度学习库可从光合开发者社区下载安装。

DTK驱动：dtk24.04
python：python3.10
jax: ≥ 0.4.23
gcc: ≥ 9.4.0

Tips：以上dtk驱动、python等DCU相关工具版本需要严格一一对应

其它非深度学习库参照requirements.txt安装：

pip install -r requirements.txt

数据集

暂无

训练

官方github未开源微调代码

推理

推理所需硬件最低配置参考：

device(s)：K100 64G 8卡

Driver version: 5.16.29.22.20

Vbios version: 5.223.001200k.548101

linux: ubuntu20.04

处理器： Hygon C86 7375 32-core Processor

CPU(s)：128

内存：≥ 600G, 1T以上最优

python run.py

result

日志信息可以参考run.log文件，测试结果如下图所示。

精度

暂无

应用场景

算法类别

对话问答

热点应用行业

制造,广媒,家居,教育

预训练权重

下载地址SCNet_grok-1

源码仓库及问题反馈

• ModelZoo / Grok-1_jax · GitLab

参考资料

• https://github.com/xai-org/grok-1