Grok-1:马斯克的xAI开源LLM(尝鲜体验)

前言¶

Grok-1 (格鲁克) 是马斯克的 xAI 开源的大语言模型(LLM)，它没有针对某个特殊行业进行微调，使用 JAX 库和 Rust 从底层开发的基础模型，有 314B 个参数。

• OpenAI Grok论文代码: https://github.com/openaI/grok
• xAI Grok开源代码: https://github.com/xai-org/grok-1
• 讨论区: https://github.com/xai-org/grok-1/discussions

JAX与Rust¶

开源文档说明信息里面特别提到这两个技术：

• JAX库
• Rust语言

这里顺便简单了解一下。

JAX库¶

• https://github.com/google/jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

JAX is NumPy on the CPU, GPU, and TPU, with great automatic differentiation for high-performance machine learning research.

JAX is a Python library for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning.

Rust语言¶

• https://www.rust-lang.org/

A language empowering everyone to build reliable and efficient software.

$ curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
$ rustup update

安装Grok-1¶

$ git clone https://github.com/xai-org/grok-1.git

安装依赖软件包¶

$ pip install -r requirements.txt

获取checkpoint和权重文件(chpt-0/)¶

• 磁力下载(319GB):

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-cli 工具:

$ 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

测试¶

$ python run.py

默认是8个GPU:

# line 60
local_mesh_config=(1,8)

如果只有1个GPU，则改成:

local_mesh_config=(1,1)

总结(运行失败)¶

在存储空间方面，权重文件大概需要320GB左右磁盘空间，这至少需要一天才能完全下载。

另外，为了运行，磁盘还需要有额外的300GB以上空间，因为运行期间会创建临时目录。

其它硬件配置，据说至少要满足:

GPU: 8 x A100 80G
Total VRAM: 640G
CPU: 2 x Xeon 8480+
RAM: 1.5 TB

目前在自己可怜的单个RTX3080机器上没有运行成功。

这应该不是平民玩家可以玩得起的，目前只能持续关注，等待Mini简配版本发布。