在 TPU 上部署 LLaMA 3-70B 服务

LLaMA 3-70B 有 8 个 KV 头，因此每个 token 的大小为 2 * K * H * L = 2 * 8 * 128 * 80 = 160kB。

请注意它有多大！ 如果序列长度为 32k token（这很常见），每个序列会占用 160e3 * 32,768 = 5.3GB / sequence。当 BS=240 时，总计为 1.3TB！由于每颗 TPU v5e 只有 16GB，我们至少需要约 (70e9 + 1.3e12) / 16e9 = 86 颗 TPU v5e 芯片，才能容纳这么多内存。还要注意，与 70GB 的模型参数相比，这个数有多大。

因为 int8 下每个 token 的 KV 为 160e3 字节，所以 KV 总内存为 160e3 * 8192 * 32 = 41.9e9 字节。参数为 70e9 字节，因为每个参数占 1 字节。因此，总内存用量为 41.9e9 + 70e9 = 112GB。

可用的最小切片需要 112e9 / 16e9 = 7 颗 TPU；将其向上取整为合理的偶数规模，就是 TPU v5e 4x2。这会非常紧张；把其他开销考虑进去后，我们可能无法真正放得下，因此至少可能需要 4x4（或者降低批大小）。

可以调用上一节中的公式：

这里的临界批大小约为 120，因为参数为 int8，而 FLOPs 采用 bfloat16。我们也可以手动计算等式右侧的最大值，但本质上这个计算已经做过好几次了。所以，无论对矩阵乘法还是 FLOPs 而言，我们都已深入内存带宽受限区间。

如果严格只看内存带宽，步骤时间基本为 (KV size + param size) / (8 * HBM bandwidth) = 112e9 / (8 * 8.2e11) = 17ms。因此，理论步骤时间约为 17ms。 吞吐量为 32 / .017 = 1882 tokens / sec，即 1882 / 8 = 235 tokens / sec / chip。

这里有一个注意事项：要检查矩阵乘法是否可能受 ICI 限制。这里可以为它投入 2 个轴，因此理论上当 $Y > 2 * F / 2200 = 2 * 28672 / 2200 = 26$ 时才会受 ICI 限制，所以完全没问题！

如果改在 4x4 上运行，ICI 方面仍然没问题，因此延迟会降至 17 / 2 = 8.5ms，但每颗芯片的吞吐量保持不变。

如第 7 节所述，对于任何满足 $B \ll D, F$ 的 bfloat16 矩阵乘法，都有

当权重采用 int8 时，分母会少一个 2 倍因子，于是得到 $2BDF / DF = 2B > 240$，等价于 $B > 120$，也就是之前临界批大小的一半。这对我们非常有帮助！当权重和 FLOPs 都采用 int8 时，必须使用 TPU 的 int8 FLOPs/s 数值；它从 bfloat16 的 1.97e14 提升到 3.94e14，接近翻倍。这意味着我们又回到了起点，约为 $B > 240$。

int8 权重配合 bfloat16 FLOPs 的情况十分常见，因为对参数做无损量化通常比执行低精度算术更容易。

这很简单！如果可以接受极小的批大小，那么唯一限制就是参数内存必须装入 HBM；也就是说，芯片数就是 ceil(num_params * sizeof(dtype) / HBM per TPU)，或者写成 ceil(70e9 * sizeof(dtype) / 16e9)，再向上取整到最近的合理拓扑（2 的某个倍数）：

这相当不错！它告诉我们，只要愿意，就能把 LLaMA 70B 放进 TPU v5e 2x2。只不过你会注意到 KV 缓存数量很小。那就是我们的批大小！这意味着 FLOPs 利用率会非常糟糕。为了把批大小推高到 240，我们会很乐意改用更大的拓扑。

这也很简单，因为我们选择的批大小恰好会用满全部 HBM！问题只是在于，把整颗 TPU v5e 所能容纳的字节载入 MXU 需要多久。它就是 v5e HBM / v5e HBM memory bandwidth = 16GB / 8.2e11 = 19ms，所以结果是每步 19ms。假设生成结果的中位长度为 512 个 token，则每次解码大约需要 9 秒。请注意，更小的批大小可以带来略低的延迟；例如，如果只考虑 int4 模型参数，最低延迟约为每步 10ms，因为此时 HBM 不再是满的。

这是一个重要问题，因为它与每个 token 的成本完全相关。

按照对解码长度中位数的假设，吞吐量就是 $B / (\text{per-step latency} \cdot \text{median steps} \cdot N) \approx 43 / (0.019 * 512 * N)$。这大约为 $(4.42 / N)$ QPS，因此代入 $N$ 可得：

请注意，这个估计相当乐观，因为它完全忽略了前向传播的工作内存（分配给激活和注意力的内存）。使用 Flash Attention 时，这并非荒谬的假设，但也不现实。真实数值很可能只有这里的一半左右。要获得绝对最高的吞吐量，我们可能需要把芯片数量增加到两倍以上，并显著增大批大小。

如果在 bfloat16 下使用 4x8 切片，就会剩余 372GB 可供 KV 缓存使用，从而能把批大小增至 140。由于步骤时间保持不变，吞吐量将为 14.39 / num_chips，即：

进一步增加规模还会带来更大的收益！最重要的结论是：如果受到 KV 缓存大小限制，那么最小拓扑并不总是性能最高的拓扑。

如上一节所述，生成期间真正可用的分片选择只有一种：模型并行。在受到通信限制之前，最多可以做多少模型并行？如上一节讨论的那样，模型大致会在下式成立时变得通信受限：

对 LLaMA 3-70B 而言，F = 28,672；因此，如果沿 2 个轴进行模型分片，大约得到 $Y = 28672 \cdot 2 / 2200 = 26$。所以总体上可扩展至约 16 颗芯片而不受通信限制，这允许使用 4x4，但不允许使用 4x8。一般而言，因为我们无法完美重叠计算，所以就连这个估计也过于乐观。

要点：我们实际上无法用纯模型并行在 4x8 上提供服务。 这里最多只能用 4x2，或者也许能用 4x4。

不过，如前所述，当批大小较小时，往往可以采用更多模型并行而不显著损害吞吐量，因为模型受内存带宽限制，而不是受 FLOPs 限制。前面说过，这个值大约为 $Y=F / (8\cdot B)$；因此，如果批大小为 64，理论上在变为 ICI 受限之前，可以达到 Y = 28,672 / (8 * 64) = 56 路模型并行。为了对结果做合理性检查，可以考察单次矩阵乘法的 $T_\text{ici comms}$、$T_\text{hbm comms}$ 和 $T_\text{math}$。显然有：

对 4x8 而言，$T_\text{ici comms}$ = (2 * 64 * 8192) / 9e10 = 11us，$T_\text{hbm comms}$ = (2 * 8192 * 28,672) / (32 * 8.2e11) = 18us，$T_\text{math}$ = (2 * 64 * 8192 * 28,672) / (32 * 1.97e14) = 4us，因此理论上仍然受 HBM 带宽限制，这很好！请注意，从 4x4 扩展到 4x8 从吞吐量角度看可能没有帮助，但会降低延迟！

如果考察 int8 和 int4 配置，就会发现它们可以用纯模型并行实现。因此，我们来到了量化带来实际优势的地方，而且这种优势超出了更快的 FLOPs：它使我们能在变为通信受限之前采用更大的批大小。所以，这个故事的最终结论是：在 4x8 上无法达到峰值吞吐量，但对 int8 和 int4 配置而言，可以使用纯模型并行。

在 8k token 下，我们显然处于计算受限状态，所以只需要从 FLOPs 的角度推理。已知模型有 70e9 个参数，因此每次前向传播使用 2 * 70e9 * B FLOPs。假设 MFU（FLOPs 利用率）为 40%，运行时间约为 2 * 70e9 * 8192 / (16 * 1.97e14 * 0.4) = 0.91s。与前面考察的数值相比，这其实相当可观！

这个问题相当直接。因为解码长度中位数为 4096 个 token，所以一个序列大约每 1 / 4096 个 token 完成一次。给定批大小 32，这意味着每个步骤会逐出 32 / 4096 个序列。由于 KV 缓存长度约为 8192 + 4096，所以每个步骤会逐出 32 * (8192 + 4096) / 4096 = 96 个 token。通用公式为 $B * (P + G) / G$，其中 $P$ 和 $G$ 分别是预填充长度和生成长度。

这个问题颇有意思。设 $P$ 为预填充服务器数量，$G$ 为生成服务器数量。总体来说，这是一个流水线问题：以 P / prefill_latency 的速率输入序列，以 B * G / (generate_latency * median_decode_length) 的速率消费序列。我们之前算得，每个预填充步骤为 910ms，批大小 43（这里就按 32 计）时每个解码步骤为 19ms。因此需要 P / 0.91 = 32 * G / (0.019 * 512)，也就是 P = 3G；换言之，预填充服务器数量大约要是生成服务器的 3 倍！

下面是计算这些 Roofline 的代码：

import numpy as np
 
num_chips = 16  # we fix 16 as the amount of total model parallelism we do
bytes_per_param = 1  # int8 means 1 byte per param
param_count = 70e9
param_size = bytes_per_param * param_count
sequence_length = 8192  # can vary this
 
hbm_bandwidth = 8.20E+11  # v5e
flops = 1.97E+14  # v5e
 
def kv_cache_size(bs):
    return 2 * bs * 128 * 8 * 80
 
def min_topology(bytes):
    return 2 ** np.ceil(np.log2(bytes / 16e9))
 
def get_max_batch_size(
    num_chips: int,
    sequence_length: int,
    param_size: float,
) -> int:
  batch_sizes = np.arange(1, 1024, 4)
  kv_sizes = kv_cache_size(sequence_length * batch_sizes)
  required_chips = min_topology(kv_sizes + param_size)
  max_idx = np.where(required_chips <= num_chips)[0][-1]
  return max_idx
 
max_idx = get_max_batch_size(
    num_chips=num_chips,
    sequence_length=sequence_length,
    param_size=param_size,
)  # get the largest batch size that can fit
batch_sizes = np.arange(1, 512, 1)[:max_idx]
kv_sizes = kv_cache_size(sequence_length * batch_sizes)
 
kv_comms_time = kv_sizes / (num_chips * hbm_bandwidth)
 
param_comms_time = param_size / (num_chips * hbm_bandwidth)
param_comms_time = np.asarray([param_comms_time] * batch_sizes.shape[0])
 
flops_time = 2 * param_size * batch_sizes / (num_chips * flops)  # roughly true in a 2ND sense
 
mlp_time = np.maximum(flops_time, param_comms_time)
attn_time = kv_comms_time  # always bandwidth-bound for generate
 
latency = 1000 * (mlp_time + attn_time)
throughput = batch_sizes / (latency * num_chips)

import numpy as np
 
num_chips = 16  # we fix 16 as the amount of total model parallelism we do
bytes_per_param = 1  # int8 means 1 byte per param
param_count = 70e9
param_size = bytes_per_param * param_count
sequence_length = 8192  # can vary this
 
hbm_bandwidth = 8.20E+11  # v5e
flops = 1.97E+14  # v5e
 
def kv_cache_size(bs):
    return 2 * bs * 128 * 8 * 80
 
def min_topology(bytes):
    return 2 ** np.ceil(np.log2(bytes / 16e9))
 
def get_max_batch_size(
    num_chips: int,
    sequence_length: int,
    param_size: float,
) -> int:
  batch_sizes = np.arange(1, 1024, 4)
  kv_sizes = kv_cache_size(sequence_length * batch_sizes)
  required_chips = min_topology(kv_sizes + param_size)
  max_idx = np.where(required_chips <= num_chips)[0][-1]
  return max_idx
 
max_idx = get_max_batch_size(
    num_chips=num_chips,
    sequence_length=sequence_length,
    param_size=param_size,
)  # get the largest batch size that can fit
batch_sizes = np.arange(1, 512, 1)[:max_idx]
kv_sizes = kv_cache_size(sequence_length * batch_sizes)
 
kv_comms_time = kv_sizes / (num_chips * hbm_bandwidth)
 
param_comms_time = param_size / (num_chips * hbm_bandwidth)
param_comms_time = np.asarray([param_comms_time] * batch_sizes.shape[0])
 
flops_time = 2 * param_size * batch_sizes / (num_chips * flops)  # roughly true in a 2ND sense
 
mlp_time = np.maximum(flops_time, param_comms_time)
attn_time = kv_comms_time  # always bandwidth-bound for generate
 
latency = 1000 * (mlp_time + attn_time)
throughput = batch_sizes / (latency * num_chips)

请注意，我们非常明确地把延迟拆分为两个来源：KV 加载与参数加载；而延迟受 FLOPs 或通信中的较大者所限制。