优化故事：Bloom推理

Optimization Story Bloom Reasoning

本文介绍了我们如何制作一个高效的推断服务器，为https://huggingface.co/bigscience/bloom提供动力的幕后情况。

在几个星期内，我们实现了5倍的延迟降低（和50倍的吞吐量增加）。我们想要分享我们为实现这种速度改进而经历的所有挣扎和史诗般的胜利。

许多不同的人在许多阶段都参与其中，因此这里不会涵盖所有内容。请谅解，由于我们仍在学习如何优化非常大的模型以及许多新的硬件特性和内容经常出现，因此某些内容可能已过时或完全错误。

如果您喜欢的优化方法未被讨论或不正确地表示，我们很抱歉，请与我们分享，我们非常乐意尝试新的东西并纠正我们的错误。

• MTEB 大规模文本嵌入基准测试
• 从PyTorch的DDP到Accelerate再到Trainer，轻松掌握分布式训练
• 使用同态加密实现对加密数据的情感分析

毫无疑问，如果首先无法访问大型模型，就没有真正的理由为其优化推断。这是许多不同人共同努力的结果。

为了在训练过程中最大化GPU的利用率，我们尝试了几种解决方案，最终选择了Megatron-Deepspeed来训练最终模型。这意味着代码本身不一定与transformers库兼容。

由于原始训练代码的原因，我们开始做我们经常做的事情之一：将现有模型转移到transformers。目标是从训练代码中提取相关部分，并在transformers中实现。Younes完成了这项工作。这绝不是一项小任务，因为这花费了近一个月的时间和200次提交才完成。

有几点需要注意，稍后会再次提到：

我们需要有更小的模型bigscience/bigscience-small-testing和bigscience/bloom-560m。这非常重要，因为它们体积更小，所以在处理它们时速度更快。

首先，您必须放弃所有希望，以便最终的logits完全相同。PyTorch的版本可能会更改内核并引入微妙的差异，不同的硬件可能会产生不同的结果，因为不同的架构（出于成本考虑，您可能不想一直在A100 GPU上开发）。

对于所有模型来说，获得一个良好的严格测试套件非常重要

我们发现的最好的测试方法是使用一组固定的提示。您知道提示，您知道需要确定性的完成。如果两个生成的结果相同，您基本上可以忽略较小的logits差异。每当您看到漂移时，您需要进行调查。可能是您的代码没有做到应该做的事情，或者您实际上超出了该模型的领域，因此该模型对噪声更敏感。如果有几个提示并且提示足够长，您不太可能意外触发所有提示。提示越多越好，越长越好。

第一个模型（small-testing）与大型bloom一样都是<bfloat16，所以一切都应该非常相似，但它没有经过很多训练或者表现不好，所以输出非常波动。这意味着我们在生成测试时遇到了问题。第二个模型更加稳定，但是在保存时使用的是float16而不是bfloat16。这两者之间存在更大的误差空间。

为了完全公平起见，bfloat16-> float16转换在推理模式下似乎是可行的（bfloat16主要用于处理较大的梯度，在推理中不存在）。

在该步骤中，我们发现并实施了一个重要的权衡。因为bloom是在分布式环境中进行训练的，部分代码在一个线性层上进行了Tensor并行操作，这意味着在单个GPU上以单个操作方式执行相同操作会得到不同的结果。这花了一段时间才能找到问题所在，要么我们选择100%的一致性，模型速度要慢得多，要么我们会在生成结果上有一点差异，但运行速度要快得多，代码也更简单。我们选择了可配置的标志。

注意：在此上下文中，流水线并行（Pipeline Parallelism，PP）意味着每个GPU将拥有
一些层，因此每个GPU在将数据传递给下一个GPU之前将在给定的数据块上工作。

现在我们有一个可用的干净版本的transformers，可以开始运行它。

Bloom是一个352GB（176B参数以bf16表示）的模型，我们至少需要那么多的GPU内存才能使其适应。我们曾经在较小的机器上尝试将其转移到CPU，但推理速度慢了几个数量级，所以我们放弃了这个想法。

然后我们想基本上使用pipeline。所以这是dogfooding，这是API在幕后一直使用的东西。

然而，管道没有分布感知（这不是它们的目标）。经过简要讨论，我们最终使用了加速新创建的device_map="auto"来管理模型的分片。我们必须解决一些错误，并稍微修改transformers代码以帮助accelerate正确工作。

它的工作原理是将transformers的各个层进行分割，并将模型的一部分分配给每个GPU。因此，GPU0开始工作，然后将其交给GPU1，依此类推。

最后，在一个小型HTTP服务器的基础上，我们可以开始提供bloom（大模型）！

但是我们甚至还没有开始讨论优化！

我们实际上有很多优化，整个过程就像是一座纸牌堡。在优化过程中，我们将对底层代码进行修改，确保不以任何方式破坏模型非常重要，而且比你想象的更容易。

所以我们现在处于优化的第一步，我们需要开始测量并持续测量性能。所以我们需要考虑我们关心的内容。对于一个支持多种选项的开放式推理服务器，我们希望用户发送许多带有不同参数的查询，我们关心的是：

我们可以同时为多少用户提供服务（吞吐量）一个平均用户被服务需要多长时间（延迟）？

我们在locust中制作了一个测试脚本，就是这样：

from locust import HttpUser, between, task
from random import randrange, random


class QuickstartUser(HttpUser):
    wait_time = between(1, 5)

    @task
    def bloom_small(self):
        sentence = "Translate to chinese. EN: I like soup. CN: "
        self.client.post(
            "/generate",
            json={
                "inputs": sentence[: randrange(1, len(sentence))],
                "parameters": {"max_new_tokens": 20, "seed": random()},
            },
        )

    @task
    def bloom_small(self):
        sentence = "Translate to chinese. EN: I like soup. CN: "
        self.client.post(
            "/generate",
            json={
                "inputs": sentence[: randrange(1, len(sentence))],
                "parameters": {
                    "max_new_tokens": 20,
                    "do_sample": True,
                    "top_p": 0.9,
                    "seed": random(),
                },
            },
        )

**注意：这不是我们使用的最好也不是唯一的负载测试，但它总是第一个运行的，以便可以公平地比较各种方法。在这个基准测试中表现最好并不意味着它是最好的解决方案。除了实际的真实性能之外，还必须使用其他更复杂的场景。

我们希望观察各种实现的启动情况，并确保在负载下服务器适当地断开连接。断开连接意味着服务器可以（快速）回答，它不会回答您的查询，因为有太多人同时尝试使用它。避免过度负载非常重要。

在这个基准测试中，初始性能为（在16xA100 40Go on GCP上，这是整个机器使用的机器）：

每秒请求数：0.3（吞吐量）延迟：350ms / token（延迟）

这些数字并不算太好。在开始工作之前，让我们估计我们可以想象实现的最佳结果。操作数量的公式是24Bsh^2 + 4𝐵s^2h24Bsh^2 + 4𝐵s^2h，其中B是批量大小，s是序列长度，h是隐藏维度。

让我们进行计算，我们得到单次前向传递的17 TFlop。看看A100的规格，它声称单个卡可以达到312 TFLOPS。这意味着单个GPU的潜在运行速度为17 / 312 = 54ms/token。我们使用了其中的16个，因此在整个机器上是3ms/token。请对这些数字保持一定的保留，永远无法达到这些数字，而且实际性能很少能与规格相匹配。此外，如果计算不是限制因素，那么您可以得到更低的数字。只是好的做法是了解您离目标有多远。在这种情况下，我们差了两个数量级，所以相当远。此外，此估计将所有浮点运算用于延迟服务，这意味着一次只能处理一个请求（这没问题，因为您正在最大化机器的使用，所以没有其他很多事情可做，但通过批处理可以更容易地获得吞吐量回归）。

注意：在此上下文中，张量并行（TP）意味着每个GPU将拥有部分权重，因此所有GPU始终处于活动状态且工作量较少。
通常，这会带来一些轻微的开销，即一些工作被重复执行，更重要的是GPU定期需要相互通信以继续计算。

既然我们对我们的立场有了很好的理解，现在是时候开始工作了。

我们尝试了许多不同的方法，基于人员和我们的各种知识。

所有的努力都值得一篇博客文章，所以我只会列出它们，解释一些最终的经验教训，并深入了解当前服务器中的详细信息。从管道并行（PP）到张量并行（TP）的转变对延迟来说是一个重要而有趣的变化。每个GPU将拥有一部分参数，并且所有GPU将同时工作。因此，延迟应该大大减少，但要付出的代价是通信开销，因为它们需要定期相互通信关于它们的结果。

值得注意的是，这是一种非常广泛的方法，并且故意的目的是更多地了解每个工具以及它如何适应以后的努力。

将代码迁移到在TPU上运行的JAX/Flax

• 预计选择并行类型更容易。因此，测试TP应该更容易。这是Jax设计的好处之一。
• 硬件上更受限制，TPU上的性能可能优于GPU，并且对于TPU的供应商选择较少。
• 缺点是需要进行另一次移植。但是，无论如何，这在我们的库中都是受欢迎的。

结果：

• 迁移并不是一项容易的任务，因为一些条件和内核很难正确重现。不过还是可以应付。
• 一旦迁移完成，并行性就很容易获得。感谢Jax，这个声明是正确的。
• Ray/与TPU工作节点进行通信对我们来说确实很痛苦。我们不知道是工具的问题，网络的问题还是简单地缺乏知识，但它比我们预期的更大程度地减慢了实验和工作的进度。我们启动了一个需要运行5分钟的实验，等了5分钟什么都没发生，10分钟后仍然没有任何变化，结果发现有些工作节点已经停止/没有响应，我们不得不手动进入，弄清楚发生了什么，修复问题，重新启动某些东西，然后重新启动，我们刚刚浪费了半个小时。重复这样的过程足够多次，失去的时间很快就会累积起来。让我们强调一下，这不一定是我们使用的工具的批评，而是我们的主观经验。
• 无法控制编译。一旦我们使事情运行起来，我们尝试了几个设置来找出最适合我们心目中推理的设置，结果发现从设置中很难猜测延迟/吞吐量会发生什么。例如，我们在batch_size=1（因此每个请求/用户都是独立的）时的rps为0.3，延迟为15ms/token（不要与本文中的其他数字进行太多比较，因为它是在一个具有非常不同配置文件的不同机器上），这非常好，但是总吞吐量并不比我们之前的旧代码好多少。因此，我们决定添加批处理，使用BS=2时，延迟增加了5倍，但吞吐量只增加了2倍…进一步调查后，发现在batch_size=16之前，每个batch_size具有相同的延迟特性。因此，我们可以在5倍延迟成本下获得16倍的吞吐量。不错，但从我们的目标数字来看，我们真的希望有更精细的控制。我们瞄准的数字来自于100ms、1s、10s、1mn规则。

使用ONNX/TRT或其他编译方法

• 它们应该处理大部分优化工作
• 缺点是通常需要手动处理并行性。

结果：

• 结果发现，为了能够跟踪/即时编译/导出内容，我们需要重新设计部分PyTorch，以便它可以轻松地与纯PyTorch方法融合。总的来说，我们发现通过保持在PyTorch世界中，我们可以获得大部分所需的优化，这使我们能够在不需要太多编码工作的情况下保持灵活性。还需要注意的是，由于我们在GPU上运行并且文本生成有许多前向传递，我们需要张量保留在GPU上，有时很难将张量发送到某个库中，然后再得到结果，执行逻辑运算（如argmax或采样）并再次将其反馈。将循环放在外部库中意味着失去了类似Jax的灵活性，因此在我们的用例中没有考虑这种情况。

DeepSpeed

• 这是用于训练的技术，使用它进行推断是合理的。
• 缺点是以前从未使用过/准备过用于推断。

结果：

• 我们取得了非常令人印象深刻的快速结果，与我们目前运行的最后一次迭代大致相同。
• 我们不得不发明一种方法，在DeepSpeed之上放置一个Web服务器（处理并发），而DeepSpeed本身也有多个进程（每个GPU一个）。由于存在一个优秀的库Mii，它不符合我们当初设想的极度灵活的目标，但我们现在可能会开始在其上工作（当前解决方案稍后讨论）。
• 我们在使用DeepSpeed时遇到的最大问题是缺乏稳定性。当我们在CUDA 11.4上运行它时，会出现代码构建为11.6的问题。而长期存在的问题是我们无法真正解决的，即会出现定期的内核崩溃（Cuda非法访问，尺寸不匹配等）。我们解决了其中一些问题，但在我们的Web服务器压力下，我们无法完全实现稳定性。尽管如此，我想向帮助我们的Microsoft团队大喊一声，我们进行了非常好的交谈，提高了我们对发生情况的理解，并为我们进行了一些后续工作提供了真正的见解。
• 我感到痛苦的一点是，我们的团队大部分在欧洲，而Microsoft位于加利福尼亚，因此合作在时间上很棘手，因此我们可能因此失去了很大一部分时间。这与技术无关，但承认合作的组织部分也非常重要。
• 另一个需要注意的是，DeepSpeed依赖于transformers来进行优化，由于我们不断更新代码，这使得DeepSpeed团队很难在我们的main分支上保持正常运作。我们对此感到抱歉，我想这就是为什么它被称为前沿技术。

Web服务器的想法

• 考虑到我们将运行一个免费服务器，用户将发送长文本、短文本、希望获取几个标记或整个配方，每个请求都具有不同的参数，这里必须做一些处理。

结果：

• 我们使用优秀的绑定tch-rs，将所有东西都转换为Rust。Rust并不旨在获得性能提升，而只是为了更细粒度地控制并行性（线程/进程），并在Web服务器并发性和PyTorch并发性上进行更细粒度的操作。由于GIL的存在，Python在处理低级细节方面非常困难。
• 事实证明，大部分痛苦来自端口，之后实验变得轻松。我们发现，通过对循环有足够的控制，即使在各种具有不同特性的请求的背景下，我们也可以为每个人提供出色的性能。这里有代码供好奇的人查看，但它不提供任何支持或良好的文档。
• 由于它在并行性方面更宽容，它在生产环境中使用了几周，我们可以更有效地使用GPU（使用GPU0处理请求1，而GPU1处理请求0）。我们的吞吐量从0.3 RPS增加到约2.5 RPS，延迟保持相同。最理想的情况是将吞吐量提高16倍，但这里显示的数字是真实的工作负载测量结果，所以这也不算太差。

纯PyTorch

• 通过删除reshape等操作，使用更好优化的内核，纯粹修改现有代码使其更快。
• 缺点是我们必须自己编写TP的代码，并且我们的代码仍然需要适应我们的库（主要是）。

结果：

• 下一章。

编写更高效的PyTorch

清单上的第一项是在最初的实现中删除不必要的操作。有些可以通过查看代码并找出明显的缺陷来看到：

• Alibi在Bloom中用于添加位置嵌入，并且在太多地方进行了计算，我们只需要计算一次并更高效地计算。

旧代码：链接 新代码：链接

这是一个10倍的加速，最新版本还包括填充！由于这一步只计算一次，实际速度并不重要，但是总体上减少操作和张量创建的数量是一个好的方向。

当您开始进行性能分析时，其他部分会更清晰，我们广泛使用了tensorboard扩展

这提供了一种能够提供洞察力的图像:

注意，这需要很长时间，请注意这是CPU视图，因此长条并不意味着时间长，而是意味着CPU正在等待前一步骤的GPU结果。 我们在`baddbmm`之前看到了许多`cat`操作。

例如，删除了很多reshape/transpose，我们发现： – 注意力是热点路径（这是预期的，但验证总是好的）。 – 在注意力中，由于大量的reshape，许多内核实际上是副本 – 通过重新设计权重和过去，我们可以消除reshape。这是一个破坏性的改变，但它确实大大提高了性能！

支持TP

好的，我们已经消除了大部分低悬挂的水果，现在我们从PP的每个标记延迟大致减少了15%，从350ms/token降到300ms/token。 这是延迟的15%减少，但实际上提供的比这更多，但我们最初在测量方面不是非常严格，所以让我们坚持这个数字。

然后我们开始提供TP实现。结果比我们预期的要快得多，实现只用了半天时间（由一名有经验的开发人员完成）。结果在这里。我们还能够重用其他项目的代码，这有助于提高效率。

延迟直接从300ms/token降至91ms/token，这在用户体验方面是巨大的改进。一个简单的20个标记的请求从6秒变为2秒，从“慢”体验变为稍有延迟。

此外，吞吐量大大提高到10RPS。吞吐量来源于批量大小为1的查询所需的时间与批量大小为32的查询相同，此时吞吐量在延迟成本上基本上是免费的。

低悬挂的水果

现在我们有了TP实现，我们可以再次进行性能分析和优化。这是一个足够重大的变化，我们不得不从头开始。

首先显眼的是同步（ncclAllReduce）开始成为负载的主要部分，这是预期的，这是同步部分，它确实需要一些时间。我们从未尝试过查看和优化它，因为它已经使用了nccl，但在那里可能仍然有一些改进的空间。我们认为要做得更好可能会很困难。

其次是Gelu运算符，它启动了许多逐元素的内核，总体上它占据了比我们预期的更大的计算份额。

我们从以下方式进行了更改：

def bloom_gelu_forward(x):
    return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))

到

@torch.jit.script
def bloom_gelu_forward(x):
    return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))

这将操作从多个小的逐元素内核（因此张量副本）转换为单个内核操作！

这将延迟从91ms/token降低了10%，达到了81ms/token！

但要小心，这不是一个可以随处使用的神奇黑匣子，内核融合不一定会发生，或者先前使用的操作已经非常高效。

在我们发现它效果良好的地方：

• 您有很多小的/逐元素操作
• 您有一些难以删除的热点，例如reshape、复制等
• 融合发生时。

史诗级失败

在我们的测试期间，我们还遇到了一些情况，即Rust服务器的延迟与Python服务器相比始终较低25%。这相当奇怪，但由于它经过了一致的测量，并且因为删除内核提供了加速，所以我们有一种印象，也许去除Python的开销可以提供很好的提升。

我们开始了一个为期3天的工作，重新实现了torch.distributed的必要部分，以在Rust世界中运行。我们之前已经实现了这个版本，但是与Python版本相比，有些东西不对劲。在调查问题的过程中，我们发现… 我们忘记了在PyTorch测量中移除分析器 …

这是一个史诗级的失败，因为移除它后我们恢复了25%，然后两个版本都以同样的速度运行。这正是我们最初的预期，因为Python主要运行torch cpp的代码，所以它不应该降低性能。最终，3天并不是世界末日，它可能在将来有用，但还是很糟糕。这在优化过程中很常见，错误的或者不准确的测量结果最终会令人失望，甚至对整体产品产生不利影响。这就是为什么要分步骤进行，并尽快对结果有期望，以帮助控制风险。

我们还需要格外谨慎的地方是初始前向传递（没有过去的信息）和后面的前向传递（带有过去的信息）。如果优化了第一个前向传递，那么后面的前向传递肯定会变慢，因为后面的前向传递更重要，占用了大部分运行时间。另一个常见的问题是测量CPU时间而不是实际的CUDA时间，所以在运行时需要使用torch.cuda.synchronize()确保内核完成。

自定义内核

到目前为止，我们在PyTorch之外几乎没有使用任何自定义代码就实现了接近DeepSpeed的性能！非常棒。而且我们也不需要在运行时批处理大小的灵活性上做出任何妥协！

但是鉴于DeepSpeed的经验，我们想尝试编写一个自定义内核，在热点路径中融合几个操作，而torch.jit.script无法为我们做到这一点。基本上是以下两行：

attn_weights = attention_scores.masked_fill_(attention_mask, torch.finfo(attention_scores.dtype).min)
attention_probs = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(input_dtype)

第一个masked fill创建了一个新的张量，这个张量只是告诉softmax运算符忽略这些值。此外，softmax需要在float32上计算（为了稳定性），但在自定义内核中，我们可以限制所需的类型转换数量，只限制在实际求和和累加上。

代码可以在这里找到。请记住，我们有一个单独的GPU架构来进行目标定位，所以我们可以专注于这一点，而且我们并不是（现在）在编写内核方面的专家，所以可能有更好的方法来实现。

这个自定义内核使延迟降低了另外10%，从81ms/token降低到71ms/token。同时保持了我们的灵活性。

之后，我们还调查和探索了其他事项，比如融合更多的操作、去除其他的reshape或将它们放在其他地方。但是没有任何尝试能够产生足够大的影响，达到最终版本。

Web服务器部分

与Rust对应物一样，我们需要实现具有不同参数的请求批处理。由于我们处于PyTorch世界中，我们对正在进行的工作有很大的控制权。因为我们在Python中，我们面临的限制因素是torch.distributed需要在多个进程而不是线程上运行，这意味着在进程之间进行通信稍微困难一些。最后，我们选择使用Redis pub/sub在所有进程之间传递原始字符串来分发请求。由于我们处于不同的进程中，与其传递张量（大小更大），这种方式更容易。

然后，我们不得不放弃使用generate，因为这会将参数应用于批处理的所有成员，而我们实际上想要应用不同的参数集。幸运的是，我们可以重复使用低级别的项目，如LogitsProcessor，以节省大量工作。

因此，我们重构了一个generate函数，它接受参数列表，并将其应用于批处理的每个成员。

最终用户体验中的另一个非常重要的方面是延迟。由于不同的请求具有不同的参数集，我们可能有一个请求包含20个令牌，另一个请求包含250个令牌。由于每个令牌的延迟为75ms，一个请求需要1.5秒，而另一个请求需要18秒。如果我们一直批处理，那么请求者将被迫等待18秒，而看起来我们的运行速度只有900ms/token，这相当慢！

由于我们处于一个极具灵活性的PyTorch世界中，我们可以做的是在生成前20个标记后，立即从批处理中提取第一个请求，并在请求的1.5秒内返回给用户！我们还刚好节省了230个标记的计算。

因此，灵活性对于获得最佳延迟是很重要的。

优化是一项永无止境的工作，就像任何其他项目一样，通常情况下，20%的工作通常能产生80%的结果。在某个时候，我们开始有一个小的测试策略来确定我们提出的某个想法的潜在收益，如果测试没有产生显著的结果，那么我们就放弃了这个想法。花1天时间增加10%是足够有价值的，花2周时间增加10倍也是足够有价值的。花2周时间增加10%并不是那么有趣。

你试过…吗？

以下是我们知道存在但由于各种原因而未使用的工具。可能是因为感觉它不适合我们的用例，工作量太大，收益不够可期，甚至仅仅是因为我们有太多的选择而放弃了一些，只是因为时间不够。以下没有特定的顺序：

• Cuda图
• nvFuser（这是torch.jit.script的动力源，所以我们确实使用过它。）
• FasterTransformer
• Nvidia的Triton
• XLA（Jax也在使用xla！）
• torch.fx
• TensorRT

如果您喜欢的工具不在这里，请随时告诉我们，或者如果您认为我们错过了一些重要的东西，那可能会有用！

闪光关注

我们曾经简要研究过集成闪光关注，虽然在第一次前向传递（没有past_key_values）时表现非常好，但在使用past_key_values时没有带来很大的改进。由于我们需要将alibi张量包含在计算中，我们决定不进行这项工作（至少目前不进行）。

OpenAI Triton

Triton是一个在Python中构建自定义内核的强大框架。我们希望能够更多地使用它，但迄今为止我们还没有。我们非常希望能够看到它是否比我们的Cuda内核表现更好。当我们考虑该部分的选择时，直接使用Cuda编写似乎是最直接的路径。

填充和重塑

正如本文中所提到的，每个张量副本都有一定的成本，而运行生产过程的隐藏成本则是填充。当两个查询的长度差异很大时，您必须进行填充（使用一个虚拟标记）使它们适应一个方形。这可能会导致很多不必要的计算。详细信息。

理想情况下，我们将能够根本不进行这些计算，并且永远不会进行重塑。Tensorflow有RaggedTensor和Pytorch Nested tensors的概念。这两者似乎不像常规张量那样简化，但可能使我们能够进行较少的计算，这总是有利的。

在理想的世界中，整个推断过程将是用CUDA或纯GPU实现的。考虑到我们能够合并操作时产生的性能改进，这看起来是令人向往的。但是这会带来多大的效果，我们不知道。如果更聪明的GPU人有想法，我们会倾听！

所有这些工作都是许多HF团队成员的合作成果。没有特定的顺序，@ThomasWang @stas @Nouamane @Suraj @Sanchit @Patrick @Younes @Sylvain @Jeff（Microsoft）@Reza和所有BigScience组织。