大模型推理为什么 4090 很香
推理和训练有什么区别？
首先，训练不仅需要存储模型参数，还需要存储梯度、优化器状态、正向传播每一层的中间状态（activation），后面几个比参数更大，对模型内存的需求量也更大。
其次，训练任务是一个整体，流水线并行的正向传播中间结果是需要存下来给反向传播用的。为了节约内存而使用流水线并行，流水级越多，要存储的中间状态也就更多，反而加剧内存的不足。而推理任务中的各个输入数据之间并没有关系，正向传播每一层的中间状态也不需要保存下来，因此流水线并行不需要存储很多中间状态。
首先我们需要计算一下推理需要多少算力。前面针对训练算力的估算，为了简单起见，忽略了两个事情，首先是没有考虑 KV Cache，其次是没有考虑内存带宽。
KV Cache
什么是 KV Cache？对于每个输入的 prompt，在计算第一个 token 输出的时候，每个 token 的 attention 肯定是都要从头计算。但是在后续 token 的生成中，都需要计算 self-attention，也就是输入 prompt 以及前面输出的 token 的 attention。这是就需要用到前面每一个 token 的 K 和 V，由于每一层的参数矩阵是不变的，此时只有刚生成的那个 token 的 K 和 V 需要从头计算，输入 prompt 和之前生成的 token 的 K 和 V 其实是跟上一轮一样的。
这时，我们就可以把每一层的 K、V 矩阵缓存起来，生成下一个 token 的时候不再需要重新计算，这就是所谓的 KV Cache。Q 矩阵每次都不一样，没有缓存的价值。前面讲的训练中的选择性保存正向 activation 是个拿计算换内存的把戏，这里的 KV Cache 就是一个拿内存换计算的把戏。
KV Cache 需要多少存储容量呢？每一层，每个 token 的 K、V 矩阵都是 embedding size 这么大，再乘上 token 数量和 batch size，就是这一层的 KV Cache 所需的存储容量了。一定要记住 batch size，在正向和反向传播的几乎所有阶段，都不会涉及到对 batch size 中各个 sample 的合并处理，因此它始终是存储量和计算量计算中的一个系数。
例如，如果 batch size = 4，在 LLaMA 2 70B 中，假设输入和输出的 token 数量达到了模型的极限 4096，80 层的 KV Cache 一共需要 2 (K, V) * 80 * 8192 * 4096 * 8 * 2B = 80 GB。如果 batch size 更大，那么 KV Cache 占据的空间将超过参数本身占的 140 GB。
KV Cache 能省下来多少计算量？每一层计算 K、V 矩阵一共需要 2 (K, V) * 2 (mult, add) * embedding size * embedding size = 4 * 8192 * 8192 这么多计算量，乘以之前输入过的 token 数量、层数和 batch size，就是 4096 * 80 * 8 * 4 * 8192 * 8192 = 640 Tflops。相当于每存储 1 个字节，节约了 16K 次计算，还是很划算的。
事实上，KV Cache 节约的远远不止这些。计算 K、V 矩阵的过程是个典型的内存密集型过程，它需要加载每一层的 K、V 参数矩阵。也就是如果不做任何缓存，假设 prompt 长度很短而输出长度接近 token 的最大长度 4096，到了最后一个 token 的时候，单是重复计算前面每个 token 的 K、V 矩阵，就需要读取内存 4096 * 80 * 2 * 8192 * 8192 = 40T 次，每次 2 个字节，要知道 H100 的内存带宽只有 3.35 TB/s，4090 更是只有 1 TB/s，这单是最后一个 token 就得耗掉一张卡几十秒的时间来做重复计算。这样，token 的输出就会越来越慢，整个输出时间是输出长度平方级别的，根本没法用。
推理是计算密集还是存储密集
接下来我们就可以计算推理所需的计算量了。总的算力很好算，前面讲过，大概就是 2 * 输出 token 数量 * 参数数量 flops。如果想看细节，可以看下面这张图，来源是这里。
但算力并不能说明一切，模型还需要访问 GPU 内存，内存带宽也可能成为瓶颈。至少需要把参数从内存里面读出来吧？事实上，内存带宽的估算就这么简单，内存访问量 = 参数数量 * 2 bytes。中间结果有一部分是可以放在缓存里面的，缓存放不下的部分也需要占内存带宽，我们先不算。
如果不做任何批量输入，也就是模型专门服务一个 prompt，batch size = 1，整个 context 的长度很短（例如只有 128），那么整个推理过程中，每载入一个参数（2 字节），就只进行 128 次乘法和加法计算，那么计算 flops 和访问内存 bytes 的比例就只有 128。基本上任何 GPU 在这种情况下都会变成 memory bound，时间都耗在加载内存上了。
对于 4090 来说，计算 flops 和内存带宽之比是 330 / 1 = 330；对于 H100 来说，计算 flops 和内存带宽之比是 1979 / 3.35 = 590。也就是说，如果 context 中的 token 数量小于 330 或者 590，那么内存访问就会成为瓶颈。
虽然 LLaMA 2 的理论上限是 4096 个 token，但很多输入 prompt 用不了这么多，因此内存访问是有可能成为瓶颈的。此时，就需要靠 batch size 来补足了。推理中的批量处理，就是把几乎同时到达后端服务的 prompt 放到一起处理。不用担心，batch 里面的不同 prompt 的处理是完全独立的，不用担心会互相干扰。但这些 prompt 的输出是步调整齐划一的，每一轮整个 batch 中的每个 prompt 都会输出一个 token，因此如果有的 prompt 先输出完了，那就只能等其他的输出结束，造成一定的算力浪费。
有的人问，批量处理所需的算力跟分别单独处理所需的算力是一样的呀，那推理时为什么需要批量处理？答案就在访问内存的带宽上。
如果同时到达服务器的 prompt 很多，是不是 batch size 越大越好？也不是，因为 KV Cache 的大小可是正比于 batch size 的，batch size 大了，KV Cache 占据的 GPU 内存容量就很可观，比如在 LLaMA-2 70B 中，每个 prompt 都要占据 5 GB 的 KV Cache，如果 batch size 搞到 32，那么 KV Cache 就会占掉 160 GB 的 GPU 内存，比参数都大了。
70B 推理需要多少张卡？
总的存储容量也很好算，推理的时候最主要占内存的就是参数、KV Cache 和当前层的中间结果。当 batch size = 8 时，中间结果所需的大小是 batch size * token length * embedding size = 8 * 4096 * 8192 * 2B = 0.5 GB，相对来说是很小的。
70B 模型的参数是 140 GB，不管 A100/H100 还是 4090 都是单卡放不下的。那么 2 张 H100 够吗？看起来 160 GB 是够了，但是剩下的 20 GB 如果用来放 KV Cache，要么把 batch size 压缩一半，要么把 token 最大长度压缩一半，听起来是不太明智。因此，至少需要 3 张 H100。
对于 4090，140 GB 参数 + 40 GB KV Cache = 180 GB，每张卡 24 GB，8 张卡刚好可以放下。
推理用流水线并行可以吗？
推理使用流水线并行，最主要的问题是串行处理的推理延迟，网络延迟倒是小问题。
首先是推理延迟。虽然流水线的不同阶段可以塞进不同的 prompt，但同一个 prompt 的处理仍然永远在单个 GPU 上轮转，这样相比 Tensor parallelism 而言，单个 prompt 的延迟就增大了。
对于很小的 batch size，GPU 内存带宽是瓶颈，此时每张卡计算每个 token 的时延就是 2 byte * 参数量 / 卡的数量 / 内存带宽，例如 8 卡 4090 跑 LLaMA-2 70B，就是 2 * 70G / 8 / 1 TB/s = 0.0175 秒。这里没有考虑 KV Cache 带来的节约。注意，8 张卡是串行处理的，因此每个 token 的时延还要乘以 8，也就是 0.14 秒。每秒只能输出 7 个 token，对于 70B 这么小的模型来说是有点慢了。
对于很大的 batch size，GPU 算力是瓶颈，此时每张卡计算每个 token 的时延就是 batch size * 2 * 参数量 / 卡的数量 / 算力，例如 batch size = 1024，同样的 8 卡例子，就是 1024 * 2 * 70G / 8 / 330 Tflops = 0.0543 秒。事实上，对于这么大的 batch size，KV Cache 和正向传播的中间结果先把 GPU 内存给吃满了。
那么要平衡利用 GPU 算力和内存带宽，batch size 需要是多少呢？这就是 2 byte * 参数量 / 卡的数量 / 内存带宽 = batch size * 2 * 参数量 / 卡的数量 / 算力，左右两边参数量和卡的数量互相抵消，得到 batch size = 算力 / 内存带宽。对于 4090，就是 330 / 1 = 330；对于 H100，就是 1979 / 3.35 = 590。也就是说，对 4090 而言，batch size 小于 330 的时候 GPU 内存带宽是瓶颈，大于 330 的时候 GPU 算力是瓶颈。当 batch size = 330 的时候，理想情况下，内存带宽和算力恰好都打满，每张卡处理每个 token 的时间就是 17.5 ms。
其次是网络延迟。流水线并行相比张量并行的优点就是网络传输量小，流水级之间只需要传输 batch size * embedding size 这么多数据。例如 batch size = 8，embedding size = 8192，只需要传输 128 KB 数据，在 32 GB/s 的 PCIe Gen4 x16 上，只需要 4 us 就可以传输完成。当然，还需要考虑到通信库本身的开销，加上 4090 不支持 GPU 之间 P2P 传输，需要通过 CPU 中转，实际上需要几十 us 的时间，相比计算部分动辄几十 ms 的时延，可以忽略不计。
即使 batch size = 330，这 5.28 MB 数据在 PCIe 上也只需要传输 0.16 ms，相比计算部分的 17.5 ms 仍然可以忽略不计。
如果可以忍受流水线并行的推理延迟，甚至可以用多台主机来做流水线并行。我们假设主机间只有 1 Gbps 的普通以太网络，每台主机只有一张 4090。对于 batch size = 1，16 KB 数据需要 0.25 ms 才能传输完成，再加上 0.25 ms 两端网络协议栈的处理时间，每个流水级就需要 0.5 ms 的时延，8 张卡花在通信上的时间只有 4 ms，相比整体计算时延 140 ms 来说可以忽略，不会显著影响系统的推理延迟。
当 batch size 很小时，流水线推理中的网络流量是突发性（bursty）的，每过 18 ms 只会进行 0.25 ms 数据传输，只有 1/72 的占空比，不用担心流水线推理把局域网全部给占满了，搞得没法正常上网了。
如果为了充分利用算力，把 batch size 设置得很大，比如 330，那么 16 KB * 330 = 5.28 MB 数据需要传输 41 ms，8 张卡花在通信上的时间高达 0.33 秒，这样就只有 3 token/s 的输出速度了，难以忍受。因此，如果用主机间通信来做流水线并行，主机间又没有很高的通信带宽，就势必需要牺牲一定的吞吐量。
例如，我们设置输出速度不小于 5 token/s，这时留给通信的时间是 60 ms，每个流水级至多 7.5 ms，1 Gbps 网络可以传输 960 KB 数据，这时 batch size 至多设置为 60，也就是这 8 张 4090 的总吞吐量是 2400 token/s。此时的有效算力利用率只有不到 20%。
最近有一个比较火的 Petals 开源项目，就是利用流水线并行，把 GPU 做成了一个类似 BitTorrent 的分布式网络。虽然推理延迟确实比较高，但至少说明了分布式 GPU 推理的可行性。
推理用张量并行怎么样？
前面讲到，流水线并行的最大缺点是 GPU 串行处理，延迟较高，导致输出 token 比较慢。而张量并行的最大缺点是传输数据量大，网络带宽低的设备不一定 hold 得住。
但是推理要传输的数据量跟训练要传输的数据量可不是一回事啊！推理只需要传输正向传播的中间结果（activation），而训练还需要传输所有参数的梯度，梯度才是数据量的大头。
在推理中，如果使用张量并行，Transformer 的每一层都需要传输把自己负责的结果向量（大小为 batch size * embedding size / num GPUs）广播给其他所有 GPU，并接受来自所有其他 GPU 广播来的数据。计算 attention 的时候需要传输一次，计算 feed-forward network 的时候又需要传输一次，也就是总共需要传输 2 * 层数这么多次。
每次发送就是 batch size * embedding size（发送和接收是不同的方向，不能算两次），对于 batch size = 1, embedding size = 8192，只需要传输 16 KB 数据，在 32 GB/s 的 PCIe Gen4 上传输只需要 1 us。当然，考虑到前面讨论的 CPU 中转开销，还是需要大约 30 us 的。一共 160 次传输，需要 4.8 ms。
我们再考虑计算的开销。还是考虑 batch size = 1 的情形，GPU 内存带宽是瓶颈，此时每张卡计算每个 token 的时延就是 2 byte * 参数量 / 卡的数量 / 内存带宽，代入我们前面的数值，仍然是 17.5 ms。但是这里 8 张卡是并行处理的，因此总的处理时长就是计算时间 + 通信时间 = 17.5 ms + 4.8 ms = 22.3 ms。这就意味着每秒可以生成 45 个 token，这个 token 生成速度已经很不错了，至少人类的阅读速度是很难赶上生成的速度了。
如果 batch size 更大会怎样？例如 batch size = 330，把 GPU 算力和内存带宽都充分利用起来，每次需要传输的数据量是 330 * 8192 * 2 = 5.4 MB，在 32 GB/s 的 PCIe Gen4 上需要 0.17 ms。一共 160 次传输，就是 27 ms。这下网络通信开销成了延迟的大头，总处理时长为 27 + 17.5 = 44.5 ms，每秒只能生成 22 个 token 了，但也不算慢。
注意，不管用多少个 GPU 做并行推理，只要用的是张量并行，网络传输的总数据量是相同的，因此增加 GPU 的数量只能加速计算，不能加速通信。
因此，A100/H100 的 NVLink 在降低推理延迟方面还是有很大作用的。如果用 A100/H100，取 batch size = 590 达到算力和带宽的平衡利用，这 9.44 MB 数据只需要 9.44 MB / 450 GB/s = 0.02 ms。一共 160 次传输，也只有 3.2 ms。由于内存带宽大了，计算时间也可以大幅缩短，例如 H100 的计算时间为 2 * 70G / 8 / 3.35 TB/s = 5.2 ms。总处理时长只有 5.2 ms + 3.2 ms = 8.4 ms，每秒可以生成 119 个 token，非常棒！
可以说，如果论单个 prompt 的 token 生成速度，无论用多少块 4090 也追不上 8 卡 H100。
用 4090 做推理的成本怎么样？
对于推理，不管用流水线并行还是张量并行，batch size 不算高到太离谱的情况下内存带宽都是瓶颈。
假如 batch size 能够高到把算力 100% 利用起来，并且还能解决 KV Cache 不够大的问题，能解决中间结果占用内存过多的问题，那么这 8 张 4090 可以达到多少吞吐量？
当然，这两个问题都不好解决，因此推理优化才是一个热门的研究领域，存在很多的 trade-off 和奇技淫巧。如果只是用标准的 PyTorch，那推理性能距离把算力 100% 利用起来还远得很呐。
假设都解决了，在张量并行的通信过程中我们可以利用 double buffer 做另外一个 batch 的计算，也就是计算和通信并行，进一步提高吞吐量。通信和计算分别是 27 ms 和 17.5 ms，传输的 27 ms 是瓶颈，也就是每 27 ms 输出一组 token，一个 batch 330 个 prompt，那这 8 张 4090 真是可以达到每秒 330 / 0.027 = 12.2K token 的吞吐量。
8 张 4090 的成本是 12800 美金，8 卡 PCIe Gen4 服务器本身要 2 万美金，加上网络设备，平均每台 4 万美金的设备成本。固定资产按照 3 年摊销，每小时 1.52 美元。整机功耗大约 400W * 8 + 2 kW = 5 kW，按照 0.1 美元一度电算，每小时 0.5 美元。这 2 美元一小时的机器，满打满算能生成 12.2K * 3600 = 44M tokens，也就是说 1 美元能生成 22M tokens。
是不是比 GPT-3.5 Turbo 的 $0.002 / 1K tokens，也就是 1 美元 0.5M tokens 便宜 44 倍？当然，账不能这么算。
- 首先，GPU 的算力利用率到不了 100%；
- 其次，如同所有 SaaS 服务一样，用户的请求数量有波峰有波谷，用户是按量付费的，平台提供方可是不管有没有人用都在烧钱的；
- 此外，每个 batch 中不同 prompt 的长度和响应 token 数量都不同，消耗的算力是 batch 中最大的那个，但收的钱是用户实际用的 token 数；
- 再次，GPT-3.5 是 175B 的模型，比 70B 的 LLaMA 很可能推理成本更高；
- 最后，OpenAI 开发 GPT-3.5 是烧了不知道多少钱的，人家至少要赚回训练成本和研发人员的工资吧。
其实 GPT-3.5 Turbo 的 $0.002 / 1K tokens 真的挺良心的，有的卖 API 的，LLaMA-2 70B 都敢比 GPT-3.5 Turbo 卖得贵。
如果换成用 H100 做推理，重新算一下这笔账。一张 H100 至少要 3 万美金，一台 8 卡 H100 高配服务器加上配套的 IB 网络，起码要 30 万美金，同样按照 3 年摊销，每小时 11.4 美元。10 kW 功耗，电费每小时 1 美元。一共 12.4 美元一小时。
这其实已经是非常良心的价格了，你在任何云服务商都不可能租得到这么便宜的 8 卡 H100。所以说从云服务商租卡卖没有护城河的 SaaS 服务，比如开源模型的推理 API，除非有一种提高推理性能的独门绝技，很难赚得了什么大钱，二房东的生意不是这么好做的。
再算算这台 8 卡 H100 机器的吞吐量，张量并行也采用传输和计算并行，H100 的通信比较快，因此计算是瓶颈，每 5.2 ms 可以输出一组 token，一个 batch 590 个 prompt，满打满算可以达到每秒 590 / 0.0052 = 113K token 的吞吐量。理想情况下，一小时能生成 407M tokens，也就是 1 美元能生成 33M tokens，H100 这单位 token 的成本比 4090 还要低 30%。
为什么 8 卡 H100 机器是 4090 机器价格的 6 倍，性价比却比 4090 高？因为一张 H100 的算力是 4090 的 6 倍，内存带宽是 4090 的 3.35 倍，当 batch size 够大，算力达到瓶颈的时候，单卡的性能就是 6 倍。而且，H100 比 4090 的网络带宽强太多了，导致 4090 在张量并行中网络通信成了瓶颈，浪费了有效算力。因此，同样的 8 卡机器吞吐量几乎可以达到 4090 的 10 倍。虽然一张 H100 卡的价格是 4090 的 20 倍以上，但算上服务器本身的成本和电费，整机的成本只是 6 倍左右。
用最便宜的设备搞出最高的推理性能
我们发现在 8 卡 4090 机器中，3 万美金的设备成本，GPU 卡只占了 1.28 万美金，不像 H100 机器那样 GPU 成本占了大头。还有办法进一步降低吗？
如果我们可以忍受 5 token/s 的输出速度，甚至可以利用流水线并行，用家用台式机和 4090 攒出个推理集群来。
遥想我当年在 MSRA 的时候，在一台只用 1000 美金攒出来的机器上插了 10 块 FPGA，做出个世界最快的 Key-Value Store。其实如果让我去设计一个性价比最高的 4090 推理集群，有很多种方案可以尝试：
- 用流水线并行，台式机 + 10 Gbps 网卡，足够在 5 ms 内传输 batch size = 330 的 5.28 MB 数据了，通信 40 ms，计算 140 ms，达到 5 token/s 的单 prompt 输出速度，同时又能充分利用 4090 的算力。10 Gbps 的网卡和交换机都很便宜，Intel X710 网卡只要 150 美金，20 口交换机只要 1500 美金（每 8 个口 750 美金），一台家用台式机 700 美金，这只要 2 万美金就可以搞定原本需要 4 万美金的设备。
- 用张量并行，台式机 + 200 Gbps ConnectX-6 网卡，上 RoCE，可以把 batch size = 330 的 5.28 MB 数据在 0.22 ms 内传完，160 次传输是 35 ms，加上计算的 17.5 ms，一个 token 52.5 ms，可以达到 19 token/s 的单 prompt 输出速度，这个速度已经不错了。网卡 1000 美金，200G 交换机 2 万美金 40 个端口，平均每 8 个端口 4000 美金，一台家用台式机 700 美金，这只要 3 万美金就能搞定原本 4 万美金的设备。
- 主机内用张量并行，主机间用流水线并行，4 卡 PCIe Gen4 服务器主板只要 1000 美金而且能跑满 PCIe 带宽（因为 8 卡就需要 PCIe switch 了，价格会贵很多），两台主机之间用 25 Gbps 网卡直连，主机内张量并行的时延是 27 ms，主机间流水线并行只需 2 次 8 ms 的传输（注意 25G 的网络带宽是 4 张 GPU 卡共享的），加上两次流水线计算各 17.5 ms，总共 78 ms，可以达到 13 token/s 的单 prompt 输出速度。网卡 300 美金 * 2，服务器 3000 美金 * 2，这只要 1.95 万美金就可以搞定原本需要 4 万美金的设备。
2 万美金按照 3 年摊销是每小时 0.76 美元。按照 0.1 美元/度的电价，每小时的电费都要 0.5 美元，接近设备成本了，这有点挖矿的味道了。这 1.26 美元一小时的机器如果跑满了 44M tokens 的吞吐量，1 美元能生成 35M tokens，终于赶上 8 卡 H100 的 33M token per dollar 了。
为什么 H100 以 20 倍于 4090 的 GPU 价格，9 倍的性能，却仍然能在系统性价比上打个平手，首先是因为能耗成本更低，8 卡 H100 的功耗是 10 kW，但 9 台 8 卡 4090 的功耗是 45 kW；其次是因为主机和网络设备成本更低，一台 8 卡 H100 准系统虽然贵，但只占整机价格的 20% 左右；但 4090 因为卡多，除非像 GPU 矿机那样压成本，只要还是用数据中心级的设备，准系统价格就要占到 35% 以上。
其实，这个世界上不止有 A100/H100 和 4090，还有 A10 等计算卡和 3090 等游戏卡，还有 AMD 的 GPU 和很多其他厂商的 AI 芯片。H100 和 4090 大概率都不是性价比的最优解，例如 A10 和 AMD GPU 的性价比有可能就更高。
我都想搞一个推理性价比挑战赛，看谁能用最便宜的设备搞出最强的推理吞吐量，同时延迟不能太高；或者用最便宜的设备搞出最低的推理延迟，同时吞吐量不能太低。
这一切都是在假设使用 LLaMA-2 70B 模型，没有做量化压缩的前提下。如果做了量化压缩，那性能就更高，甚至在 Unified Memory 够大的 MacBook Pro 上都能单机跑了。