关于人工智能:优化故事-BLOOM-模型推理

通过“九九八十一难”，大模型终于炼成。下一步就是架设服务，筹备开门营业了。真这么简略？恐怕未必！行百里者半九十，推理优化又是新的雄关漫道。如何进行提早优化？如何进行老本优化 (别忘了 OpenAI 8K 上下文的 GPT-4 模型，提醒每 1000 词元只需 0.03 美金，补全每 1000 词元只需 0.06 美金)？如何在提早和吞吐量之间折衷？如何解决大模型特有的分布式推理后端和网络服务前端的合作问题……要不入手之前还是先看看 BLOOM 推理服务踩过的坑吧！

本文介绍了咱们在实现 BLOOM 模型高效推理服务的过程中产生的幕后故事。

在短短数周内，咱们把推理提早升高了 5 倍 (同时，吞吐量减少了 50 倍)。咱们将分享咱们为达成这一性能改良而经验的所有奋斗和史诗般的胜利。

在此过程中，不同的人参加了不同的阶段，尝试了各种不同的优化伎俩，咱们无奈一一列举，还请多多包涵。如果你发现本文中某些内容可能已过期甚至齐全谬误，这也不奇怪，因为一方面对于如何优化超大模型性能咱们仍在努力学习中，另一方面，市面上新硬件性能和新优化技巧也层出不穷。

咱们很道歉如果本文没有探讨你最中意的优化技巧，或者咱们对某些办法表述有误。但请通知咱们，咱们十分乐意尝试新货色并纠正错误。

训练 BLOOM

这是显而易见的，如果不先获取到大模型，那推理优化就无从谈起。大模型训练是一项由很多不同的人独特领导的超级工程。

为了最大化 GPU 的利用率，咱们摸索了多种训练计划。最初，咱们抉择了 Megatron-Deepspeed 来训练最终模型。这意味着训练代码与 transformers 库并不齐全兼容。

移植至 transformers

因为上文提及的起因，咱们第一件事是将现有模型移植到 transformers 上。咱们须要从训练代码中提取相干代码并将其实现至 transformers 里。Younes 负责实现了这项工作。这个工作量相对不小，咱们大略花了将近一个月的工夫，进行了 200 次提交 才最终实现。

有几点须要留神，咱们前面还会提到:

小版的模型，如 bigscience/bigscience-small-testing 和 bigscience/bloom-560m 十分重要。因为模型构造与大版的一样但尺寸更小，所以在它们下面所有工作 (如调试、测试等) 都更快。

首先，你必须放弃那种最终你会失去比特级统一的 logits 后果的空想。不同的 PyTorch 版本间的算子核函数更改都会引入细微差别，更不用说不同的硬件可能会因为体系架构不同而产生不同的后果 (而出于老本起因，你可能并不能始终在 A100 GPU 上开发)。

一个好的严格的测试套件对所有模型都十分重要

咱们发现，最佳的测试形式是应用一组固定的提醒。从测试角度，你晓得提醒 (prompt)，而且你想要为每个提醒生成确定性的补全 (completion)，所以解码器用贪婪搜寻就好了。如果两次测试生成的补全是雷同的，你基本上能够忽视 logits 上的小差别。每当你看到生成的补全产生漂移时，就须要考察起因。可能是你的代码没有做它应该做的事; 也有可能是你的提醒不在该模型的常识域内 [译者注: 即模型的训练数据中并不蕴含提醒所波及的话题]，所以它对噪声更敏感。如果你有多个提醒且提醒足够长，不太可能每个提醒都触发上述不在常识域的问题。因而，提醒越多越好，越长越好。

第一个模型 (small-testing) 和大 BLOOM 一样，精度是 bfloat16 的。咱们原以为两者应该十分类似，但因为小模型没有通过太多训练或者单纯只是性能差，最终体现进去的后果是它的输入稳定很大。这意味着咱们用它进行生成测试会有问题。第二个模型更稳固，但模型数据精度是 float16 而不是 bfloat16，因而两者间的误差空间更大。

偏心地说，推理时将 bfloat16 模型转换为 float16 仿佛问题不大 (bfloat16 的存在次要是为了解决大梯度，而推理中不存在大梯度)。

在此步骤中，咱们发现并实现了一个重要的折衷。因为 BLOOM 是在分布式环境中训练的，所以局部代码会对 Linear 层作张量并行，这意味着在单 GPU 上运行雷同的操作会失去 不同的数值后果。咱们花了一段时间才查明这个问题。这个问题没方法彻底解决，要么咱们谋求 100% 的数值一致性而就义模型运行速度，要么咱们承受每次生成时都会呈现一些小的差别但运行速度更快，代码更简略。咱们为此设了一个标记位供用户本人配置。

首次推理 (PP + Accelerate)

留神: 这里，流水线并行 (Pipeline Parallelism, PP) 意味着每个 GPU 将分得模型的一些层，因而每个 GPU 将实现一部分操作，而后再将其后果交给下一个 GPU。

当初咱们有了一个能反对 BLOOM 的 transformers，咱们能够开始跑了。

BLOOM 是一个 352GB (176B bf16 参数) 的模型，咱们至多须要那么多显存能力放下它。咱们花了一点工夫试了试在小显存的 GPU 上应用 CPU 卸载的形式来推理，然而推理速度慢了几个数量级，所以咱们很快放弃了它。

而后，咱们转而想应用 transformers 的 pipeline API，吃一下这个 API 的狗粮。然而，pipeline 不是分布式感知的 (这不是它的设计指标)。

通过短暂的技术计划探讨，咱们最终应用了 accelerate 的新性能 device_map="auto 来治理模型的分片。咱们不得不解决一些 accelerate 以及 transformers 的 bug，才使得这一计划能失常工作。

它的工作原理是将 transformer 模型按层进行切分，每个 GPU 分到一些层。真正运行时，是 GPU0 先开始工作，而后将后果交给 GPU1，顺次上来。

最初，在前端架一个小型 HTTP 服务器，咱们就能够开始提供 BLOOM (大模型) 推理服务了！！

终点

至此，咱们甚至还没有开始探讨优化！

咱们其实做了不少优化，这所有过程有点像纸牌叠城堡游戏。在优化期间，咱们将对底层代码进行批改，所以肯定要确保咱们不会以任何形式毁坏模型，这一点十分重要，而且其实比设想中更容易做到。

优化的第一步是测量性能。在整个优化过程中，性能测量贯通始终。所以，首先须要思考咱们须要测量什么，也即咱们关怀的是什么。对于一个反对多种选项的开放式推理服务而言，用户会向该服务发送各种不同的查问申请，咱们关怀的是:

1. 咱们能够同时服务的用户数是多少 (吞吐量)？
2. 咱们均匀为每个用户服务的工夫是多少 (提早)？

咱们用 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(),},
            },
        )

留神: 这不是咱们最佳的也不是惟一的负载测试，但始终是咱们第一个运行的负载测试，因而它可用于偏心地比拟不同计划。在此基准测试体现最好并不意味着它相对是最好的解决方案。咱们还须要应用其余更简单的测试场景来模仿实在场景的真实性能。

咱们想察看各种实现计划部署时如何爬坡，并确保在熔断时适当地升高服务器负载。熔断意味着本来能 (疾速) 响应你的申请的服务不再响应你的申请，因为同一时间有太多人想要应用它。防止 死亡之拥 (hug of death) 是极其重要的。[译者注: 死亡之拥是一个互联网畛域的隐喻，意指因为极其峰值流量而导致互联网服务宕机]

在上述基准测试中，咱们失去的初始性能是 (应用 GCP 上的 16xA100 40G 环境测得，本文后续所有测试都基于该环境):

每秒解决申请数 (吞吐量): 0.3
每词元提早: 350ms

这两个值并不是很好。在正式开始工作之前，咱们能够预估一下咱们能失去的最好后果。BLOOM 模型所需的计算量公式为 $24Bsh^2 + 4Bs^2h * 24Bsh^2 + 4Bs^2h$，其中 B 是 batch size，s 是序列长度，h 是隐含层维度。

让咱们算一下，一次前向流传须要 17 TFlop。A100 的 规格 为单卡 312 TFLOPS。这意味着单个 GPU 最多能达到 17 / 312 = 54 毫秒 / 词元 的提早。咱们用了 16 个 GPU，因而可得 3 毫秒 / 词元。这只是个下限，咱们永远不可能达到这个值，况且事实中卡的性能很少能达到其规格所声称的数字。此外，如果你的模型并不受限于计算 [译者注: 如受限于内存带宽、受限于 IO 带宽等]，那么这个值你也达不到。晓得现实值，只是为了让咱们对优化指标心里有个数。在这里，咱们到目前为止与现实值差 2 个数量级。此外，这个预计假如你将所有算力都用于提早型服务，这意味着一次只能执行一个申请 (没关系，因为你正在最大化你的机器利用率，所以没有太多其余事件要做; 但另一个思路是，咱们能够就义一点提早，通过批处理形式来取得更高的吞吐量)。

摸索多条路线

留神: 这里，张量并行 (Tensor Parallelism，TP) 意味着每个 GPU 将领有局部权重，因而所有 GPU 始终处于工作状态，专一于分给它的局部工作。通常这会带来十分轻微的开销，因为会有一些工作是反复的，更重要的是，GPU 必须定期互相通信交换它们的后果，而后再持续计算。

当初咱们曾经比较清楚地理解了咱们的处境，是时候开始工作了。

咱们依据咱们本人及其他人的各种教训和常识尝试了各种办法。

每次尝试都值得写一篇专门的博文，因为篇幅所限，在这里咱们仅将它们列出来，并只深刻解释并钻研那些最终利用到以后服务中去的技术的细节。从流水线并行 (PP) 切换到张量并行 (TP) 是提早优化的一个重要一步。每个 GPU 将领有局部参数，并且所有 GPU 将同时工作，所以提早应该会迅速降落。然而付出的代价是通信开销，因为它们的两头后果须要常常相互通信。

须要留神的是，这里波及的办法相当宽泛。咱们会无意识地学习更多对于每个工具的常识，以及在后续优化中如何应用它。

将代码移植到 JAX/Flax 中以在 TPU 上运行

• 并行计划的抉择更加容易。因而 TP 的测试会更不便，这是 JAX 的设计带来的益处之一。
• 对硬件的限度更多，JAX 上 TPU 的性能可能比 GPU 更好，但 TPU 比 GPU 更难获取 (只在 GCP 上有，数量也没有 GPU 多)。
• 毛病: 须要移植工作。但无论如何，把它集成到咱们的库外面这件事必定是受欢迎的。

后果:

• 移植比拟麻烦，因为某些条件语句和核函数很难精确复制，但尚可勉力为之。
• 一旦移植完后，测试各种并行计划就比拟不便。感激 JAX，没有食言。

• 事实证明，在 Ray 集群里与 TPU worker 通信对咱们来讲真的太苦楚了。

  不晓得是工具起因还是网络的起因，或者仅仅是因为咱们不太懂，但这事实上减慢了咱们的试验速度，而且须要的工作比咱们预期的要多得多。
  咱们启动一个须要 5 分钟工夫运行的试验，等了 5 分钟没有产生任何事件，10 分钟之后依然没有任何事件产生，后果发现是一些 TPU worker 宕机了或者是没有响应。咱们不得不手动登进去看，弄清楚产生了什么，修复它，重启一些货色，最初再重新启动试验，就这样半小时过来了。几次下来，几天就没了。咱们再强调一下，这未必真的是咱们应用的工具的问题，但咱们的主观体验的确如此。

• 无法控制编译

  咱们运行起来后，就尝试了几种设置，想找出最适宜咱们心目中想要的推理性能的设置，后果证实很难从这些试验中揣测出提早 / 吞吐量的法则。例如，在 batch_size=1 时吞吐量有 0.3 RPS (Requests Per Second, RPS) (此时每个申请 / 用户都是独立的)，提早为 15 毫秒 / 词元 (不要与本文中的其余数字进行太多比拟，TPU 机器与 GPU 机器大不相同)，提早很好，然而总吞吐量跟之前差不多。所以咱们决定引入批处理，在 batch_size=2 的状况下，提早减少到原来的 5 倍，而吞吐量只进步到原来的 2 倍…… 通过进一步考察，咱们发现始终到 batch_size=16，每个 batch_size 之间的提早都差不多。
  因而，咱们能够以 5 倍的提早为代价取得 16 倍的吞吐量。看上去挺不错的，但咱们更心愿对提早有更细粒度的管制，从而使得提早能满足 100ms, 1s, 10s, 1mn 规定中的各档。

应用 ONNX/TRT 或其余编译办法

• 它们应该能解决大部分优化工作
• 毛病: 通常须要手动解决并行性

后果:

• 事实证明，为了可能 trace/jit/export 模型，咱们须要重写 PyTorch 相干的一部分代码，使其可能很容易与纯 PyTorch 办法相交融。总体来讲，咱们发现咱们能够通过留在 PyTorch 中取得咱们想要的大部分优化，使咱们可能放弃灵活性而无需进行太多编码工作。另一件值得注意的事件是，因为咱们在 GPU 上运行，而文本生成有很多轮前向过程，所以咱们须要张量留在 GPU 上，有时很难将你的张量输给某个库，返回后果，计算 logits (如 argmax 或采样)，再回输给那个库。

  将循环放在内部库外面意味着像 JAX 一样失去灵活性，这不是咱们构想的推理服务利用场景的应用办法。

DeepSpeed

• 这是咱们训练 BLOOM 时应用的技术，所以用它来推理也很偏心
• 毛病: DeepSpeed 之前从未用于推理，其设计也没筹备用于推理

后果:

• 咱们很快就失去了很不错的后果，这个后果与咱们现行计划的上一版性能大致相同。
• 咱们必须想出一种办法，在多过程上架设用于解决并发申请网络服务，因为当初一个推理工作是由多个 DeepSpeed 过程实现的 (每个 GPU 一个过程)，。有一个优良的库 Mii 可供使用，它尽管还达不到咱们所构想的极致灵便的指标，但咱们当初能够在它之上开始咱们的工作。(以后的解决方案稍后探讨)。

• 咱们在应用 DeepSpeed 时遇到的最大问题是不足稳定性。

  咱们在 CUDA 11.4 上运行基于 11.6 编译的代码时遇到了问题。而其中一个由来已久的、咱们永远无奈真正解决的问题是: 常常会产生核函数解体 (CUDA 非法拜访、尺寸不匹配等)。咱们修复了其中一些问题，但在压测咱们的网络服务时，咱们永远无奈齐全实现稳定性。尽管如此，我想向帮忙过咱们的 Microsoft 人员说，感激那些十分欢快的交换，它们进步了咱们对正在产生的事件的了解，并为咱们的后续工作提供了远见卓识。

• 另一个痛点是咱们的团队次要在欧洲，而微软在加利福尼亚，所以单干工夫很辣手，咱们因而损失了大量工夫。这与技术局部无关，但咱们的确意识到单干的组织局部也十分重要。
• 另一件须要留神的事件是，DeepSpeed 依赖于 transformers 来注入其优化，并且因为咱们始终在更新咱们的代码，这使得 DeepSpeed 团队很难在咱们的主分支上工作。很道歉让它变得艰难，这也可能是 transformers 被称为技术最前沿的起因。

无关 Web 服务的想法

• 鉴于咱们筹备运行一个收费服务，反对用户向该服务发送长短不一的文本，并要求获取短至几个词，长至如整个食谱那么长的回应，每个申请的参数也能够各不相同，web 服务须要做点什么来反对这个需要。

后果:

• 咱们应用绑定库 tch-rs 在 Rust 中重写了所有代码。Rust 的指标不是进步性能，而是对并行性 (线程 / 过程) 以及 web 服务和 PyTorch 的并发性进行更细粒度的管制。因为 GIL 的存在，Python 很难解决这些底层细节。
• 结果表明，大部分的苦楚来自于移植工作，移植完后，试验就轻而易举了。咱们认为，通过对循环进行准确的管制，即便在具备大量不同属性的申请的场景中，咱们也能够为每个申请提供杰出的性能。如果你感兴趣的话，能够查看 代码，但这份代码没有任何反对，也没有好的文档。
• Rust web 服务投入生产了几周，因为它对并行性的反对更宽松，咱们能够更无效地应用 GPU (如应用 GPU0 解决申请 1，而 GPU1 解决申请 0)。在放弃提早不变的状况下，咱们把吞吐从 0.3 RPS 进步到了 ~2.5 RPS。尽管在最现实状况下，咱们能将吞吐进步到 16 倍。但理论工作负载上的测进去能到 8 倍左右的话也还算不错。

纯 PyTorch

• 纯正批改现有代码，通过删除诸如 reshape 之类的操作、应用更优化的核函数等办法来使其运行速度更快。
• 毛病: 咱们必须本人编写 TP 代码，并且咱们还有一个限度，即批改后代码最好依然适宜咱们的库 (至多大部分)。

后果

• 在下一章详述。

最终路线: PyTorch + TP + 1 个自定义内核 + torch.jit.script

编写更高效的 PyTorch

第一件事是在代码中删除不必要的操作。能够通过代码走查并找出显著可被删除的某些操作:

• Alibi 在 BLOOM 中用于增加地位嵌入 (position embeddings)，源代码中计算 Alibi 的中央太多，每次都从新计算一次，咱们优化成只计算一次，这样效率更高。

旧代码: 链接

新代码: 链接

这个改变取得了 10 倍的减速，最新版本还减少了对填充 (padding) 的反对！
因为此步骤仅计算一次，因而在这里，运算自身理论速度并不重要，而总体上缩小操作和张量创立的次数更重要。

当你开始 分析 代码性能时，其余局部会越来越清晰，咱们大量地应用了 tensorboard 来帮忙咱们进行性能分析。它提供了如下图所示的这类图像，能够提供无关性能的洞见:

<img src=”https://devrel.andfun.cn/devrel/posts/2023/04/17/SaHLyP.jpg”>

注意力层占用了很多工夫，留神这是一个 CPU 视图，所以条形很长并不意味着核函数执行工夫很长，它只意味着 CPU 正在期待上一步的 GPU 后果。

<img src=”https://devrel.andfun.cn/devrel/posts/2023/04/17/hmf54c.jpg”>

咱们还在 baddbmm 操作之前看到许多 cat 操作。

再举个例子，在删除大量 reshape / transpose 后，咱们在 tensorboard 中发现:

• 注意力是性能热点 (这是预期的，但可能通过测量数据来验证总是好的)。
• 在注意力中，因为大量的 reshape，很多核函数其实是显存拷贝函数。
• 咱们 能够 通过批改权重和 past_key_values 的内存布局来移除 reshape。这个改变有点大，但性能的确有肯定的进步！

反对 TP

好了，咱们曾经拿到了大部分唾手可得的成绩，当初咱们的 PP 版本的提早从大概 350 毫秒 / 词元升高到 300 毫秒 / 词元。提早升高了 15%，理论状况收益更大，但因为咱们最后的测量并不是十分严格，所以就用这个数吧。

而后咱们持续实现一个 TP 版。进度比咱们预期的要快得多，一个 (有教训的) 开发人员仅花了半天工夫就实现进去了，代码见 此处。在此过程中，咱们还重用了一些其余我的项目的代码，这对咱们很有帮忙。

提早从 300 毫秒 / 词元间接变为 91 毫秒 / 词元，这是用户体验的微小改良。
一个简略的 20 个词元的申请提早从 6 秒变成了 2 秒，用户体验间接从“慢”变成了轻微提早。

此外，吞吐量回升了很多，达到 10 RPS。batch_size=1 和 batch_size=32 提早基本相同，因而，从这种意义上来讲，在雷同的提早下，吞吐量的回升基本上是 收费 的。

唾手可得的果实

当初咱们有了一个 TP 版本的实现，咱们能够再次开始进行性能分析和优化。因为并行计划产生了扭转，咱们有必要再从头开始剖析一遍。

首先，同步 (ncclAllReduce) 开始成为次要热点，这合乎咱们的预期，同步须要花工夫。但咱们不打算优化这一部分，因为它曾经应用了 nccl。尽管可能还有一些改良空间，但咱们认为咱们很难做得更好。

第二个是 Gelu 算子，咱们能够看到它启动了许多 element-wise 类的核函数，总体而言它占用的计算份额比咱们预期的要大。

咱们对 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)))

咱们应用 jit 将许多小的 element-wise 核函数交融成了一个核函数，从而节俭了核函数启动开销和内存拷贝开销。

该优化升高了 10% 的提早，从 91 毫秒 / 词元到 81 毫秒 / 词元，搞定！

不过要小心，这种办法可不是任何时候都无效，算子交融不肯定每次都会产生。另外如果原来的算子实现曾经十分高效了，就算交融了也不能带来很多的增益。

咱们发现它在上面几个场合有用:

• 你有很多小的、element-wise 的操作
• 你的性能热点里有一些难以去除的 reshape 算子，这些算子个别就是拷贝
• 算子能交融时

滑铁卢

在测试期间，有一段时间，咱们察看到 Rust 服务的提早比 Python 服务低 25%。这很奇怪，但因为它们的测试环境是统一的，而且去除了核函数后咱们还是能测到这个速度增益，咱们开始感觉，兴许升高 Python 开销能够带来不错的性能晋升。

咱们开始了为期 3 天的从新实现 torch.distributed 局部代码的工作，以便在 Rust 里运行 nccl-rs。代码能工作，但生成的句子与 Python 版有些不一样，于是咱们开始考察这些问题，就在这个过程中，咱们发现 …… 在测量 PyTorch 版性能时，咱们遗记删除 PyTorch 里的 profiler 代码了 ……

咱们遭逢了滑铁卢，删除 profiler 代码后提早升高了 25%，两份代码提早一样了。其实咱们最后也是这么想的，Python 肯定不会影响性能，因为模型运行时运行的次要还是 torch cpp 的代码。尽管 3 天其实也不算啥，但产生这样的事还是挺蹩脚的。

针对谬误的或不具代表性的测量数据进行优化，这很常见，优化后果最终会令人悲观甚至对整个产品带来反成果。这就是为什么 小步快走 以及 设立正确预期 有助于管制这种危险。

另一个咱们必须分外小心的中央是产生第一个新词的前向过程 [译者注: 第一个新词 past_key_values 为 None ] 和产生后续新词的前向过程 [译者注: 此时 past_key_values 不为空] 是不一样的。如果你只针对第一个词优化，你反而会拖慢后续的那些更重要并且占大部分运行工夫的词的生成工夫。

另一个很常见的罪魁祸首是测量工夫，它测量的是 CPU 工夫，而不是理论的 CUDA 工夫，因而运行时须要用 torch.cuda.synchronize() 来确保 GPU 执行实现。

定制核函数

到目前为止，咱们曾经实现了靠近 DeepSpeed 的性能，而无需任何自定义代码！很简洁。咱们也不用在推理 batch size 的灵活性上做出任何斗争！

但依据 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 架构 (即 A100)，所以该核函数不适用于其余 GPU 架构; 同时咱们也不是编写核函数的专家，因而很有可能有更好的实现办法。

这个自定义核函数又提供了 10% 的提早晋升，提早从 81 毫秒 / 词元升高到 71 毫秒 / 词元。同时，咱们持续放弃了灵活性。

在那之后，咱们考察、摸索了更多优化伎俩，比方交融更多的算子来删除剩下的 reshape 等等。但还没有哪个伎俩能产生足够大的晋升而值得被放入最终版本。

Web 服务局部

就像咱们在 Rust 里做的一样，咱们必须实现对具备不同参数的申请的批处理。因为咱们处于 PyTorch 世界中，咱们简直能够齐全管制正在产生的事件。
而又因为咱们处于 Python 世界中，咱们有一个限度因素，即 torch.distributed 须要多过程而不是多线程运行，这意味着过程之间的通信有点麻烦。最初，咱们抉择通过 Redis 公布 / 订阅来传递原始字符串，以便同时将申请分发给所有过程。因为咱们处于不同的过程中，所以这样做比进行张量通信更容易、通信量也很小。

而后咱们不得不放弃应用 generate 函数，因为这会将参数利用于 batch 中所有的序列，而实际上每个序列的参数可能各不相同。值得庆幸的是，咱们能够重用较底层的 API，如 LogitsProcessor，以节俭大量工作。因而，咱们重构了一个 generate 函数，它承受一个参数列表并将列表中的参数别离利用于 batch 中的各个序列。

最终用户体验次要还是看提早。因为咱们反对不同的申请有不同的参数，因而可能呈现这样的状况: 一个申请想要生成 20 个词元，而另一个申请想要生成 250 个词元。因为每个词元须要 75 毫秒的提早，因而一个申请须要 1.5 秒，而另一个须要 18 秒。如果咱们始终进行批处理的话，咱们会让第一个用户期待 18 秒，因而看起来如同咱们正在以 900 毫秒 / 词元的速度运行，太慢了！

因为咱们处于具备极大灵活性的 PyTorch 世界中，咱们能够做的是在生成前 20 个词元后立刻从批处理中提取第一个申请，并在 1.5 秒内返回给该用户！这同时也节俭了 230 个词元的计算量。

因而，灵活性 对于获得最佳提早十分重要。

最初的笔记和疯狂的想法

优化是一项永无止境的工作，与任何其余我的项目一样，20% 的工作通常会产生 80% 的后果。

从某个工夫点开始，咱们开始制订一个小的测试策略来确定咱们的某个想法的潜在收益，如果测试没有产生显著的后果，咱们就会放弃这个想法。1 天减少 10% 足够有价值，2 周减少 10 倍也足够有价值。2 周进步 10% 就算了吧。

你试过……吗？

因为各种起因，有些办法咱们晓得但咱们没应用的。可能起因有: 感觉它不适宜咱们的场景、工作量太大、收益后劲不够大、或者甚至仅仅是因为咱们有太多的抉择要试而工夫不够所以就放弃了一些。以下排名不分先后:

• CUDA graphs
• nvFuser (它是 torch.jit.script 的后端，所以从这个角度来讲，咱们也算用了它。)
• FasterTransformer
• Nvidia’s Triton
• XLA (JAX 也应用 XLA!)
• torch.fx
• TensorRT

如果你最喜爱的工具没有列在这儿，或者你认为咱们错过了一些可能有用的重要工具，请随时与咱们分割！

Flash attention

咱们简略集成过 flash attention，尽管它在生成第一个词元 (没有 past_key_values) 时体现十分好，但在有了 past_key_values 后，它并没有产生太大的改良。而且如果咱们要用上它，咱们须要对其进行调整以反对 alibi 张量的计算。因而咱们决定临时不做这项工作。

OpenAI Triton

Triton 是一个用于在 Python 中构建定制核函数的杰出框架。咱们前面打算多用它，但到目前为止咱们还没有。咱们很想晓得它的性能是否优于咱们手写的 CUDA 核函数。过后，在做计划抉择时，咱们认为间接用 CUDA 编写仿佛是实现目标的最短门路。

填充和 reshape

正如本文通篇所提到的，每次张量拷贝都有老本，而生产环境中运行时的另一个暗藏老本是填充。当两个查问的长度不同时，你必须应用填充 (应用虚构标记) 以使它们等长。这可能会导致很多不必要的计算。更多信息。

现实状况下，咱们能够永远 不 做这些计算，永远不做 reshape。
TensorFlow 有 RaggedTensor 而 PyTorch 也有 嵌套张量 的概念。这两者仿佛都不像惯例张量那样精简，但能使咱们的计算更好，这对咱们有益处。
现实的状况下，整个推理过程都能够用 CUDA 或纯 GPU 代码来实现。思考到咱们在交融算子时看到性能改良，这种办法看起来很迷人。但咱们不晓得性能晋升能到什么水平。如果有更聪慧的 GPU 专家晓得，咱们洗耳恭听！

致谢

所有这些工作都是许多 HF 团队成员单干的后果。以下排名不分先后，@ThomasWang @stas
@Nouamane @Suraj
@Sanchit @Patrick
@Younes @Sylvain
@Jeff (Microsoft) @Reza
以及 BigScience 我的项目中的所有人。

英文原文: https://hf.co/blog/bloom-inference-optimization

作者: Nicolas Patry

译者: Matrix Yao (姚伟峰)，英特尔深度学习工程师，工作方向为 transformer-family 模型在各模态数据上的利用及大规模模型的训练推理。

排版 / 审校: zhongdongy (阿东)