Grok-1

本仓库包含用于加载和运行Grok-1开放权重模型的JAX示例代码。

确保下载检查点并将"ckpt-0"目录放置在"checkpoints"中 - 参见下载权重

然后，运行

pip install -r requirements.txt
python run.py

以测试代码。

该脚本加载检查点并基于测试输入从模型中采样。

由于模型规模较大（3140亿参数），需要具有足够GPU内存的机器才能使用示例代码测试模型。 本仓库中MoE层的实现并不高效。选择这种实现方式是为了避免使用自定义内核来验证模型的正确性。

模型规格

Grok-1目前设计有以下规格：

• 参数数量： 3140亿
• 架构： 8专家混合（MoE）
• 专家利用： 每个token使用2个专家
• 层数： 64
• 注意力头： 查询48个，键/值8个
• 嵌入大小： 6,144
• 分词： 使用131,072个token的SentencePiece分词器
• 其他特性：
  • 旋转位置编码（RoPE）
  • 支持激活分片和8位量化
• 最大序列长度（上下文）： 8,192个token

下载权重

您可以使用torrent客户端和以下磁力链接下载权重：

magnet:?xt=urn:btih:5f96d43576e3d386c9ba65b883210a393b68210e&tr=https%3A%2F%2Facademictorrents.com%2Fannounce.php&tr=udp%3A%2F%2Ftracker.coppersurfer.tk%3A6969&tr=udp%3A%2F%2Ftracker.opentrackr.org%3A1337%2Fannounce

或直接使用HuggingFace 🤗 Hub：

git clone https://github.com/xai-org/grok-1.git && cd grok-1
pip install huggingface_hub[hf_transfer]
huggingface-cli download xai-org/grok-1 --repo-type model --include ckpt-0/* --local-dir checkpoints --local-dir-use-symlinks False

许可证

本发布中的代码和相关Grok-1权重采用Apache 2.0许可证。该许可证仅适用于本仓库中的源文件和Grok-1的模型权重。