Meta 把推荐性能拆到子模块，83ms 的编译机会估成 50ms

Component Benchmark: Hierarchical Model Profiling for Large-scale Recommendation Systems

Dharak Kharod, Yuzhen Huang, Zhou Wang, Jackie Xu, Fuzail Khan, Jacky Zhou, Hao Yan, Lidong Zhao, Xizhou Feng, Yvonne Liu, Karthik Jayaraman, Praveen Ramachandran, Vishwa Karia, Yashasvi Makin

cs.IR

2026-09-25

Meta 用 hook 把每个子模块连同真实输入单独 profile；生产里把整模 PT2 机会估成 83ms 到 50ms 并换来 34% QPS，一次 2.8 倍显存回归也拆到两个模块。

这篇在解决什么

推荐模型还没收敛到 Transformer 一种骨架。DLRM、DCN、DHEN、HSTU 混着 embedding 查表、小尺寸 GEMM、jagged 特征和低算术强度算子。生产模型是 TB 级，跑在数千 GPU 上，每天吃大约 10¹¹ 条样本。个位数的吞吐提升会变成机房容量。

现有工具答错粒度。Nsight 和 Kineto 给 kernel 轨迹。PyTorch Profiler 的 withmodules 在整网前向里归因，模块之间的 cache、融合、调度叠在一起；PyTorch 2 编译还会抹掉模块边界。ML 工程师问的是「哪个 nn.Module 拖后腿」，工具答的是「这条 kernel」。

方法

Component Benchmark（CB）核心大约 1000 行，PyTorch 原生。hook 在选定边界抓住子模块的运行时输入，按广度优先把每个模块单独跑一遍。父模块塞不进单卡就跳过，用子模块结果聚合。

插件分四层：

Provider：从训练栈抽出模型和真实 batch。embedding 表先缩小以塞进单 GPU，捕获输入后再复制回原 batch 大小，降低抽取阶段 OOM。
Preprocessor：混合精度、batch 调整。
Profiler：FLOP、显存快照、Kineto、kernel 分析。
Result：冰柱图加表，按模块路径给出 latency、模型体积、TFLOPS、MFU、带宽。

演示在一张 H100 上，包了开源 DeepFM-DLRM 和 NanoGPT。生产模型通过自定义 RecSysProvider 接入。

结果

开源 DLRM 上，冰柱图直接读出异构：sparsearch 几乎纯带宽；interarch 里 deepfm 的 MFU 为 27.06%，做 FM 交叉项的兄弟模块只有 0.06%。dense MLP 里两层 Linear 的 MFU 是 25.91% 对 20.74%，形状分别是 1024×4096 和 4096×512。ReLU 报告 0 FLOP，颜色整段红。

NanoGPT 上 12 个 block 几乎一样。同一 block 里注意力 MFU 大约 10%，MLP 大约 20%，layer-norm / GeLU / dropout 接近 0。框架不绑推荐结构。

生产故事没有对照实验设计，数字来自使用记录：

若干低 MFU 的 GEMV 改写成 elementwise。单 kernel 更慢，PT2 融合区变大，训练加服务吞吐升 3%。整网轨迹能看到「慢」，测不了「把一块变慢让整网变快」。
当时只编译了 overarch。隔离测整模开/关 PT2，迭代时间从 83ms 降到 50ms，支撑后来把编译铺开换来的 34% QPS。
生成式排序相对前任吃 2.8 倍显存、FLOP 多 39%。CB 拆到更重的 event model 和放大的 set-transformer；改回尺寸后 QPS 升 10%、显存降 22%。
资源再分配时，轨迹只显示通信和计算都 stall。CB 标出两个低性能子模块，只缩小这两块，算力 stall 基本消失。

为什么重要

性能对话从「看 Perfetto」改成「这个 Linear 的 MFU 是 20%」。推荐这种异构网特别吃这套：同一层里有人在算矩阵，有人在搬 embedding。对 LLM 也能跑，只是同质 block 里信息密度低一些。

这是内部工具论文。开源 DLRM 和 NanoGPT 的数字是演示，不是 SOTA 对比。能信的是生产里那几个归因故事。

局限与存疑

独立 profile 丢掉了模块并行时的重叠与融合，可能高估或低估真实 stall；「把 GEMV 变慢让整网变快」本身就说明隔离数字不是全局最优。单卡、缩小 embedding，和数千 GPU 训练有差距。没有与 PyTorch Profiler withmodules 的系统误差表。生产数字缺实验协议，也没有开源仓库地址。冰柱图在论文静态图里几乎不可读，交互界面才是产品。

术语

MFU：Model FLOPs Utilization，实际 FLOP 吞吐相对硬件峰值的比例，低通常表示带宽或启动开销卡住。
jagged 特征：每个样本类别特征长度不同，动态 shape 让编译和 profile 都变难。
冰柱图：把模块树画成从上往下的嵌套条，颜色映射 MFU 或耗时。
PT2：PyTorch 2 编译栈，会融合算子并抹掉模块边界。

原文与代码

社区讨论

Meta 开源推荐系统分析工具,按子模块剖析 TB 级模型性能 _reachsumit