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的模型权重。