如何在家用设备上运行 Llama 3.1 405B？搭建 AI 集群

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

Tensor Parallism

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 项目。

Distributed Llama

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

Root Node ---充当集群根节点的应用程序，用于协调集群。
Worker Node ---充当 worker 的应用程序，从根节点执行指令。

目前，分布式 Llama 仅支持 CPU 推理，但将来会发生变化。

因此，如果您的主集群由 4 台设备组成，则应在第一台设备上运行根节点，在其余设备上运行 3 个 Worker 节点。分布式 Llama 在所有设备上拆分 RAM 使用。例如，如果 LLM 模型需要 238 GB 的 RAM，则每个节点都应该有 238 GB /n RAM。根节点是一个例外，它需要的 RAM 238 GB /n 比它在内存中保留几个额外的层多出几个百分点。

运行 405B 模型

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

复制代码

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

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

复制代码

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

该--nthreads参数定义应使用多少个 CPU 内核进行处理。您应该将此设置为设备中的 CPU 内核数。如您所见，worker 不需要模型文件。只有根节点才需要这些文件。开始时，根节点将模型的所有切片分发给 worker 节点。

在运行根节点之前，我们需要将 Llama 3 405B 模型下载到根设备并将其转换为分布式 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参数接受带有 worker 节点端口的 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 disk参数，该参数通过将 KV 缓存移动到磁盘来减少 RAM 使用量。Llama 3.1 模型需要 ~34 GB 的 RAM 才能将完整上下文（F32）存储在内存（131k 令牌）中。通过设置此参数，可以减少 RAM 使用量，但需要额外的磁盘空间。请注意，KV 缓存被拆分到所有节点上，因此您需要为每个节点设置此选项。

减少 RAM 使用率的第二个选项是使用 argument --max-seq-len 2048 。如果您不需要完整的上下文大小，则可以减小它，这同时会减少内存消耗。