第6章:训练侧远程内存与参数池化 —— Offload、Checkpoint、PMEM Hybrid

训练侧的”远程内存”看起来朴素——把装不下的 optimizer state 倒去 CPU DDR、把 checkpoint 写到 NVMe——但底下有同一套 disaggregation 哲学:让物理位置和逻辑位置解耦,用低延迟互联缝合。本章梳理三条工程线:ZeRO-Infinity 的 HBM/DDR/NVMe 三层 offload、HugeCTR 推荐系统的 GPU+PMEM hybrid embedding 表、RDMA 加速的 checkpoint 写盘。读完这章你会知道为什么 ZeRO-Infinity 能训练 32T 参数模型,以及 PS 架构在大模型时代的位置。

📑 目录

1. 训练侧”内存不够”的三种表现
2. ZeRO 系列演化:从 ZeRO-1 到 ZeRO-Infinity
3. ZeRO-Infinity 的三层 offload:HBM / DDR / NVMe
4. 推荐系统:HugeCTR 与 PMEM Hybrid Embedding
5. Checkpoint 远端写:RDMA-based 训练 fault tolerance
6. Parameter Server 在大模型时代:还没死
7. 训练侧 disaggregation 的局限
自我检验清单
参考资料

1. 训练侧”内存不够”的三种表现

训练 70B 参数模型的内存账(混合精度 + Adam 优化器,简化估算):

状态	单参数字节	70B 总和
Model parameters(fp16)	2 B	140 GB
Gradients(fp16)	2 B	140 GB
Optimizer state(Adam: m, v, master fp32 = 12 B)	12 B	840 GB
Activation(取决于 batch + seq_len, 估算)	—	几十~几百 GB
合计		>1.2 TB

数据为合理估算。8 张 H100(80GB) = 640 GB,显然装不下,必须并行 + offload。

三种”装不下”的具体表现:

参数装不下:模型 fp16 size > 单卡 HBM。需要 ZeRO-3 / FSDP 切分参数到多卡
优化器状态装不下:Adam 的 m/v/master 占 12B/param,比参数本身大 6×。这是 ZeRO 设计的重点
Activation 装不下:长 seq + 大 batch 时 activation 会爆掉。Gradient checkpointing 节省但增加 ~30% compute

🌟 结论:1B+ 参数的训练就开始吃紧,百 B+ 必须 offload。“远程内存”在训练领域不再是高端选项,而是事实标配。

🍎 直觉比喻:GPU HBM 是 L1 cache(快但小),CPU DDR 是 L2 cache(几百 GB,带宽足够),NVMe 是 L3(TB 级,慢但够廉价)。ZeRO-Infinity 就是 GPU 训练的”三级缓存层次”。

2. ZeRO 系列演化:从 ZeRO-1 到 ZeRO-Infinity

ZeRO(Zero Redundancy Optimizer)由 DeepSpeed 提出,核心思想:不让每张卡都存一份完整副本,而是切分,通信时再凑。

传统 DDP:
  Card 0: [全 params] [全 grads] [全 optim] [activations]
  Card 1: [全 params] [全 grads] [全 optim] [activations]
  ...
  N 张卡 ⇒ N×(P+G+O) 内存

ZeRO-1:           切分 optimizer
ZeRO-2:           切分 optimizer + gradients
ZeRO-3:           切分 optimizer + gradients + parameters (对应 PyTorch FSDP)
ZeRO-Infinity:    在 ZeRO-3 基础上加 CPU/NVMe offload

Stage	切分内容	通信代价	内存节省(N 卡)
ZeRO-1	optimizer	与 DDP 相当	optim 状态 N 倍节省
ZeRO-2	optimizer + grads	+ gradient reduce-scatter	+ grads N 倍节省
ZeRO-3	optim + grads + params	+ params all-gather (forward/backward 各一次)	params 也 N 倍节省
ZeRO-Infinity	同 3 + CPU/NVMe offload	+ CPU↔GPU 搬运	HBM 不限模型大小

关键数字:ZeRO-3 把 70B 训练的单卡内存从 ~150 GB 压到 ~20 GB(8 卡)。ZeRO-Infinity 进一步把 optimizer state 倒去 CPU/NVMe,单 GPU 训练 1T 参数也成立。

⭕ PyTorch FSDP 实质等价 ZeRO-3,工程实现更接近 PyTorch 原生 API,生态更广(2023 后 Meta 主推)。

3. ZeRO-Infinity 的三层 offload:HBM / DDR / NVMe

ZeRO-Infinity(SC’21)的核心是把训练状态分到三层物理介质:

       GPU HBM (TB/s)        ─── parameters (forward/backward 时按需 fetch)
       ────────────────
       CPU DDR (~100GB/s)    ─── gradients + optimizer working set
       ────────────────
       NVMe (~10GB/s)        ─── optimizer cold state (Adam m, v)

带宽和容量量级(2026 视角的代表硬件):

层	带宽	单节点容量	延迟
HBM3(H100/MI300)	3-5 TB/s	80-192 GB	几十 ns
DDR5(server)	100-200 GB/s	512 GB-2 TB	~100 ns
PCIe NVMe(Gen5)	12-14 GB/s	几 TB-几十 TB	10s of µs

关键工程:prefetch + overlap——下一层的数据要在上一层用完前预先 fetch 上来,GPU 不能停等。ZeRO-Infinity 在前向 layer N 时,异步 fetch layer N+2 的参数,让 compute 和 IO 重叠。

🧠 关键洞察:训练比推理更适合 offload,因为 batch 大、计算密集——单步 forward+backward 的 GPU 时间足够掩盖 NVMe 几十 µs 的读延迟。推理 decode 阶段单 token 计算只需几毫秒,基本无法 hide NVMe 延迟,这是为什么推理领域反而用 RDMA 内存池而不是 NVMe。

Activation memory 的处理:

ZeRO 不切 activation(activation 数据流式产生消耗,切了无意义)
用 Gradient Checkpointing(activation 不全存,反向时部分重算)解决
极端情况下 activation 也 offload 到 CPU(代价大,只在 OOM 时启用)

ZeRO-Infinity 的能力极限(数据来自 SC’21 论文):

32 TB optimizer state offload to NVMe → 训练 32T 参数模型(理论)
单节点(8x V100)能跑下 1T 参数 fine-tune

这些数字在 2026 年硬件下可以等比放大(H100 + Gen5 NVMe + CXL 内存扩展)。

4. 推荐系统:HugeCTR 与 PMEM Hybrid Embedding

推荐系统训练的”内存不够”和大模型不一样——不是参数密集,而是 embedding 表巨大。

Click-Through-Rate(CTR)模型:
  - 用户 ID embedding:亿级 user × 64 dim = 几十 GB
  - 商品 ID embedding:百万级 item × 64 dim = 几百 MB
  - 类目 / 标签 / cross feature embedding × N 个:总计 TB 级
  - 但 dense 部分只几 GB(MLP)

特征:

大头是 sparse embedding 表,TB 级
单次 forward 实际查 lookup 的只是其中一小部分(cache locality 极差)
dense MLP 部分计算重(GPU friendly)

HugeCTR(NVIDIA Merlin)的 hybrid embedding:

       ┌────────────────────┐
       │  GPU HBM           │   <-- hot embedding(高频 item)
       └────────────────────┘
       ┌────────────────────┐
       │  CPU DDR           │   <-- warm embedding
       └────────────────────┘
       ┌────────────────────┐
       │  PMEM(Optane)      │   <-- cold embedding(长尾 user/item)
       └────────────────────┘
       ┌────────────────────┐
       │  NVMe              │   <-- 全量持久化
       └────────────────────┘

为什么 PMEM 在 embedding 上比纯 GPU 便宜:

Embedding lookup 是随机小读(单次 64-256 B),延迟敏感但不要超高带宽
PMEM(Intel Optane,虽然 2022 年宣布停产但生态仍部分存在)单次随机读 ~300 ns,比 NVMe 快 30×
TB 级 PMEM 的 GB 单价比 HBM 低 ~50×

⚠️ 2026 现状:Intel Optane 停产后,PMEM hybrid 路线的硬件供给变难。接力者是 CXL Type 3 Memory Expander——同样提供”几百 GB-TB 级,延迟比 NVMe 低”的中间层。

🌟 NVIDIA Merlin SOK(Sparse Operations Kit)的关键贡献:在 GPU 上原生实现 hybrid embedding 的查询逻辑,避免 PCIe 来回。SOK + HBM/DDR/PMEM 多级存储是工业推荐系统(阿里、字节、Pinterest)的标配。

5. Checkpoint 远端写:RDMA-based 训练 fault tolerance

千亿模型 checkpoint 大小:TB 级。每隔 N 步全量写一次盘,写延迟直接拖训练 throughput。

传统 checkpoint:
  step done → block all GPU → write all params/optim to disk
              ──────────  几分钟级别 ──────────  (千亿模型 TB 级)
                                             → 训练 GPU 全在等

优化路径(几招):
  1. Async checkpoint:  写盘和训练 overlap(占 ~1% throughput)
  2. Sharded checkpoint:每个 rank 只写自己那份,用 RDMA 汇聚
  3. RDMA-direct write:GPU 直接 RDMA 写到远端 storage(GPUDirect Storage)
  4. Memory-only checkpoint:写到另一组节点的 DRAM,异步刷盘
  5. Differential checkpoint:只写变化部分(利用稀疏性)

关键工作:

GPUDirect Storage(NVIDIA):GPU HBM ↔ NVMe via RDMA 直连,绕过 CPU
Pollux(SOSP’21):弹性训练 + 检查点优化
Bamboo(NSDI’23):基于 redundant computation 的 checkpoint 优化
Check-N-Run(NSDI’22):工业级 checkpoint 经验论文

🌟 结论:RDMA 让 checkpoint 不再是训练的瓶颈——TB 级 checkpoint 在分布式 + RDMA + 异步配合下,可以做到 <1% 训练时间损失。

6. Parameter Server 在大模型时代:还没死

Parameter Server(PS)架构曾经是 DL 训练主流(2014-2018),后来被 collective(NCCL/AllReduce)架构取代。但在两个场景下 PS 仍然活跃:

6.1 推荐系统(MoE-like sparse models)

推荐 embedding 表巨大,collective allreduce 没法处理稀疏更新
PS 架构天然支持”按 key 拉取/推送”,对稀疏访问友好
阿里 PAI / 字节 BytePS / Pinterest 的内部训练栈都是 PS-based

6.2 MoE(Mixture of Experts)训练

MoE 路由后每个 expert 只被部分 token 触达
全 collective 浪费带宽
部分系统用 PS-style “expert parameter server”

PS 架构:                      Collective(NCCL)架构:
                                       
   Worker 1                       Worker 1 ←→ Worker 2
       │                              ↑           ↓
   pull/push                       Worker 4 ←→ Worker 3
       │                          (ring/tree allreduce)
   PS Server 1
   (params shard)

维度	PS	NCCL/Collective
适合	稀疏更新 / MoE / 推荐	dense LLM / CV / 同构
通信模式	点对点 pull/push	ring/tree allreduce
容错	PS 副本 + heartbeat	一卡挂全断
工业落地	阿里 PAI、字节、Pinterest	DeepSpeed、Megatron、PyTorch DDP

🌟 结论:PS 不死,但改了形态——从”大模型主路线”退到”稀疏 / MoE 专用”,仍是大厂训练栈的活跃组件。

7. 训练侧 disaggregation 的局限

7.1 Bandwidth wall(NVMe 写盘)

NVMe Gen5 单盘 ~14 GB/s,1 TB optimizer state 全量交换需要 ~70 秒——这是 ZeRO-Infinity 实际能 sustain 的工作集大小上限。超过这个量级就要靠 layer-wise overlap 或多盘并行。

7.2 Prefetch 不准的代价

ZeRO-Infinity 的 overlap 假设是层级预测——下一层马上要用。但有些模型(MoE 的稀疏路由、动态图)预测不准,fetch 错的数据被丢弃,GPU 等真数据——这种 stall 在长 seq 训练里会累积放大。

7.3 为什么 CXL pool 在训练侧暂时没大爆发

CXL 在推理侧(KV pool、embedding lookup)非常合适——随机小读、延迟敏感。但训练侧:

训练 IO 模式是大块顺序(layer parameters 整块拉),NVMe + RDMA 已经够用
训练 batch 大,几十 µs 的 NVMe 延迟可被掩盖,CXL 100ns 的优势不明显
CXL 容量上限(几 TB)反而不如 NVMe(几十 TB)

训练侧 CXL 真正发力可能在:checkpoint to memory pool——把检查点写到 CXL pool(几百 ns 一致访问)而不是 NVMe(µs 级)。但这是 frontier,2026 年仍是研究阶段。

⚠️ 失败模式总结:

Optimizer offload 错时机 → 反而比纯 GPU 慢(IO 没法 hide)
Prefetch 模型不准 → GPU 周期性 stall
Checkpoint 太频繁 → 把 RDMA 带宽占满,训练通信反受影响

✅ 自我检验清单

三种状态:能徒手算 70B 模型训练时 HBM 需求(参数 + Adam state + activation)
ZeRO 演化:能讲清 ZeRO-1/2/3 各自切了什么
三层 offload:能默写 HBM / DDR / NVMe 各层带宽和延迟量级
推荐 PMEM:能讲清为什么 PMEM 在 embedding 表上比纯 GPU 便宜
Checkpoint 优化:能列出至少 3 种加速 checkpoint 写的工程招(async / sharded / RDMA-direct / mem-first / differential)
PS vs ZeRO:能讲清两者各自适合什么场景
CXL 在训练:能讲清”为什么 CXL 在训练侧暂时没大爆发”

📚 参考资料

关键论文

ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning(Rajbhandari et al., SC’21):arXiv 2104.07857 ⭐ —— HBM/DDR/NVMe 三层 offload 奠基
ZeRO: Memory Optimizations Toward Training Trillion Parameter Models(Rajbhandari et al., SC’20):arXiv 1910.02054
DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters(Rasley et al., KDD’20):KDD 链接
Pollux: Co-adaptive Cluster Scheduling for Goodput-Optimized Deep Learning(Qiao et al., OSDI’21):USENIX 链接
Bamboo: Making Preemptible Instances Resilient for Affordable Training(Thorpe et al., NSDI’23):USENIX 链接
Check-N-Run: A Checkpointing System for Training Deep Learning Recommendation Models(Eisenman et al., NSDI’22):工业级 checkpoint 经验

框架文档

DeepSpeed:deepspeed.ai —— ZeRO 系列官方实现
PyTorch FSDP:官方教程 —— ZeRO-3 等价物,生态更主流
NVIDIA Merlin / HugeCTR:github.com/NVIDIA-Merlin/Merlin —— 推荐系统训练(SOK + hybrid embedding)
GPUDirect Storage:NVIDIA 官方文档

行业讨论

Microsoft DeepSpeed 工程博客 —— ZeRO-Infinity 的实际部署经验
Meta 工程博客 —— FSDP 演进、训练规模化经验

下一章把前面四条应用线放在一张大表上,做横向系统对比与选型决策。

搜索