ZeRO-Inference：让大型模型推理民主化

介绍

图像、语音和自然语言等人工智能 (AI) 领域的当前趋势表明，通过增加模型规模可以提高模型质量。例如，在自然语言处理中，最先进 (SOTA) 模型在不到四年的时间里从 3 亿个参数 (Bert-Large) 增长到 5000 亿个参数 (Megatron-Turing-530B)。然而，模型规模的这种急剧增长极大地提高了训练、微调或推断这些模型的 GPU 成本，使大多数用户无法负担。为了让每个人都能使用 AI 创新，Hugging Face (BigScience)、Meta 和 Yandex 等大型组织最近公开了预训练的大型模型。不幸的是，即使这些公开可用的模型也没有被广泛使用，因为许多用户负担不起数十个 GPU 来为推断计算提供支持。例如，Megatron-Turing-530B (自然语言的 SOTA 模型) 的半精度推断计算需要至少 40 个 A100-40GB GPU，这对许多学生、模型科学家、爱好者和小企业来说是无法负担的，而这些群体可以从使用这些强大的模型中获益。因此，一个真正的担忧是，如果模型规模的急剧增长持续下去，那么越来越多的用户将被排除在这些 AI 创新的益处之外。

DeepSpeed 是微软 AI at Scale Initiative 的一部分，开发了 ZeRO-Inference 技术来解决 AI 民主化面临的这些障碍。ZeRO-Inference 来自 ZeRO 技术家族，ZeRO 技术家族是一组强大的内存和并行性优化，可用于在现代 GPU 集群上高效地进行大规模模型训练和推理。DeepSpeed 之前开发了 ZeRO-Infinity，这项技术利用异构内存 (GPU、CPU 和 NVMe) 将模型训练高效地扩展到极端水平。ZeRO-Inference 适应和优化了 ZeRO-Infinity 的技术，用于在 GPU 上进行模型推理，方法是在 CPU 或 NVMe 内存中托管模型权重，从而在 GPU 中不托管任何 (零) 权重。这种方法的灵感来自于这样的观察：大多数商品计算设备 (例如笔记本电脑、台式机、工作站等) 中 CPU 和 NVMe 内存的总容量约为 TB 级，足以托管已知最大的模型以进行推断计算。通过利用这种非 GPU 内存，ZeRO-Inference 使得即使在单个 GPU 上也能进行大型模型 (具有数百亿个参数) 的推理计算，从而使几乎所有人都能够使用大型模型推理。此外，通过使用 CPU 或 NVMe 内存 (其价格明显低于 GPU 内存) 极大地降低了 GPU 内存需求，它极大地降低了大型模型推理的成本，为 SOTA 模型提供了一条经济实惠的推断路径。

ZeRO-Inference 的工作原理

大型模型推理对计算资源的要求很高，这意味着需要 GPU 等加速器才能高效地执行。因此，对于 GPU 预算有限的大型模型推理来说，一个重要的设计决策是如何在模型权重、推理输入和中间结果之间分配 GPU 内存。

卸载所有模型权重

ZeRO-Inference 将整个模型权重固定在 CPU 或 NVMe 中 (只要足够容纳整个模型即可)，并将权重逐层流式传输到 GPU 中进行推断计算。在计算完一层后，输出将保留在 GPU 内存中，作为下一层的输入，而层权重占用的内存将被释放，以便下一层使用。因此，模型推理时间由在 GPU 上计算层的时间和通过 PCIe 获取层的时间组成。对于大型模型推理，这种方法提供了如下所述的扩展和效率优势。

ZeRO-Inference 通过两种方式提供扩展优势。首先，通过在任何时候只在 GPU 内存中保留一个 (或几个) 模型层，ZeRO-Inference 显着减少了推断大型模型所需的 GPU 内存量。对于当前具有大约一百个层的 SOTA 模型 (例如，GPT3-175B 和 Megatron-Turing-530B 中分别有 96 个和 105 个层)，ZeRO-Inference 将 GPU 内存需求降低了两个数量级。例如，使用 ZeRO-Inference，Megaton-Turing-530B 的半精度推理的 GPU 内存消耗从 1TB 降至 10GB。其次，通过将模型放入 CPU 或 NVMe 内存中 (其价格比 GPU 内存低几个数量级)，与将整个模型放入 GPU 内存中的方法相比，ZeRO-Inference 使得扩展到未来的 SOTA 模型 (例如，具有万亿或十万亿个参数的模型) 更为经济实惠。

尽管通过 PCIe 互连从 CPU 或 NVMe 获取模型权重存在延迟，但 ZeRO-Inference 仍能为吞吐量导向的推理应用提供高效的计算。这主要是因为，通过将模型对 GPU 内存的使用限制在一个或几个权重层，ZeRO-Inference 可以使用大部分 GPU 内存来支持大量以长序列或大批量形式存在的输入 token。一个大型模型层需要大量的计算，尤其是在处理具有许多输入 token 的输入时。例如，一个 GPT3-175B 层需要大约 7 TFlops 来处理批量大小为 1 且序列长度为 2048 的输入。因此，对于具有长序列长度和大批量大小的推理场景，计算时间会主导获取模型权重的延迟，最终提高效率。总之，ZeRO-Inference 的策略是利用 GPU 内存来支持大量输入 token，这使得大型模型能够实现高性能推理。

优化

为了进一步提高系统效率，ZeRO-Inference 利用了另外两种优化措施来减少从 CPU 或 NVMe 内存获取层权重到 GPU 内存的延迟。

第一个优化涉及将层的获取与先前层的计算重叠，也称为层预取。层预取使 ZeRO-Inference 能够隐藏预取层的部分传输延迟。这在计算时间不够长或无法足够增加 (例如，使用更大的批量大小) 以主导获取层权重的延迟时尤其有用。

第二个优化措施适用于在多个 GPU 上进行推理，它涉及通过使用每个 GPU 仅获取层的一部分来在多个 GPU 上并行化每个层的获取。以这种方式利用 GPU 的总 PCIe 链接基本上会线性地提高传输带宽，从而降低延迟。使用这种方法，将层获取到 GPU 内存中将分为两个阶段。首先，每个 GPU 独立地通过 PCIe 将层的某个分区获取到其内存中。此时，每个 GPU 上只会驻留层的某个分区。接下来，每个 GPU 通过通过高带宽 GPU-GPU 互连 (例如 NVLink、xGMI 等) 从其他 GPU 中获取缺少的层片段来组装完整的层。由于 GPU-GPU 互连带宽通常比 PCIe 带宽高一个数量级，因此可以使用高效的多 GPU 或多节点通信原语 (例如 NCCL 或 RCCL 全部收集) 来高效地将完整层组装到所有 GPU 上，与 PCIe 延迟相比，其延迟可以忽略不计。

替代方法：在 GPU 内存中托管部分模型权重

ZeRO-Inference 的另一种方法是将尽可能多的模型权重固定到 GPU 内存中，并在计算需要时从 CPU 或 NVMe 中获取其余权重。这种方法的好处是可以避免获取已固定在 GPU 内存中的权重的延迟。但是，这种方法有两个缺点： (i) 对于数百亿个参数的模型，延迟节省可以忽略不计，因为只有少量权重可以容纳在 GPU 内存中，(ii) 即使可以容纳模型权重的相当一部分 (例如，对于约 100 亿个参数的模型，> 50%)，剩余的 GPU 内存只能容纳小批量大小，这会影响推理吞吐量。我们将在后面展示评估结果，以证明这种方法不是最佳选择。

单个 GPU 上的模型扩展

与将模型托管在 GPU 内存中 (即 HBM) 的基线相比，ZeRO-Inference 使得单个 GPU 上的推理能够实现显着的模型扩展。例如，我们考虑使用 NVIDIA DGX2 系统中的单个 NVIDIA Tesla V100 GPU 进行半精度模型推理。虽然 V100 GPU 具有 32GB 的内存，但该系统配备了 1.5TB 的 CPU DRAM 和 30TB 的 NVMe 存储器。在 GPU 上进行推断计算时，支持的最大模型大小取决于托管模型的内存。图 1 显示了该系统中使用 ZeRO-Inference 进行 GPU 推断时可以实现的模型规模。相比之下，基线无法支持超过 160 亿个参数的模型进行 GPU 推断1。相比之下，ZeRO-Inference 具有将模型托管在与 HBM 不同的内存 (DRAM 或 NVMe) 中的灵活性。这种灵活性使 ZeRO-Inference 能够支持比基线大得多的模型。例如，通过在 NVMe 内存中托管模型，Zero-Inference 可以支持模型，其 GPU 推断参数最多可以达到 15 万亿个，这几乎是基线的 1000 倍。从图 1 中可以得到一个实用的结论，即 ZeRO-Inference 使得当前 SOTA 模型能够在单个 GPU 上进行推断计算，因为这些模型的参数少于 15 万亿个。

Model-Scaling

Token 生成性能

一个重要的推理工作负载是基于输入提示的 token 生成。在此工作负载中，模型将文本序列作为输入提示提供，并且基于此提示，模型生成可配置长度的输出文本。我们使用此工作负载来演示 ZeRO-Inference 的性能。此工作负载包括两个阶段： (1) 提示处理阶段，模型在其中处理输入提示，以及 (2) 生成阶段，模型在其中生成输出 token。

ZeRO-Inference 面向吞吐量导向的推理应用，因此我们用于此工作负载的性能指标是在生成阶段每秒生成的 token 数量。在我们的实验中，我们使用 Hugging Face token 生成管道来衡量使用贪婪搜索算法生成十个输出 token (给定四个 token 的输入提示) 的性能。我们实验中的生成管道使用 KV 缓存优化来提高性能，方法是缓存生成的 token 以避免重新计算。我们考虑了 ZeRO-Inference 设计选择和优化的三个方面对性能的影响： (1) 完全卸载模型权重而不是部分卸载，(2) 在使用之前预取层权重，以及 (3) 使用多个 GPU 来并行化通过 PCIe 获取层。此外，我们还测量了改变输出 token 数量对性能的影响。

模型

在我们的实验中，我们使用了表 1 中列出的三个公开可用的大型语言模型。我们将这些模型配置为进行半精度推理计算。由于这些模型的参数量大于 GPU 内存，因此需要 ZeRO-Inference 才能在单个 V100-32GB 上进行推理。

Public-models

完全卸载与部分卸载模型权重

ZeRO-Offload 的一个关键设计选择是将所有超过 GPU 内存的模型权重卸载到 CPU 或 NVMe 内存中，而不是在 GPU 内存中保留部分权重。我们直觉地认为，对于以吞吐量为导向的推理应用程序，由完整卸载实现的更大批次大小将带来比部分卸载更好的性能。在表 2中，我们展示了在单个 V100-32GB 上进行 OPT-30B token 生成时，将模型权重完全卸载与在 GPU 内存中保留部分权重（即 100 亿和 120 亿个参数2）的比较结果。结果表明，对于 CPU 内存（每秒 43 个 token）和 NVMe 内存（每秒 30 个 token），完全卸载都提供了最佳性能。对于 CPU 和 NVMe 内存，完全卸载分别比 180 亿和 200 亿参数的部分卸载快 1.3 倍和 2.4 倍。完全卸载的性能优势来自与部分卸载选项相比更大的批次大小。因此，当模型无法完全放入 GPU 时，使用 GPU 内存来增加批次大小而不是部分容纳模型，将导致更快的 token 生成。

Full-offload

预取层权重

ZeRO-Inference 在使用前预取层，并与当前层的计算重叠，以隐藏层传输延迟。我们在单个 V100-32GB 上测量了预取对 token 生成性能的影响，并在表 3中汇总了结果。我们观察到预取并没有改善 CPU 卸载。这是因为 token 生成中相对较短的序列（即少于 50 个 token）导致层计算时间不足以隐藏从 CPU 获取层的时间的很大一部分。相反，预取通过 1.13 倍、1.14 倍和 1.21 倍分别提高了 OPT-30B、OPT-175B 和 BLOOM-176B 的 NVMe 卸载性能。这是因为从 NVMe 通过 CPU 内存传输权重，使预取能够将从 CPU 到 GPU 内存的传输与从 NVMe 到 CPU 的传输重叠，从而提高了有效的传输带宽。

Prefetch-Layer

在多个 GPU 上并行化层获取

ZeRO-Inference 利用 GPU 和 CPU 内存之间的四个 PCIe 互连来并行化层获取，以便在多个 GPU 上进行更快的推理计算。在表 4中，我们报告了与单个 GPU3 相比，在两个和四个 GPU 上进行 token 生成的吞吐量改进。这些结果是在启用层预取的情况下收集的。报告的吞吐量数字是每个 GPU 的，表明当聚合的 PCIe 链接减少层获取延迟时，每个 GPU 上的 token 生成速度会变快。每个 GPU 吞吐量的提高转化为超线性扩展性能。此外，这些结果表明未来 PCIe 世代的带宽提高将有助于提高 ZeRO-Inference 的性能。

Multi-GPU

生成输出长度的影响

我们测量了输出 token 数量对性能的影响，因为 KV 缓存优化的内存开销随着输出 token 的增加而增加，并且可能限制批次大小。首先，我们考虑了 token 长度 10、20、50 和 100 对可以容纳一个 V100-32GB GPU 的批次大小的影响。表 5中的结果表明，对于 token 数量的 5 倍增加（与 10 的基线数量相比），批次大小减少了 2 倍。

Token-count-batch-size

接下来，我们测量了对使用四个 V100-32GB GPU 的生成吞吐量的影响。表 6 展示了 CPU 卸载的结果，表 7 展示了 NVMe 卸载的结果。我们观察到跨模型和卸载内存的一致影响，即增加输出 token 数量会按比例减少吞吐量，以匹配批次大小的减少。这些结果还表明了大型批次大小对 ZeRO-Inference 性能的重要性。

Token-count-cpu-throughput

Token-count-nvme-throughput

使用 ZeRO-Inference

我们简要讨论了用户如何确定 ZeRO-Inference 是否适合他们的应用程序，以及如何在 DeepSpeed 中启用 ZeRO-Inference。

何时使用 ZeRO-Inference

ZeRO-Inference 旨在用于需要 GPU 加速但缺乏足够的 GPU 内存来容纳模型的推理应用程序。此外，ZeRO-Inference 针对以吞吐量为导向且允许大型批次大小的推理应用程序进行了优化。其他技术，如Accelerate、DeepSpeed-Inference和DeepSpeed-MII，它们将整个模型放入 GPU 内存中（可能使用多个 GPU），更适合于对延迟敏感或批次大小很小的推理应用程序。

如何使用 ZeRO-Inference

ZeRO-Inference 在 DeepSpeed 库版本 >= 0.6.6 中可用。将 ZeRO-Inference 集成到 token 生成管道（如Hugging Face generate）中，需要更新 DeepSpeed 配置，将ZeRO 优化设置为阶段 3，并将参数卸载设置为 CPU 或 NVMe。

以下是使用 CPU 内存卸载启用 ZeRO-Inference 的配置片段。

    "zero_optimization": {
        "stage": 3,
        "offload_param": {
            "device": "cpu",
            ...
        },
        ...
    }

以下是卸载到“/local_nvme”上挂载的 NVMe 设备的配置片段。

    "zero_optimization": {
        "stage": 3,
        "offload_param": {
            "device": "nvme",
            "nvme_path": "/local_nvme",
            ...
        },
        ...
    }

结论

人工智能技术的最新进展主要来自模型规模的极度扩展。然而，极端的模型扩展也使得训练和推断的硬件成本对于除最大组织以外的所有组织来说都难以承受，严重限制了对人工智能创新的访问。为了帮助民主化人工智能，我们开发了 ZeRO-Inference，这是一项技术，能够在单个 GPU 上执行大型模型的推理计算。ZeRO-Inference 通过将模型托管在 CPU 或 NVMe 内存中并将模型层流式传输到 GPU 内存中进行推理计算，降低了 SOTA 模型推理的 GPU 成本。ZeRO-Inference 补充了大型组织公开发布预训练 SOTA 模型的民主化工作，确保大多数用户（例如学生、爱好者、模型科学家等）能够负担得起这些模型的推理计算。

致谢

DeepSpeed 团队感谢 Stas Bekman 预览了这篇博文并提供了宝贵的反馈。

  1. 160 亿参数模型无法在 V100-32GB 上进行半精度推理，因为没有剩余的内存用于输入和中间结果。 ↩

  2. 在 GPU 内存中固定更多参数会导致小批次大小的内存不足错误。 ↩

  3. 对于多 GPU 运行，我们选择具有独立 PCIe 互连到 CPU 内存的 GPU。 ↩

更新：