这个仓库包含了加载和运行 Grok-1 开源权重模型的 JAX 示例代码。

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

然后运行

shell
1pip install -r requirements.txt
2python run.py

来测试代码。

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

由于模型的大小(314B 参数),需要具有足够 GPU 内存的机器才能使用示例代码测试模型。此存储库中 MoE 层的实现效率不高。选择这种实现是为了避免需要自定义内核来验证模型的正确性。

模型规格

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

• 参数: 314B
• 架构: 8 专家混合 (MoE)
• 专家利用率: 每个令牌使用 2 个专家
• 层数: 64
• 注意力头: 查询 48 个,键/值 8 个
• 嵌入大小: 6,144
• 分词: 131,072 个令牌的 SentencePiece 分词器
• 额外特性:
  • 旋转嵌入 (RoPE)
  • 支持激活分片和 8 位量化
• 最大序列长度(上下文): 8,192 个令牌

下载权重

你可以使用种子客户端和这个磁力链接下载权重:

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