Llama3.1 405B分布式推理

在开源 LLM 模型和封闭 LLM 模型的竞争中,开源模型的最大优势在于你可以在本地运行它们。你不需要依赖外部提供商,也不需要支付除电力和硬件成本之外的任何额外费用。然而,随着模型规模的增加,这种优势开始减弱。运行需要大量内存的大型模型并不容易。幸运的是,张量并行和分布式推理可能会有所帮助。

1、张量并行

LLM 中的大多数计算都涉及矩阵乘法,这占所有计算的 97-98% 左右。矩阵乘法很容易在多个 CPU/GPU 核心上并行化。我们可以在多个设备上做同样的事情。设备可以这样分割,每个设备只计算矩阵乘法的一部分。如果单个设备可以在 n 秒内计算矩阵乘法,那么两个设备应该在 n/2 秒内计算它!这就是张量并行。

这听起来很有希望,但这里的主要瓶颈是同步。我们可以加快乘法,但在某个时候,我们需要同步神经网络的状态。这需要时间。专业的 AI 集群使用高级链接在实现非常高传输速度的 GPU 之间进行通信(如 NVLink)。然而,家用设备的以太网速度很慢。但令人惊讶的是,如果模型执行器的架构旨在减少传输大小,那么同步 LLM 所需的数据量可以非常低。例如,如果集群由 2 个设备组成,则量化的 Llama 3 8B 到 Q40 格式(6.3 GB)每个令牌只需要 1 MB 的数据来同步。这是非常非常低的。

我们在这里。张量并行性加快了推理速度,但同步减慢了推理速度。这两个因素的结合将决定最终的性能。如果有 8 台设备,并能通过快速链接将它们连接起来,你将观察到显著的加速(通过 USB4 进行同步似乎非常有希望,可以实现 10 到 20 Gbps 的速度)。

那么,我们如何在家中运行大型模型?你需要一个实现这些想法的项目。让我介绍一下 Distributed Llama 项目。

2、Distributed Llama

Distributed Llama 是一个允许你在多个设备上运行 LLM 模型的项目。它使用张量并行性,并针对同步所需的少量数据进行了优化。Distributed Llama 区分了你可以在设备上运行的两种类型的节点:

  • 根节点 — 充当集群根节点的应用程序,协调集群。
  • 工作节点 — 充当工作器的应用程序,从根节点执行指令。

目前,Distributed Llama 仅支持 CPU 推理,但这在未来会发生变化。

AI 集群拓扑,4 个设备,共 256 GB RAM

因此,如果你的主集群由 4 个设备组成,则应在第一个设备上运行根节点,在其余设备上运行 3 个工作节点。Distributed Llama 将 RAM 使用量分配到所有设备。例如,如果 LLM 模型需要 238 GB 的 RAM,则每个节点应具有 238 GB /n 的 RAM。根节点是个例外,它需要的 RAM 比 238 GB /n 多几个百分点,因为它需要在内存中保留一些额外的层。

3、运行 405B 模型

要运行 Llama 3.1 405B,我们需要克隆 Distributed Llama 存储库并在你要用于推理的所有设备上运行 dllama 应用程序。需要 G++ 或类似的编译器。

git clone https://github.com/b4rtaz/distributed-llama.git
make dllama

然后,需要将所有设备连接到同一个本地网络。你可以使用任何以太网交换机来实现这一点。正如我之前提到的,同步时间是一个重要因素,因此应该使用尽可能快的交换机。千兆以太网是最低要求。你还可以考虑通过 USB4 连接设备并创建 USB4 网状网络。接下来,需要在工作设备上运行工作节点:

./dllama worker --port 9998 --nthreads 4

--nthreads 参数定义应使用多少个 CPU 核心进行处理。你应该将其设置为设备中的 CPU 核心数。如你所见,工作节点不需要模型文件。这些文件仅对根节点是必需的。一开始,根节点将模型的所有切片分发到工作节点。

在运行根节点之前,我们需要将 Llama 3 405B 模型下载到根设备并将其转换为 Distributed Llama 格式。你可以手动执行此操作,也可以从 Huggingface 下载预转换的权重。使用 Distributed Llama 存储库中的 launch.py​​ 脚本,你可以通过执行单个命令下载模型和标记器。所有文件都将放置在 models 文件夹中。

launch.py llama3_1_405b_instruct_q40

确保你已接受 Huggingface 上的 Llama 3.1 许可证。确保你的磁盘上有大约 240GB 的可用空间。

现在可以在根节点上运行推理。

./dllama inference \
   --model models/llama3_1_405b_instruct_q40/dllama_model_llama3_1_405b_instruct_q40.m \
   --tokenizer models/llama3_1_405b_instruct_q40/dllama_tokenizer_llama3_1_405b_instruct_q40.t \
   --buffer-float-type q80 \
   --prompt "Hello world" \
   --steps 64 \
   --nthreads 4 \
   --workers 10.0.0.1:9998 10.0.0.2:9998 10.0.0.3:9998

请注意, --workers 参数接受带有工作节点端口的 IP 地址。地址以空格分隔。此外,你可以通过设置 --steps N 参数来定义你期望的单词预测数量。

如果想运行支持 /v1/chat/completions 端点的 API 服务,你应该构建 dllama-api 应用程序并在根设备上运行它,而不是 dllama inference 。

./dllama-api \
  --model models/llama3_1_405b_instruct_q40/dllama_model_llama3_1_405b_instruct_q40.m \
  --tokenizer models/llama3_1_405b_instruct_q40/dllama_tokenizer_llama3_1_405b_instruct_q40.t \
  --buffer-float-type q80 \
  --max-seq-len 2048 \
  --nthreads 4

Distributed Llama 还支持 --kv-cache-storage 磁盘参数,该参数通过将 KV 缓存移至磁盘来减少 RAM 使用量。Llama 3.1 型号需要约 34 GB 的 RAM 来将完整上下文 (F32) 存储在内存中 (131k 个令牌)。通过设置此参数,你可以减少 RAM 使用量,但需要额外的磁盘空间。请注意,KV 缓存分布在所有节点上,因此你需要为每个节点设置此选项。

减少 RAM 使用量的第二个选项是使用 --max-seq-len 2048 参数。如果你不需要完整的上下文大小,则可以减小它,这将同时减少内存消耗。


原文链接:How to Run Llama 3.1 405B on Home Devices? Build AI Cluster!

BimAnt翻译整理,转载请标明出处