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翻译整理，转载请标明出处