""" ============================================================================= CS349E · 第 4 讲 · 04_latency_throughput_cost.py —— 部署一个 70B 模型：切几份、批多大、延迟、吞吐和成本 ============================================================================= 【这个脚本要回答的问题】 把前面的工具合起来，回答一个真实的部署问题：Llama-3.1-70B 放在 H100 上（一台机器 8 块，NVLink 交换机）， 用几块卡做张量并行、批开多大、用什么精度？每个选择怎样影响"每个用户的速度"和"每个 token 的成本"？ 演示 1：张量并行度与批大小：每一步 = 读权重 / KV + 计算 + 每层两次 all-reduce。通信给延迟设了一个下限。 演示 2：延迟–成本前沿：把所有（并行度、批大小、精度）组合画成一张表，找出"同样延迟下最便宜"的那些点。 演示 3：张量并行切多少份才划算（How to Scale Your Model 第 7 章的 ICI 上限）。 演示 4：一台机器放不下或者不够快时：跨机器的张量并行 vs 流水线并行 vs 多开几个副本。 【模型的假设】roofline 下界：峰值按 100%，计算 / 访存 / 通信分别算，计算和访存取较大者、通信不重叠（偏保守）； all-reduce 用前置讲 P5 的 α–β 公式（NVLink：450 GB/s、α = 2 µs；InfiniBand：每卡 50 GB/s、α = 5 µs，α 是量级假设）。 上下文统一取 4096。所有输出每次运行完全相同。 【运行】python3 04_latency_throughput_cost.py （< 1 秒） 【依赖】同目录的 shard_lab；../prereq5/comm_sim；../lecture01 的 console、hw_specs、llm_math。 """ import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) sys.path.insert(0, os.path.join(HERE, "..", "prereq5")) from comm_sim import IB_NDR, Comm, allreduce_hierarchical, allreduce_ring, two_level # noqa: E402 from console import fmt_time, print_header, print_observations, print_table # noqa: E402 from hw_specs import H100, TPU_V5E # noqa: E402 from llm_math import LLAMA31_70B # noqa: E402 from shard_lab import NVLINK4, tp_decode_step # noqa: E402 M = LLAMA31_70B CTX = 4096 def fits(n, w_bytes, batch, kv_bytes=2.0, reserve=0.1): """n 块 H100 放不放得下：权重 + batch 个 CTX 长的序列的 KV（考虑 n > KV 组数时的复制）。""" w = M.n_params() * w_bytes kv = batch * CTX * M.kv_bytes_per_token(kv_bytes) * max(1.0, n / M.n_kv_heads) return w + kv <= n * H100.mem_cap * (1 - reserve) # --------------------------------------------------------------------------- # 演示 1：张量并行度与批大小 # --------------------------------------------------------------------------- def demo1_tp_degree(): print_header("演示 1：张量并行度 × 批大小——Llama-3.1-70B、BF16、H100（NVLink），上下文 4096") rows = [] for n in (2, 4, 8): for B in (1, 16, 64, 256): if not fits(n, 2.0, B): rows.append([n, B, "放不下", "", "", "", "", "", ""]) continue r = tp_decode_step(M, H100, NVLINK4, n, B, CTX, algo="ring") d = tp_decode_step(M, H100, NVLINK4, n, B, CTX, algo="rd") best = min(r["t_step"], d["t_step"]) rows.append([n, B, fmt_time(max(r["t_math"], r["t_mem"])), fmt_time(r["t_comm"]), fmt_time(d["t_comm"]), fmt_time(best), f"{r['t_comm'] / r['t_step'] * 100:.0f}%", f"{1 / best:.0f}", f"{B / best / n:,.0f}"]) print("每一步：计算 / 访存取较大者 + 160 次 all-reduce（80 层 × 2）；最后三列按 ring 和递归倍增中更快的那个算：") print_table(["n", "B", "计算 / 访存", "通信（ring）", "通信（递归倍增）", "一步", "ring 时通信占比", "tok/s/用户", "tok/s/芯片"], rows) print_observations([ "B = 1 时，n 从 2 加到 8，读权重的时间从 21 ms 降到 5.3 ms（8 块卡的 HBM 一起读）——张量并行**能降低单用户延迟**，" "这是它在推理里最重要的用途。", "但 ring all-reduce 的步数随 n 增长：n = 8 时 160 次 all-reduce 一共 4.5 ms，和读权重差不多；" "换成递归倍增（每次 3 步）只要 1 ms。小消息的集合通信算法选错，延迟几乎翻倍（前置讲 P5 §5.3、§7.1）。", "批变大以后通信的**占比**下降（n = 8 时从 46% 降到 28%），但通信时间本身在涨：B = 256 时每次 all-reduce 的消息有 4 MB，" "带宽项也不小了。**小批量、低延迟的部署对通信的延迟最敏感**。", "每芯片吞吐（tok/s/芯片）随 n 增大略降：同样的批，切得越细，每块卡分到的活越少，通信却没少。", ]) # --------------------------------------------------------------------------- # 演示 2：延迟–成本前沿 # --------------------------------------------------------------------------- def demo2_frontier(): print_header("演示 2：延迟–成本前沿（Llama-3.1-70B，H100，上下文 4096，all-reduce 用递归倍增）") pts = [] for prec, wb, cp in (("BF16", 2.0, "bf16"), ("FP8", 1.0, "fp8")): for n in (1, 2, 4, 8): for B in (1, 2, 4, 8, 16, 32, 64, 128, 256): if not fits(n, wb, B): continue st = tp_decode_step(M, H100, NVLINK4, n, B, CTX, w_bytes=wb, algo="rd", compute_prec=cp) t = st["t_step"] cost = n * t / B # 芯片·秒 / token pts.append((t, cost, prec, n, B)) pts.sort() frontier, best = [], float("inf") for p in pts: if p[1] < best: frontier.append(p) best = p[1] rows = [[f"{t * 1e3:.2f} ms", f"{1 / t:.0f}", f"{c * 1e6 / 3600:.3f}", prec, n, B] for t, c, prec, n, B in frontier] print(f"一共 {len(pts)} 种放得下的组合；下面是前沿上的点（没有别的组合能同时更快、更便宜）：") print_table(["每 token 时间", "tok/s/用户", "芯片·时 / 百万 token", "精度", "张量并行度 n", "批大小 B"], rows) print_observations([ "前沿的最左端（最快）是 FP8、n = 8、B = 1：每个用户约 280 tok/s，但 8 块卡只服务一个用户——成本最高。", "沿着前沿往右，主要是**加大批**：B 从 1 到 128，每个 token 的成本降了约 38 倍，每个用户只慢了约 3.4 倍；" "最后一个点换成 n = 4、B = 128，再便宜一点。最便宜的点比最快的点便宜约 46 倍，每个用户慢约 5.6 倍。", "前沿上全是 FP8：字节减半，同样的延迟需要的卡更少（第 3 讲）。n = 1（FP8 刚好放下权重）剩下的显存只够一个 4K 序列，" "成本反而很高，不在前沿上；更大的批（B = 256）只有 n = 8 放得下，但并不比 n = 4、B = 128 便宜。", "**没有'最好'的配置，只有'在你的延迟要求下最便宜'的配置**（第 1 讲 §1.3、§4.5）。Pope et al. 图 1 画的就是这样的前沿。", ]) # --------------------------------------------------------------------------- # 演示 3：张量并行的上限 # --------------------------------------------------------------------------- def demo3_ici_bound(): print_header("演示 3：decode 时张量并行能切多少份——How to Scale Your Model 第 7 章的上限 Y > F / (B·β)") rows = [] for name, chip, link_bw in (("H100 + NVLink", H100, NVLINK4.bw), ("TPU v5e + ICI（一个轴，双向）", TPU_V5E, 9e10)): beta = chip.mem_bw / link_bw for B in (8, 32, 128, 512): rows.append([name, f"{beta:.1f}", B, f"{M.d_ff / (B * beta):,.0f}"]) print_table(["硬件", "β = HBM 带宽 / 链路带宽", "批大小 B", "切到多少份通信才赶上读权重"], rows) lat_bytes = NVLINK4.bw * 1e-6 print(f"\n另一条线：消息小于 链路带宽 × 1 µs ≈ {lat_bytes / 1e3:.0f} KB 时，通信基本是延迟主导；" f"Llama-3.1-70B 每次 all-reduce 的消息 = B × {M.d_model} × 2 字节，B < {lat_bytes / (M.d_model * 2):.0f} 时都在这一区。") print_observations([ "推导（memory-bound 的 decode）：每层读权重 2·D·F / (Y·W_hbm)，每层通信约 2·B·D / W_link。" "通信超过读权重时 Y > F / (B·β)。批越小，能切的份数越多——小批量本来就该多切，用更多卡的 HBM 带宽换延迟。", "H100 一台机器只有 8 块卡（NVLink 域），这个上限在 B ≤ 128 时都大于 8（B = 512 时恰好约 8）：**单机内的张量并行在带宽上几乎总是划算的**，" "真正的限制是延迟（每次 all-reduce 几微秒，见演示 1）和 NVLink 域的大小。", "这和训练不同：训练要 compute-bound，张量并行只能切到 F / 2200 左右（4–8 份）；decode 是 memory-bound，可以切得更细。", ]) # --------------------------------------------------------------------------- # 演示 4：跨机器 # --------------------------------------------------------------------------- def cross_node_allreduce(msg_bytes): """两台机器 × 8 块 H100 上一次 all-reduce 的时间：ring 和层次化（前置讲 P5 §6.3）中更快的那个，用 comm_sim 模拟。""" n_el = 1 << 10 rng = np.random.default_rng(0) bufs = [rng.standard_normal(n_el) for _ in range(16)] times = [] for fn in (allreduce_ring, lambda c, b: allreduce_hierarchical(c, b, per_node=8)): c = Comm(two_level(2, 8), byte_scale=msg_bytes / (n_el * 8)) fn(c, bufs) times.append(c.time) return min(times) def demo4_cross_node(): print_header("演示 4：16 块 H100（2 台机器 × 8）同时服务 128 个序列——跨机器张量并行 vs 流水线 vs 两个副本（BF16）") C = 128 tp8_full = tp_decode_step(M, H100, NVLINK4, 8, C, CTX, algo="rd") # 参照：一台机器扛全部 128 个 # TP16：每块卡的计算 / 访存再减半（KV 有 2 份：16 > 8 组 KV），每层两次 all-reduce 都要跨机器 t_ar16 = cross_node_allreduce(C * M.d_model * 2) tp16 = tp_decode_step(M, H100, NVLINK4, 16, C, CTX, algo="rd") t_tp16 = max(tp16["t_math"], tp16["t_mem"]) + 2 * M.n_layers * t_ar16 # TP8 × PP2：每台机器 40 层；两个微批各 64 个序列，交替流过两台机器 mb = tp_decode_step(M, H100, NVLINK4, 8, C // 2, CTX, algo="rd") stage = (max(mb["t_math"], mb["t_mem"]) + mb["t_comm"]) / 2 # 半个模型、64 个序列 p2p = IB_NDR.alpha + (C // 2) * M.d_model * 2 / IB_NDR.bw t_pp = 2 * stage + 2 * p2p # 一个 token 依次经过两台机器 # 两个副本：每台机器一个完整的 TP8，各服务 64 个序列 rep = tp_decode_step(M, H100, NVLINK4, 8, C // 2, CTX, algo="rd") w_per_gpu = {"tp8": M.n_params() * 2 / 8, "tp16": M.n_params() * 2 / 16, "pp": M.n_params() * 2 / 16, "rep": M.n_params() * 2 / 8} rows = [ ["一台机器 TP8 扛全部 128 个（参照）", 8, fmt_time(tp8_full["t_step"]), f"{1 / tp8_full['t_step']:.0f}", f"{C / tp8_full['t_step']:,.0f}", f"{w_per_gpu['tp8'] / 1e9:.1f} GB"], ["TP16（每层 all-reduce 跨机器）", 16, fmt_time(t_tp16), f"{1 / t_tp16:.0f}", f"{C / t_tp16:,.0f}", f"{w_per_gpu['tp16'] / 1e9:.1f} GB"], ["TP8 × PP2（两个 64 序列的微批）", 16, fmt_time(t_pp), f"{1 / t_pp:.0f}", f"{C / t_pp:,.0f}", f"{w_per_gpu['pp'] / 1e9:.1f} GB"], ["两个 TP8 副本（各 64 个序列）", 16, fmt_time(rep["t_step"]), f"{1 / rep['t_step']:.0f}", f"{C / rep['t_step']:,.0f}", f"{w_per_gpu['rep'] / 1e9:.1f} GB"], ] print(f"跨机器一次 all-reduce（{C * M.d_model * 2 / 1e6:.2f} MB）：{fmt_time(t_ar16)}；" f"同样大小在机器内（递归倍增）：{fmt_time(tp8_full['t_comm'] / (2 * M.n_layers))}；上下文 {CTX}") print_table(["方案", "卡数", "每 token 时间", "tok/s/用户", "总吞吐 tok/s", "每块卡的权重"], rows) print_observations([ "TP16 让每块卡读的权重和 KV 再少一半，可每层两次 all-reduce 都要走网卡（每卡 50 GB/s、α = 5 µs）：" "160 次跨机器 all-reduce 一共约 8 ms，把省下的全吃掉了：比两个副本慢得多，甚至比一台机器扛全部 128 个序列还慢。**张量并行最好留在一个 NVLink 域里**。", "TP8 × PP2 和两个副本的延迟、吞吐几乎一样：流水线里两个微批交替流过两台机器，每台机器每次处理 64 个序列、半个模型——" "和'两台机器各自完整地处理 64 个序列'的工作量相同。机器之间只传一个 [64, d] 的激活（约 1 MB），很便宜。", "两者的区别在**显存**：流水线每块卡只存半个模型（8.8 GB vs 17.6 GB 的权重），省下的空间可以放更多 KV。" "所以流水线并行的用处是模型一台机器放不下、或者想把显存留给 KV 的时候；放得下时，多开副本最简单（没有任何跨机器依赖）。", "How to Scale Your Model 第 7 章：decode 时'纯数据并行不帮你更快地读参数，不如多起几个模型副本'。" "专家并行（MoE，第 3 讲 §7.6）是另一种切法：把专家分到不同机器上，用 all-to-all 换'每个专家的有效批'（第 5 讲）。", ]) if __name__ == "__main__": demo1_tp_degree() demo2_frontier() demo3_ici_bound() demo4_cross_node()