查看原文
科技

无法运行:Grok 开源了,但真的太大了...

lencx 浮之静 2024-03-18

前几天 Elon Musk 在线 @xAI,宣布本周将开源 Grok,没想到它真的来了,我满怀欣喜的打开 GitHub,然后...

Grok-1(Open Release of Grok-1[1])是一款由 xAI 开发的大型语言模型,拥有 3140 亿个参数,属于混合专家模型(MoE:Mixture-of-Experts model)。模型特点是在处理每个 token 时,只有 25% 的权重(weights)被激活。它是在 2023 年 10 月,利用 JAX 和 Rust 构建的自定义训练栈,从头开始训练的。

该模型的基础模型权重和网络架构现已在 GitHub 上公开发布,并未经过针对任何特定任务的微调,因此它是预训练阶段的原始模型。开源协议遵循 Apache 2.0 许可证。

Grok-1

下载运行

看上去很简单嘛,我来跑一个!

Grok-1(xai-org/grok-1[2])存储库提供了使用 JAX 框架加载和运行 Grok-1 模型的示例代码。要运行这些示例,用户需要先下载模型的检查点文件,将其放置在指定的目录中(将下载的 ckpt-0 目录放置在 checkpoint 目录中),然后执行以下命令来安装依赖并运行示例:

# download
git clone https://github.com/xai-org/grok-1.git
cd grok-1

# install
pip install -r requirements.txt
python run.py

注意:运行代码需要一台拥有充足 GPU 内存的机器,因为 Grok-1 模型的参数量非常庞大。此外,虽然存储库中包含的 MoE(混合专家模型)层的实现不是最高效的,但它被选用是为了避免使用自定义内核,从而更容易验证模型的正确性。

用户可以通过以下磁力链接使用种子客户端下载模型权重:

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

无法运行

无法运行,真的太大了...

本来还幻想本地运行一下,但当我打开磁力链接时,就彻底断绝了想法。七百多个文件,近 300G(这还只是模型下载)...

打开 Grok issues,不出所料,清一色在询问在本地运行 Grok,需要什么样的配置。

看到“遥遥领先”的 issues,真不知道该说些什么...,在开源社区总有一些人在不断拉低国人的下限...

运行要求

普通人缺的不是开源模型,而是显卡...

我翻阅了大量 issues,总结出运行 Grok 的最低配置要求(仅作为参考):

  • #3[3]:在 FP16 下大约需要 630GB 的显存,可能会达到 700GB。使用 8 个 H100 是否能运行还是未知数,我认为在它被优化之前,你无法在 CPU 上运行它。在 FP16 精度下,Grok-1 模型大约需要 630GB 至 700GB 的显存。即便配置了 8 个 NVIDIA H100 GPU,能否成功运行该模型仍不确定。在进行某些优化(如通过 GGUF[4] 工具)之前,这个模型可能无法在 CPU 上运行。

  • #24[5]:你需要拥有 TPU 或 NVIDIA/AMD 品牌的 GPU,且系统中必须装有 8 个此类设备。当前不支持 Apple silicon 设备(如 M1、M2、M3 等)。尽管 Jax 提供了一个 Metal 插件,让你可以在苹果芯片上运行 Jax(Accelerated JAX training on Mac[6]),但在使用 dm_haiku[7] 依赖时仍会遇到问题。即便克服了这些技术障碍,苹果芯片设备可能也没有足够的内存来运行如此庞大的 Grok-1 模型。

  • #25[8]:需要 8 个 GPU,每个 GPU 拥有 80GB 的显存,典型的选择是 A100 型号。即使是使用 4 个 NVIDIA 4090 显卡,也只能在 4 位量化的情况下勉强容纳模型的权重,而无法实际运行模型。此外,所需的硬件成本极高,单个 A100 的价格约为 12,000 美元,而一台配备 4 个 A100 GPU 的 NVIDIA DGX Station 的起价在 120,000 美元左右。因此,尽管技术上可行,但这样的配置对于大多数人来说是不切实际的。

从源码中(grok-1/model.py[9])可以获取到的其他相关信息(#14[10]):

  • 3140 亿参数(314B parameters)

  • 8 个专家的混合体(Mixture of 8 Experts)

  • 每个 token 使用 2 个专家(2 experts used per token)

  • 64 层(64 layers)

  • 查询的 48 个注意力头(48 attention heads for queries)

  • 键/值的 8 个注意力头(8 attention heads for keys/values)

  • 嵌入大小:6144(embeddings size: 6,144)

  • 旋转嵌入(rotary embeddings, RoPE)

  • SentencePiece 分词器;131,072 个令牌(SentencePiece tokenizer; 131,072 tokens)

  • 支持激活分片和 8 位量化(Supports activation sharding and 8-bit quantization)

  • 最大序列长度(上下文):8192 个 token(Max seq length (context): 8,192 tokens)

名词解释

JAX

JAX[11] 是一个专为加速器优化的数组计算和程序转换设计的 Python 库,主要目标是高性能数值计算和大规模机器学习。它通过改进的 Autograd[12] 版本,能够自动计算 Python 和 NumPy 函数的导数,支持复杂结构的自动微分,包括循环、分支、递归和闭包,并能进行多阶导数计算。

一个关键特性是 JAX 利用 XLA(Accelerated Linear Algebra)编译器将 NumPy 代码编译运行在 GPU 和 TPU 上,实现高效的计算。这种编译是即时进行的,无需用户干预。此外,JAX 提供了 jit 函数,允许用户将自定义 Python 函数即时编译成 XLA 优化的内核,极大提升性能。

JAX 支持的功能不仅限于自动微分和即时编译,它还包括自动向量化(通过 vmap)和单程序多数据(SPMD)并行编程(通过pmap),使得在多加速器环境下的编程变得简单。JAX 旨在使得复杂的算法能够在 Python 环境下高效实现,同时保持代码的可读性和可维护性。

尽管 JAX 是一个强大的工具,但它仍然是一个研究项目,并非正式的 Google 产品,因此可能存在一些错误和不足之处。开发团队鼓励用户尝试使用 JAX,并通过报告问题或提供反馈来帮助改进项目。

GGUF

@ggerganov[13] 开发,他也是 llama.cpp(一个流行的 C/C++ LLM 推理框架)的开发者。原先在 PyTorch 等框架中开发的模型可以转换为 GGUF 格式,以便在这些推理引擎中使用。

GGUF 是一种专为快速加载和保存模型而设计的二进制文件格式,主要用于与基于 GGML[14] 的执行器进行推理。这种格式是为了高效地进行模型推理而优化的,支持单文件部署,具有良好的可扩展性,兼容 mmap 以实现快速加载,易于使用,且包含加载模型所需的全部信息。

它是 GGML、GGMF 和 GGJT 格式的后续版本,采用键值结构存储元数据,以提高灵活性和扩展性。这种结构允许在不影响现有模型兼容性的前提下,为模型添加新的元数据或额外信息。

mmap

mmap(内存映射,memory-mapped files)是一种在文件和内存间直接映射的技术,允许应用程序直接从内存中访问文件内容,而无需进行传统的读写操作。这种技术可以显著提高文件访问的效率,特别是对于大型文件的处理。

当一个文件被 mmap 映射到内存时,操作系统会为该文件内容在内存中创建一个直接映射。应用程序可以像访问普通内存一样访问这部分映射的内存,而操作系统会负责将任何更改同步回文件本身,或者从文件中拉取最新的数据到内存中。

在机器学习和深度学习中,mmap 可以用于快速加载大型模型权重或数据集。例如,使用 mmap 可以在不将整个文件内容一次性加载到内存的情况下,实现对大型模型权重文件的快速随机访问。这样做可以显著减少模型加载时间和运行时内存消耗。

具体到 GGUF 格式,利用 mmap 可以实现对模型文件的快速加载和保存,尤其是在处理大型模型时,这种方法可以大大提高效率和性能。通过将模型文件映射到内存,模型加载和推理过程中对权重的访问将更加迅速,从而提升整体的推理速度。

References

[1]

Open Release of Grok-1: https://x.ai/blog/grok-os

[2]

xai-org/grok-1: https://github.com/xai-org/grok-1

[3]

#3: https://github.com/xai-org/grok-1/issues/3

[4]

GGUF: https://github.com/ggerganov/ggml/blob/master/docs/gguf.md

[5]

#24: https://github.com/xai-org/grok-1/issues/24

[6]

Accelerated JAX training on Mac : https://developer.apple.com/metal/jax

[7]

dm_haiku: https://github.com/google-deepmind/dm-haiku

[8]

#25: https://github.com/xai-org/grok-1/issues/25

[9]

grok-1/model.py: https://github.com/xai-org/grok-1/blob/main/model.py

[10]

#14: https://github.com/xai-org/grok-1/issues/14

[11]

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

[12]

Autograd: https://github.com/HIPS/autograd

[13]

@ggerganov: https://github.com/ggerganov

[14]

GGML: https://github.com/ggerganov/ggml

继续滑动看下一个
向上滑动看下一个

您可能也对以下帖子感兴趣

文章有问题?点此查看未经处理的缓存