xai-org / grok-1

xai-org/grok-1

open_in_new前往仓库

Grok 开源发布

项目概览

项目概述

Grok-1 是 xAI 发布的开源大语言模型，本仓库提供用于加载和运行 Grok-1 开放权重模型的 JAX 示例代码。模型拥有 3140 亿参数，采用混合专家（MoE）架构，是当前参数量最大的开源模型之一。

核心功能

  • 加载 Grok-1 预训练权重并进行文本采样。
  • 支持 8-bit 量化（需自行配置）。
  • 提供 HuggingFace 和 BitTorrent 两种权重下载方式。

适用与不适用场景

适用场景：

  • 研究 MoE 架构、大模型推理。
  • 在拥有充足 GPU 内存（如多张 A100/H100）的环境下进行模型验证。

不适用场景：

  • 资源受限环境（单卡消费级 GPU）。
  • 生产环境部署（示例代码未优化，MoE 层效率低）。

技术架构与依赖

  • 语言：Python
  • 框架：JAX
  • 模型架构：MoE（8 个专家，每 token 激活 2 个），64 层，48 个查询注意力头，8 个键/值头，嵌入维度 6144，上下文长度 8192。
  • 分词器：SentencePiece，词表大小 131072。
  • 依赖：见 requirements.txt（仓库未提供具体列表）。

安装与快速开始

  1. 克隆仓库： 
    git clone https://github.com/xai-org/grok-1.git
cd grok-1
  2. 安装依赖： 
    pip install -r requirements.txt
  3. 下载权重（任选一种）：
    • 使用 HuggingFace： 
      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
    • 使用 BitTorrent：通过磁力链接下载，并将 ckpt-0 目录放入 checkpoints。
  4. 运行示例： 
    python run.py

典型使用方法

运行 python run.py 会加载检查点并对测试输入进行采样。用户可修改 run.py 中的输入文本以测试不同提示。

配置与部署要点

  • 需要足够 GPU 内存（模型 314B 参数，建议多卡 A100/H100）。
  • 仓库代码未优化 MoE 层，仅用于验证正确性，不适合生产部署。
  • 可通过配置启用激活分片和 8-bit 量化（具体配置方式仓库未提供）。

限制、风险与许可证

  • 限制：模型体积巨大，推理资源要求高；示例代码效率低。
  • 风险：模型可能产生不当内容，使用者需自行评估。
  • 许可证：Apache-2.0，适用于仓库源码和 Grok-1 权重。

官方链接

信息来源和分析时间