AI Infra 专题 · 异构训推

训推分离的异构栈:省下的硬件成本,去哪了?

越来越多团队用一套成熟的硬件和软件栈做训练,用另一套更便宜的硬件做推理。硬件采购的账很好算,但算子、编译器、数值精度、通信和工具链上的成本是隐性的。这篇整理了这类方案在算子和 Infra 层会碰到的典型问题,以及业界目前的缓解办法。

01为什么要训推分离

训练和推理对硬件的要求本来就不一样,分开选型有它的道理:

训练:算力密集 + 大规模互联

  • 前向 + 反向 + 优化器,矩阵乘占主导,偏 compute-bound
  • 需要数千卡长时间稳定运行,对集合通信带宽、容错和 checkpoint 要求高
  • 模型架构迭代快,需要成熟的算子库和调试工具

推理:显存带宽 + 成本敏感

  • Decode 阶段每步只生成一个 token,偏 memory-bound,KV cache 占大量显存
  • 负载随流量波动,单位 token 成本是核心指标
  • 推理占总算力消耗的比例在持续上升,长期看它才是大头

所以思路很自然:训练继续用生态最成熟的平台,推理换到单价更低、显存带宽性价比更高的平台。问题在于,两套栈之间的"翻译成本"往往被低估了。

02两套栈:逐层对比

把一个模型从训练平台 A 搬到推理平台 B,要过的不只是"权重格式转换"这一关。下面是从上到下每一层的典型迁移难度:

训练平台 A(成熟生态)

模型定义PyTorch / JAX 模型代码
框架与并行Megatron / FSDP / DeepSpeed 等
图编译torch.compile / XLA
算子库厂商 BLAS、FlashAttention、各种手写 fused kernel
低精度格式BF16 / FP8 / 块缩放格式
集合通信厂商通信库 + 专有互联拓扑
Profiling / 调试成熟的 trace、性能分析工具
基本可复用
需适配
需适配
大量重写
语义差异
重新设计
工具断层

推理平台 B(低成本硬件)

模型定义同一份结构,但常需改写以适配推理框架
推理引擎vLLM / SGLang 或厂商自研引擎
图编译厂商编译器,算子融合规则不同
算子库覆盖面窄,新算子需自己写
低精度格式支持的格式 / 溢出语义 / 累加精度可能不同
集合通信不同的通信库和拓扑,EP/TP 策略需重调
Profiling / 调试工具链薄弱,定位问题慢
基本可复用 需要适配,工作量可控 需要大量重写或重新设计

注:这是定性判断,具体难度取决于两个平台的生态差距。差距越大,红色的层越多。

03什么是"倒挂"

"倒挂"指的是两条曲线方向相反:对算子性能要求最苛刻的环节,偏偏落在软件生态最薄的平台上。

软件生态成熟度

训练平台 A
推理平台 B

对单个算子极致性能的依赖

训练
推理

示意图,长度只表示相对方向,不是测量数据。

为什么推理反而更依赖算子?

结果 硬件单价省下来的钱,很大一部分要用在推理平台的算子工程、精度对齐和性能调优上。如果算子达不到硬件理论性能的合理比例,"便宜的卡"折算到每个 token 上可能并不便宜。

04问题一:算子覆盖与滞后

新的模型架构几乎总是先在训练平台上实现。研究员用现成的算子库搭起新结构,训练完成后,推理平台才开始补对应的算子。

典型的"长尾"算子

Attention 变体

MQA / GQA / 低秩压缩 KV、稀疏或滑窗 attention、paged attention、chunked prefill 的各种组合。

MoE 相关

Top-k 路由、token 重排(permute / unpermute)、Grouped GEMM、带负载均衡的专家调度。

低精度 GEMM

FP8 / 块缩放格式的矩阵乘,以及配套的量化、反量化和 scale 计算 kernel。

融合算子

RMSNorm + 量化、RoPE + KV 写入、激活 + 门控等,融合与否对 decode 延迟影响很大。

采样与解码

Top-p / top-k 采样、投机解码的验证 kernel、结构化输出的约束解码。

新位置编码 / 归一化

每一个"小改动"都可能需要一个新 kernel,否则就退化为多个小算子拼接。

新模型上线的时间线

T0:训练平台上模型完成所有算子都有现成实现或已在训练中调好
T1:推理平台"能跑"缺失的算子用通用实现拼起来,功能正确,但性能可能只有目标的一部分
T2:精度对齐逐层对比输出,排查格式、累加顺序、融合带来的偏差
T3:推理平台"跑得快"关键算子手写优化、融合、调优,达到可以替代原方案的单位成本

T0 到 T3 之间的这段时间就是 day-0 支持的缺口。这期间要么推迟上线,要么继续用训练平台的硬件做推理,"分离"省下的成本也就打了折扣。

05问题二:能跑 ≠ 跑得快(性能可移植性)

Triton、MLIR 这类工具解决了很大一部分功能可移植性:同一份 kernel 代码可以编译到不同硬件。但性能可移植性是另一回事。

公开数据点说明
GEMM 约 95%Triton-distributed 论文中,Triton 生成的 GEMM 在 GPU 上能达到 cuBLAS / CUTLASS 约 95% 的性能。
94–96%ML-Triton 在 Intel GPU 上,GEMM 和访存密集场景达到专家手写 XeTLA kernel 的 94–96%。
复杂 kernel 需重写Meta 在自研加速器 MTIA 上用 Triton,认为可移植性是最大收益之一;但对于复杂 kernel,只靠编译器优化做不到目标性能,还是要针对硬件重写。
调优成本Autotuning 是关键,但开销大,实践中用得不充分;不同平台上同一份 kernel 的表现并不一致。

原因很直接:决定性能的是内存层次、tile 尺寸、片上存储容量、矩阵单元形状、异步拷贝能力这些硬件细节。这些差异不会因为换了一层抽象就消失。主流编译器的成熟度通常也最贴近最主流的硬件,其他后端往往落后一截。

06问题三:数值一致性

同样的权重、同样的输入,在两套平台上算出来的结果不会逐 bit 相同。大部分时候偏差很小,但要证明"很小"本身就是一项工程。

偏差从哪里来

低精度格式的语义差异

  • FP8 E4M3 溢出时,有的实现饱和到最大有限值,有的变成 NaN
  • 同叫 "E4M3",也有 F8E4M3 和 F8E4M3FN 两种不同语义(是否有 inf、NaN 编码)
  • 块缩放格式的块大小、scale 类型不同:OCP MX 用 32 元素块 + E8M0 scale,有些私有 4-bit 格式用 16 元素块 + FP8 scale,后者还没有开放的跨厂商规范(详见下一节对比表)

计算顺序与累加精度

  • 浮点加法不满足结合律,矩阵乘的分块方式、split-K 策略不同,结果就不同
  • MX 规范把点积的内部精度和运算顺序留给实现决定,只建议输出 FP32
  • 算子融合改变中间结果的存储精度(如融合后中间值不再舍入到 BF16)

Batch 不变性

  • 同一条请求在不同 batch 大小下,kernel 会选不同的并行和归约策略,结果也就不同
  • Thinking Machines 的研究指出,这才是推理不确定性的主要来源,而不是常说的"并发 + 原子操作"
  • Batch-invariant kernel 能让同一平台上结果可复现,但做不到跨平台逐 bit 一致

MoE 的放大效应

  • 路由是离散决策:logit 上一个很小的偏差就可能让 token 选到另一个专家
  • 一旦选错专家,后续计算就完全是另一条路径,偏差不再是"很小"
  • 层数越深、专家越多,这种跳变越容易出现
工程上的含义 精度验收不能只看最终 benchmark 分数。需要逐层输出对比、路由一致率、logprob 分布差异等更细的指标,还要定义清楚"多大的偏差可以接受"。这套验证体系本身就要持续投入。

07量化格式对比

低精度格式是训推跨平台最容易"踩雷"的地方:训练侧用的格式,推理平台不一定原生支持;就算名字一样,细节语义也可能不同。下表按家族整理了常见格式。

格式元素编码缩放粒度Scale 类型等效位宽精度表现跨平台风险点
高精度基线
BF16E8M7无—16训练基线低 基本所有平台都支持;但融合 kernel 内部是否保留 FP32 中间值各家不同
FP16E5M10无—16精度高、范围窄中 BF16 训练的模型转 FP16 推理,激活值可能溢出
FP8 家族
FP8 E4M3
(per-tensor)
E4M3整个张量FP32≈8接近无损(常用于推理)高 存在 E4M3 / E4M3FN 两种语义;溢出可能饱和也可能变 NaN
FP8 E5M2E5M2整个张量FP32≈8范围大、精度低,多用于梯度中 推理一般不用;有 inf 编码,处理方式要对齐
FP8 细粒度块缩放E4M3如激活 1×128、权重 128×128FP32≈8比 per-tensor 更抗离群值高 需要 kernel 支持块级反量化;目标平台若只支持 per-tensor,就得重新量化
MXFP8E4M3 或 E5M232 元素块E8M0(2 的幂)8.25接近 FP8 块缩放中 开放规范(OCP MX),但内部累加精度和顺序由实现决定
6 / 4-bit 浮点
MXFP6E2M3 或 E3M232 元素块E8M06.25介于 FP8 和 FP4 之间中 硬件支持面较窄
MXFP4E2M132 元素块E8M04.25有可感知损失,略逊于 NVFP4中 开放规范;2 的幂 scale 只能整倍缩放,常"缩过头"
NVFP4E2M116 元素块 + 张量级E4M3 + FP324.54-bit 中较好;厂商报告相对 FP8 损失约 1% 以内高 非开放跨厂商规范,其他平台通常需要转换或重新量化
整数量化
INT8 W8A8INT8权重 per-channel、激活 per-tokenFP16 / FP32≈8成熟,需处理激活离群值低 支持最广;但 SmoothQuant 类预处理要在部署链路里复现
INT4 权重量化
(GPTQ / AWQ 类)
INT4每组 64 / 128 个权重FP16 scale + 零点≈4.1–4.25decode 访存友好,精度依赖校准数据中 依赖反量化 kernel;权重打包布局各家不同

等效位宽 = 元素位宽 + scale 位宽 ÷ 块大小,例如 MXFP4 = 4 + 8/32 = 4.25,NVFP4 = 4 + 8/16 = 4.5(未计张量级 scale)。"精度表现"是公开资料中的定性结论,具体数值随模型和任务变化。

跨平台迁移时的几个坑

格式不支持 → 重新量化

目标平台没有对应格式时,只能从高精度权重重新量化到它支持的格式。如果手上只有量化后的权重,就会发生"二次舍入",误差叠加。建议始终保留 BF16 主权重。

Scale 之间不能随便互转

E8M0 是纯 2 的幂,E4M3 带尾数但范围小得多(大约 2-9 到 448)。两者互转一般都有损失。块大小 16 和 32 之间转换也不对称:细粒度转粗粒度一定有损失。

矩阵单元的累加精度

一些硬件的 FP8 矩阵单元内部累加精度不到完整 FP32。长 K 维累加时误差会积累,需要分段把部分和提升到 FP32 再累加。不同平台在这方面的行为差异很大。

KV cache 也在量化

推理侧经常把 KV cache 量化到 FP8 或 INT8 来省显存。训练时从来没有这一步,所以它是推理独有的误差来源。在 PD 分离部署中,两侧的 KV 格式还必须一致(见后文)。

选格式的经验 跨平台部署时,优先选开放规范、两边都原生支持的格式,哪怕理论精度不是最优。私有格式在单一平台上效果更好,但换平台时的转换成本和精度风险要另算。

08问题四:通信与并行策略

大模型推理很少是单卡的事。张量并行、专家并行、Prefill/Decode 分离部署都依赖高效的通信。

维度迁移时的问题
集合通信库AllReduce、All-to-All 的实现和性能特征不同;训练侧针对特定拓扑调好的通信与计算重叠方案无法直接搬过来。
互联拓扑单节点卡数、节点内带宽、跨节点网络结构不同,最优的 TP / EP 切分方式随之改变。
MoE All-to-All专家并行的 token 分发对延迟极其敏感,通常要专门写 dispatch / combine kernel,并和计算做细粒度重叠。
PD 分离Prefill 和 Decode 分开部署时要在节点间传 KV cache,传输方式、布局都要适配新硬件(详见下一节)。
数值影响AllReduce 的归约顺序本身也是数值差异的来源之一,并行策略变了,结果也可能变。

09Prefill / Decode 分离部署

PD 分离(disaggregation)是近两年大模型推理的主流架构之一。它和"训推分离"是同一个思路,只是分得更细:连推理本身也拆成两段,分别放到最合适的硬件上。这也让它成为异构部署最自然的切入点,同时带来新的 Infra 问题。

为什么要拆

Prefill:一次性处理整个 prompt

  • 大矩阵乘,偏 compute-bound
  • 决定首 token 延迟(TTFT)
  • 长 prompt 的一次 prefill 可能占用 GPU 几百毫秒

Decode:逐 token 生成

  • 每步读全部权重和 KV cache,偏 memory-bound
  • 决定每 token 延迟(TPOT)
  • 算力利用率很低,靠大 batch 提升吞吐

两者混在同一组卡上时会互相干扰:一个长 prefill 插进来,正在 decode 的请求就会卡顿。拆开之后,两个资源池可以分别扩缩容、分别选并行策略。DistServe、Splitwise、Mooncake 等系统都报告了显著的 goodput 提升。

基本架构

请求prompt
→
调度器选 P 实例 + D 实例
考虑前缀命中、负载、拓扑
→
Prefill 池算力密集型硬件
生成首 token + KV
⇢
KV 传输RDMA / 设备直连
逐层流式发送
→
Decode 池带宽 / 显存性价比高的硬件
持续生成

可选:在 CPU 内存 / SSD 上建一层共享 KV 缓存池(Mooncake 的做法),前缀命中时可以跳过重算

异构 PD 带来的额外问题

如果 Prefill 池和 Decode 池用的是不同的硬件栈,一个请求的生命周期就跨了两套平台。前面讲的所有训推差异,都会在单个请求内部重演一遍:

问题说明
KV 布局Page 大小、head 排列、K/V 是否分开存、张量内存布局,两边的 attention kernel 可能各有要求。不一致就要在传输时重排,增加延迟。
KV 精度格式一侧用 FP8 KV、另一侧只支持 BF16 或 INT8 时,需要在传输链路上转换,引入额外误差和开销。
数值混合KV 由 A 平台的 kernel 计算,再由 B 平台的 attention kernel 消费。同一个请求的输出同时依赖两套数值行为,复现和排查都更困难。
传输通道同一厂商内常有设备直连或 GPU-Direct RDMA;跨厂商时往往只能走通用网络,甚至要经主机内存中转,带宽和延迟都打折扣。
P:D 配比最优的 Prefill / Decode 实例比例随 prompt 长度和流量变化。两池硬件不同时,不能互相借用资源,只能各自冗余。
调度与拓扑KV 传输占 TTFT 的比例可观,调度时要考虑网络拓扑和链路拥塞(TopKV、NetKV 等工作专门研究这一点)。
模型切分两池的 TP / EP 切分方式不同,KV 的分片方式也不同,传输时需要做 N:M 的重新分发。
运维两套驱动、两套监控、两种故障模式;一侧故障时另一侧的请求也会受影响。

KV 传输量估算

KV cache 的大小直接决定 PD 分离对网络的压力。可以调整下面的参数估算:

—
每 token KV
—
每请求 KV
—
集群总传输带宽需求
—
单请求单链路传输时间

标准 K/V:每 token = 2 × 层数 × KV head 数 × head 维度 × 字节数;压缩潜向量:每 token = 层数 × 潜向量维度 × 字节数。单位按 10 进制(1 GB = 109 字节)。传输时间未计协议开销,也未计逐层流式传输带来的重叠。

以 70B 级 GQA 模型为例:8K token 的 prompt 产生约 2.7 GB 的 KV,100 QPS 时集群需要约 270 GB/s 的传输带宽。KV 结构本身的设计(GQA、低秩压缩)对传输量的影响往往比链路升级更大。

常见的优化手段

逐层流式传输

Prefill 算完第 i 层就开始传第 i 层的 KV,让传输和计算重叠,而不是等全部算完再传。

压缩 KV 结构

GQA、低秩潜向量等结构从模型设计层面减少 KV 体积,对 PD 分离最友好。

KV 量化

FP8 / INT8 KV 直接减半传输量,但两侧格式必须一致,且要验证精度。

前缀缓存池

多轮对话、共享 system prompt 场景下,命中缓存就不用重算、也不用重传。

拓扑感知调度

优先把 P 和 D 放在同一交换机或同一节点组内,减少跨层网络流量。

动态 P:D 配比

根据实时负载调整两池实例数量(如 DOPD),避免一侧空闲一侧排队。

PD 分离不一定划算 短 prompt、低并发或网络带宽不足时,KV 传输开销可能超过分离带来的收益。Chunked prefill(把 prefill 切块,与 decode 混合批处理)是另一条路线,不需要跨节点传 KV。异构 PD 还要再加上双栈适配的成本,最好先用同构 PD 验证收益,再考虑引入第二种硬件。

10问题五:在 RL 训练中被放大

如果只是离线训练、在线推理,数值偏差影响的是推理质量。但在强化学习(RL)后训练里,推理引擎本身就是训练流程的一部分:

Rollout 侧

用推理引擎(vLLM / SGLang 等)大量生成样本,并记录采样时的 logprob。这部分算力消耗很大,正是最想放到便宜硬件上的环节。

Training 侧

用训练框架(FSDP / Megatron 等)重新计算 logprob 并更新策略。PPO 类算法的重要性比率依赖两边 logprob 的一致性。

研究发现,即使在同一种硬件上,推理引擎和训练框架给同一序列算出的 logprob 也会有差异,原因包括 kernel 实现、batch 大小、专家选择逻辑和归约顺序。这会让本该 on-policy 的训练悄悄变成 off-policy,MoE 模型和量化 rollout 的情况更严重。Yao 等人甚至改了 vLLM,让它输出真实采样概率、lm_head 用 FP32 计算,差异仍然存在。

如果 rollout 和 training 再跑在两种不同的硬件上,所有前面提到的差异都会叠加进来。

目前的缓解手段

方法思路代价
系统级对齐让两边所有算子逐 bit 一致,训推 KL 为 0需要 batch-invariant 等特殊 kernel,性能开销大;跨硬件基本做不到
截断重要性采样(TIS)用训推 logprob 比值修正损失,并截断上限引入少量偏差,换取稳定性
掩码重要性采样(MIS)等直接丢弃偏差过大的 token浪费部分样本
Routing replay训练侧复用 rollout 时的 MoE 路由结果需要传递和存储路由信息
量化感知 RL训练时就模拟 rollout 的量化行为训练流程更复杂
TIS 的常见形式:每个 token 的损失乘以 wt = min( exp(logptrain − logprollout), C ),C 常取 2

也有研究指出,在超大规模 MoE 上,这类基于重要性采样的修正只能推迟训练崩溃,不能根除问题。

11隐性成本账

评估训推分离时,除了硬件单价,还应该把下面这些都算进去:

成本项具体内容性质
算子工程补齐缺失算子、手写关键 kernel、持续跟进新架构持续性人力
精度对齐建立逐层对比工具、定义验收标准、排查偏差每个新模型都要做一次
性能调优融合策略、tile 参数、并行切分、通信重叠直接决定单位 token 成本
双栈维护两套推理代码路径、两套 CI、两套监控和故障处理长期运维负担
上线延迟新模型从训完到在推理平台高效上线的时间差机会成本,过渡期可能仍需用原硬件推理
研发约束为了推理平台好支持,模型设计时回避某些算子可能限制架构创新
一个粗略的判断方式 推理规模越大、模型架构越稳定、团队自研算子能力越强,硬件差价越能覆盖这些隐性成本。反过来,如果模型几个月一换代、每次都引入新算子,隐性成本就会反复出现。

12缓解思路

1. 统一中间表示和编译器

用 StableHLO、MLIR、Triton 这类跨硬件的中间层来描述模型和 kernel,让厂商只需要实现后端。StableHLO 有版本化的序列化格式,可以在不同编译器之间传递程序;Triton 把编译拆成硬件无关的前端和硬件相关的后端。

局限:IR 稳定不等于 lowering pipeline 稳定,换版本、换编译器时 pass 流程可能失效;性能关键的 kernel 仍需针对硬件手调。

2. 算子抽象层 + 分级实现

在推理引擎内部定义一层清晰的算子接口,每个算子提供"参考实现 → 通用实现 → 硬件优化实现"三级。新模型先用参考实现保证正确,再逐步替换成优化版本,缩短 T0 到 T1 的时间。

3. 把数值行为写成"合同"

明确规定每种低精度格式的编码、溢出策略、scale 计算方式、累加精度,并准备逐 bit 的一致性测试向量。业界已经有工作在做格式目录和合规测试向量(如 Golden Ruler),可以作为参考。

4. 训练时就考虑推理约束

5. 精度回归进 CI

每次算子或编译器更新都自动跑一遍逐层对比和端到端评测,记录 logprob 差异、MoE 路由一致率等指标,而不是等上线后才发现问题。

6. RL 场景的算法侧兜底

接受训推之间存在偏差,用 TIS、MIS、routing replay 等方法在算法层面修正,同时监控训推 KL 作为健康指标。

13决策清单

在决定把推理迁到另一套硬件之前,可以先回答这几个问题:

  1. 目标模型的所有算子,在推理平台上有没有优化过的实现?缺几个,补齐需要多久?
  2. 推理平台上关键算子(attention、MoE GEMM、量化 GEMM)能达到硬件理论峰值的多少?
  3. 折算成每百万 token 成本(而不是每张卡的价格)之后,还剩多少优势?
  4. 有没有逐层精度对比工具和明确的验收标准?
  5. 模型架构的迭代速度如何?每次换代要重新付出多少适配成本?
  6. 如果涉及 RL,rollout 放在推理平台上时,训推 logprob 差异有多大?算法侧能否兜住?
  7. 团队有没有足够的算子和编译器工程师来长期维护第二套栈?

14参考资料