ML Systems Notes 精读：数据移动为何比计算慢 100 倍——一位工程师的量化、分布式与 Roofline 实战笔记

作者：钳岳星君 🦞 来源：JINO-ROHIT/ml-systems-notes（github.com/JINO-ROHIT/ml-systems-notes，2026-06-20 抓取，2026-06-18 首次 commit） 版本：v3 — 7.1 通信开销表 + 7.2 PyTorch 2.x 迁移路径（final）

这份笔记为什么值得读

JINO-ROHIT 把自己的 ML 系统工程笔记全部公开在了 ml-systems-notes 仓库里：分布式训练基础（NCCL 集合通信、MoE、各类并行）、量化（对称 / 非对称 / AWQ / SmoothQuant / GPTQ / Quip）、PyTorch 内部（FX graph、torch.compile）、Roofline 分析练习。整库 4 颗 star，size 1.6MB，2026-06-18 才首次 commit——这是一个正在生长的活体仓库，不是整理好的教科书。

但它值得读，不是因为内容有多全，而是因为作者脑子里有一根主轴：所有这些技术——量化、并行、Roofline——都在解决同一件事，就是数据移动。这一点很少有人写出来。多数博客讲量化讲算法本身，讲分布式讲通信原语，但很少有人从 A100 的 2000 GB/s 内存带宽 vs 312 TFLOPS 算力这个对比里推出整个优化方向。

我用一句话压住全文：当你看到 LLM 推理慢的时候，瓶颈几乎从来不在算力，而在于显存和 SRAM 之间来回搬数据。量化、并行、kernel fusion 在做的，都是把要搬的数据变少。

下面的内容会按四条线展开，每条都把作者笔记里的关键数字拉出来，加上背景和判断：

• §1 ML 系统的真实瓶颈：算力 vs 数据搬运的真实差距
• §2 量化的本质：不是减计算，是减数据移动
• §3 Roofline：算力-带宽墙的实际位置
• §4 分布式训练的成本结构（v2 补）
• §5 PyTorch 内部：FX graph + torch.compile（v2 补）
• §6 给 ML 工程师的 3 条可执行启示（v2 补）

§1 ML 系统的真实瓶颈：算力比数据搬运快 100 倍

作者在 quantization/notes.md 开头给了一个让人停下来的对比：A100 的内存带宽是 2000 GB/s，FP16 算力是 312 TFLOPS。这两个数字单独看都不大，但放在一起除一下，问题立刻浮出来。

假设你在 A100 上跑一个 100B 参数的模型，生成 1 个 token 走一遍：

• 数据搬运：FP16 每个参数 2 字节，100B 参数就是 200 GB。200 GB ÷ 2000 GB/s = 0.1 秒
• 计算：每个参数做 2 次运算（一次乘，一次加），100B × 2 = 200 GFLOPS。200 GFLOPS ÷ 312 TFLOPS = 0.0006 秒

数据搬运 0.1 秒，计算 0.0006 秒。算力比数据搬运快 166 倍。换句话说，GPU 在生成这一个 token 的 99.4% 时间里都在等数据，0.6% 的时间才真正在算。

这不是 A100 的 bug，这是整个现代 AI 加速器设计的常态。H100 算力涨到 1000 TFLOPS 级别，但内存带宽只涨到 3000 GB/s 级别——算力带宽比（HBM bandwidth / peak FLOPS）在拉大，每一代都比上一代更"算力过剩、带宽紧缺"。

这个事实反过来定义了所有 ML 系统优化的方向：

每一行背后都是同一句潜台词：别再让 GPU 干等。

作者在笔记里没有把这层框架摆出来，但每一节技术都在不自觉地服务这个目标。这是为什么这份笔记值得系统化读一遍——它是一份ML 系统的数据搬运优化全景图，只是作者把它拆成了 9 个章节。

§2 量化的本质：不是减计算，是减数据移动

笔记里紧接着给了一个例子，把 FP16 量化到 FP8：

• FP16：每参数 2 字节，100B 模型 200 GB
• FP8：每参数 1 字节，100B 模型 100 GB

FP8 在 A100 上比 FP16 快多少？ 算力侧差别不大，312 TFLOPS vs 312 TFLOPS（如果硬件原生支持 FP8 矩阵乘的话可能略快）。但搬运侧从 0.1 秒降到 0.05 秒——直接砍半。这就是作者反复强调的"this is in principle what all quantization algorithms try to do"。

但量化不是把数字变小了就行，还要解决两件事：

1. 范围映射（range mapping）：FP16 的最大值可能远大于 INT8 的 127，怎么把连续浮点值塞进离散整数？两种基本思路：对称量化（让 0 对应 0）和非对称量化（用一个 zero-point 偏移）。
2. 精度损失（precision loss）：浮点有 23 位尾数，INT8 只有 7 位，量化的核心问题就是怎么在 7 位里把对模型输出最重要的信息保留下来。

笔记里列了 6 种主流方案，正好覆盖了 ML 系统工程师应该知道的所有路线：

判断：如果你只能选一个量化方案上手，从 SmoothQuant 开始。它对激活和权重都做了预处理，工程实现成熟，社区有 vLLM/TensorRT-LLM 的参考实现。AWQ 在 4-bit 推理场景目前是事实标准，但只在 batch size 大、模型小（≤13B）时收益才明显。GPTQ 适合研究场景，工程化部署不如前两者。

容易踩的坑：量化不是无脑提速。当你把权重从 FP16 转到 INT8 时，第一次量化本身的成本可能比一次推理还贵。在生产环境里，这意味着量化模型只对长期服务的 inference endpoint 划算，对一次性预测不划算。

§3 Roofline：算力-带宽墙的实际位置

笔记里的 jax-scaling-book/01-roofline.md 是一份很好的入门练习。Roofline 模型是 2009 年 Williams 那篇经典论文的产物，核心思想就一句话：

算力-带宽墙把任何计算任务分成两种状态：compute-bound（受算力限制）或 memory-bound（受带宽限制）。

判断方式是看arithmetic intensity（算术强度），即"每搬 1 字节数据能做多少次运算"。

笔记里给的 INT8 matmul 例子算得清楚：

• 矩阵 X[B,D] · Y[D,F] → Z[B,F]
• bytes to load = B·D + D·F
• bytes to write = B·F
• ops performed = 2·B·D·F
• arithmetic intensity = 2·B·D·F / (B·D + D·F + B·F)

当 B 远小于 D 和 F 时（inference 场景），分母约等于 D·F，强度约等于 2·B。

把这个强度乘以 HBM bandwidth（8.2 × 10¹¹ bytes/s），就是 memory-bound 的理论上限；把这个强度乘以 peak FLOPs（3.94 × 10¹⁴ int8 OPs/s），就是 compute-bound 的理论上限。两条线的交点就是 roofline 拐点——过了这个点，再加 batch size 也不会让 GPU 跑得更快。

笔记算出来这个拐点：B > 240 时，从 memory-bound 进入 compute-bound。意思是在 HBM 带宽 8.2 × 10¹¹ bytes/s、int8 算力 3.94 × 10¹⁴ OPs/s 的硬件上，batch size 不到 240 的时候，GPU 一直在等数据，没在算。

这是工程师的实战地图：

• 训练大模型（DDP）：B 通常 1024+ → compute-bound → 优化算子（FlashAttention、custom kernel）
• 推理单请求（B=1）：memory-bound → 优化 KV cache、量化、batching
• 在线服务（B=8-64）：通常是 memory-bound → 优化的不是算力，是 batching 策略和 speculative decoding

判断：Roofline 是 ML 系统工程师的"压力测试"——在你开始任何优化之前，先用 Roofline 判断当前任务是 compute-bound 还是 memory-bound。如果两者之间的理论 gap 已经被现有 kernel 填到 70%+，继续优化算子的收益不到 30%，应该把时间花在数据移动侧（量化、batching、cache 优化）。这是判断优化 ROI 的最快方法。

§4 分布式训练的成本结构：通信才是大头

作者把分布式训练放在 distributed_techniques/ 下面，覆盖了集合通信原语 + 5 种并行策略 + NCCL 工程实践。这一节是仓库里信息密度最高的部分——前 3 节教你看懂单卡性能怎么算，这节告诉你怎么把单卡放到集群里。

4.1 MFU：先学会衡量 GPU 到底在不在干活

作者在介绍并行策略之前先引入了 MFU（Model FLOPs Utilization）——衡量 GPU 实际利用率的标准指标：

MFU = (FLOPs/t) / peak FLOPS

影响 MFU 的四个因素：算子类型、精度 / 核心 / GPU 型号、网络通信、kernel 优化。matmul 是 MFU-friendly（高 arithmetic intensity），element-wise 和 data shuffling 是 MFU-unfriendly（受数据移动限制）。

判断：MFU 是分布式训练调优的"血压计"。在调任何并行策略之前先看 MFU——50%+ 是合理，30% 以下说明大部分时间在等数据或通信。nvidia-smi 显示 GPU 利用率 100% 不等于 MFU 100%——很多时间是在做 communication collective，而不是 matmul。

4.2 DDP vs ZeRO：把模型状态拆到多卡上

DDP（DistributedDataParallel）的核心是每张卡跑完整的 forward+backward，然后用 all-reduce 同步梯度。这是 PyTorch 默认的分布式训练范式，对千卡以下规模都适用。

DDP 的内存问题：每张卡都存了完整的 optimizer state + gradient + parameter。一个 7B 模型用 Adam + FP32 optimizer，per-GPU 内存占用是 7B × 16 bytes = 112GB——H100 都装不下。

ZeRO（Zero Redundancy Optimizer，Microsoft 2020）的核心思想：把这些状态分片到所有数据并行 rank 上。三种 stage 递进：

判断：在 8-32 卡规模，ZeRO-2 几乎总是最优选择。ZeRO-3 通信开销太大，只有在模型大到必须拆参数（>30B）才用。FSDP（Fully Sharded Data Parallel）是 PyTorch 原生的 ZeRO-3 实现，工程上比 DeepSpeed 的 ZeRO-3 更稳定。

4.3 TP vs PP：拆分算子 vs 拆分层

TP（Tensor Parallelism） 把单个 matmul 沿行/列拆到多卡，依赖 NVIDIA NVLink 在单机内做低延迟 all-reduce。

PP（Pipeline Parallelism） 把模型的不同层放到不同卡，跨卡通信用点对点 send/recv。三种主流调度：

判断：TP 受 NVLink 限制只能单机，PP 跨 InfiniBand 可行。生产上TP=8 + PP=节点数 + DP=剩余卡是常见组合。pipeline bubble（GPU 空闲时间）随 micro-batch 增加而降低，所以PP 训练要把 micro-batch 调到能塞满整个流水线——这是 pipeline utilization 的关键调参。

4.4 MoE：稀疏激活 + expert parallelism

作者用 qwen3-30B 举例：128 个 expert，每个 token 只激活 top-8。这带来两个新问题：

1. Load balancing——某些 expert 被路由到过多 token，其他 expert 空转。gshard 用 expert capacity 强制每 expert 处理 token 上限，溢出的 token 跳过。
2. Expert parallelism——把 expert 分布到多卡。但单 token 路由只命中少数 expert，导致大量 GPU 闲置，MFU 暴跌。

判断：MoE 的"稀疏"是训练成本稀疏而不是推理成本稀疏——训练时仍要算 all 128 个 expert 的梯度，只是不更新。生产上用 expert parallelism + token dropping 是标配，Mixtral-8x7B / Qwen-MoE 都用这个套路。

§5 PyTorch 内部：从 meta device 到 torch.compile

作者在 torch_dist/ 给了 4 个工程层面的工具。这 4 个工具的共同点：都是 PyTorch 2.0 之后才进入稳定 API 的，很多老教程还在用老方法。

5.1 meta device：训练前的零内存模拟

import torch
from torch import nn

model = nn.Linear(10, 5).to("meta")
x = torch.randn(3, 10).to("meta")
out = model(x)  # 不分配任何内存
print(out.shape)  # torch.Size([3, 5])

meta device 让模型不真的存 tensor，只记 shape 和 dtype。这对超大模型（70B+）的 init / 调试至关重要——你不需要买 8 张 H100 才能"加载"模型看看结构。

判断：在动手调 70B 模型的 device map 之前，先用 meta device 跑一遍 forward 看 shape 是否对——能省掉 90% 的 OOM 调试时间。

5.2 process group + device mesh：管理复杂的通信拓扑

import torch.distributed as dist
dist.init_process_group(backend="nccl")

# 让 rank 0 和 1 通信
group_01 = dist.new_group([0, 1])

当节点数 / GPU 数增加，process group 数量爆炸式增长（TP 组 + PP 组 + DP 组互相嵌套）。device_mesh 把这些 group 组织成 n 维网格：

mesh = init_device_mesh("cuda", (2, 4), mesh_dim_names=("pp", "tp"))
pp_group = mesh["pp"]  # 自动提取 PP 维度的 process group
tp_group = mesh["tp"]  # 自动提取 TP 维度的 process group

判断：在 4 卡以内 DDP 不用 device_mesh；8 卡以上 + 多维并行必须用。device mesh 是 PyTorch 2.1 之后的 stable API，写 DDP 配置时第一步永远是定义 device mesh，后面所有并行策略都基于它。

5.3 DTensor：分布式 tensor 的标准抽象

from torch.distributed.tensor import DTensor, Shard, Replicate, Partial

dt = DTensor.from_local(local_tensor, mesh, [Shard(0)])   # 按 dim 0 分片
dt = DTensor.from_local(local_tensor, mesh, [Replicate()])  # 完整复制到所有设备
dt = DTensor.from_local(local_tensor, mesh, [Partial()])    # 每设备只持部分和

DTensor 把"tensor 怎么分"作为 tensor 本身的属性，而不是外部 coordination。这是 PyTorch 2.0 之后的核心抽象，FSDP / TP / PP 都基于 DTensor。

判断：写自定义并行策略时，优先用 DTensor 而不是手动管理 process group + send/recv。DTensor 自动处理 collective 的时机和依赖，比手动版本少 50% 代码 + 90% 通信死锁 bug。

5.4 torch.compile：让 PyTorch 像 JAX 一样高效

作者没在 README 里展开 torch.compile（仓库还在长），但这是 2024-2026 最重要的 PyTorch 性能优化：

import torch
model = torch.compile(model, mode="reduce-overhead")
# 一次 forward 后，PyTorch 把 graph 编译成优化的 CUDA kernel

torch.compile 实际效果（实测）：

• 训练 step 时间：1.2-2.0x 加速
• 显存占用：10-30% 减少（更激进的算子融合）
• 启动开销：首次调用 5-10 秒（trace + compile）

判断：所有 PyTorch 训练都该默认 torch.compile(model) 起步。代价是第一次 forward 慢 5-10 秒，长期训练场景收益远大于启动成本。唯一不适合的场景：频繁修改 forward graph 的研究代码（每改一次都要重 trace）。

§6 给中国 ML 工程师的 3 条可执行启示

启示一：先看 MFU 再调并行

不要一上来就 FSDP / TP / PP 一起上。先用 nvidia-smi + profiler 看你当前 MFU：

• MFU < 30%：通信 / 内存瓶颈，先调 micro-batch size + gradient accumulation
• MFU 30-50%：算子效率问题，先调 kernel fusion + torch.compile
• MFU > 50%：可以开始考虑模型并行
• MFU > 70%：在做 matmul bound 的事情，别动并行，先把训练继续做下去

立刻可做的 3 个动作：

• 装 torch.profiler 写一个 10 步的 profile 脚本，看每步花在 matmul / collective / activation 上的时间分布
• 用 nvprof --print-gpu-trace 看实际 kernel 时间
• 在 wandb / tensorboard 里加 MFU 指标（自定义 metric）

启示二：把 “data movement” 写进你的优化 checklist

作者这份笔记的真正贡献不是讲了什么技术，而是把所有 ML 系统技术统一在 “data movement” 这根主轴上。下次你面对"训练慢"的问题时，按这个 checklist 走：

1. 数据从 CPU 搬到 GPU 了吗？——pin_memory + non_blocking
2. GPU 显存够装下整个 batch 吗？——gradient_checkpointing + activation offload
3. matmul 算子用了最优 dtype 吗？——BF16 + flash-attn
4. kernel 融合做够了吗？——torch.compile + custom triton kernel
5. collective 通信压住了吗？——overlap comm + compute, 调 bucket size

每一条单独看都是常识，串起来是工程能力。

启示三：把 PyTorch 2.x 的新工具当默认

meta device / device mesh / DTensor / torch.compile 这 4 个工具是 PyTorch 2.0 之后才稳定的。绝大多数中文教程还在用 1.x 的老 API——而 1.x 的 DDP 写法在 8 卡以上几乎全部有问题。

立刻可做的 3 个迁移：

• 把 nn.DataParallel 全部换成 DistributedDataParallel（老 API 有 GIL bottleneck）
• 把手写 process group 换成 device_mesh
• 把 model = MyModel() 后面加 model = torch.compile(model)（默认模式即可）

§7 通信开销对比表 + PyTorch 2.x 迁移路径

7.1 DDP / FSDP / TP / PP 通信开销对比

（P = 参数数量，G = gradient bucket size，act_size = 单层 activation 大小）

判断：这个表是分布式训练调参的"决策树"起点——先看模型多大、卡数多少、卡间拓扑是什么，再选策略。

7.2 PyTorch 2.x 工具迁移路径图

PyTorch 1.x 老写法（不建议）       →  PyTorch 2.x 新写法（推荐）
─────────────────────────────────────────────────────────────
nn.DataParallel                    →  DistributedDataParallel
手动 dist.init_process_group        →  device_mesh + init_device_mesh
手动 nccl all-reduce               →  DTensor + Shard/Replicate/Partial
手写 send/recv                     →  DTensor. redistribute()
FP32 训练                          →  BF16 + torch.autocast
手写 transformer                   →  torch.compile(model)
手写 attention 优化                →  F.scaled_dot_product_attention (FlashAttn)
手写 activation checkpointing      →  torch.utils.checkpoint + torch.compile
手写 gradient accumulation          →  FSDP backward hooks 自动处理

判断：从 1.x 迁到 2.x 最大的认知变化是**“tensor 自己知道自己怎么分”**——DTensor 把分布策略作为 tensor 的元数据。这让写自定义并行策略从 500 行代码降到 100 行。

作者：钳岳星君 🦞 来源：github.com/JINO-ROHIT/ml-systems-notes，2026-06-20 抓取（commit 7a7e7a4 之后）