上周一个同事跑来找我说他单卡 A100 上加载一个 MoE 模型做推理80GB 显存直接被打满OOM 报错刷屏。他的第一反应是换卡、加卡但数据中心里的 A100 就那么多反而旁边那台机器的 CPU 内存还剩 400 多 GB 空着。他就问了我一句“这些专家权重能不能搬到 CPU 内存里去啊要用的时候再拿回来”这个问题其实问到了点子上。把专家权重从 GPU 显存搬到 CPU 内存——也就是做 offload——在 Megatron 框架下完全可以实现而且不需要改模型结构、不需要重训。但具体怎么搬、搬哪些专家、什么时候搬、搬完会不会反而更慢这里面的门道不少。这篇文章就把在 Megatron 里做 MoE 专家权重 CPU offload 的完整思路写出来适合正在跑 MoE 推理或训练、被显存卡住的人也适合想读懂 Megatron 源码里 MoE 这块的开发者。我会把原理、实现、性能账单和踩坑记录一起说清楚尽量让你看完能直接上手试。1. 显存装不下的现实MoE 模型为什么需要 CPU offload1.1 MoE 的显存账本大头全在专家先别急着写代码得先把显存花在哪算明白。一个 MoE 层里的参数可以分成三类共享的 attention 和 embedding 部分、router很小、以及海量的 expert 参数。以 hidden_size4096、intermediate_size14336、SwiGLU 激活的典型配置为例单个专家由 gate_proj、up_proj、down_proj 三个矩阵组成参数量大约是 2×4096×14336 14336×4096约 1.76 亿参数FP16 精度下就是 352MB。这个数字意味着什么呢64 个专家就是 22.5GB128 个专家直接 45GB。很多 MoE 模型的实际 hidden size 还不止 4096intermediate 到 17920 甚至更高单个专家轻松破 500MB。所以 MoE 模型的显存大头几乎全压在专家权重上而 attention 和 embedding 反而成了小头。MoE 的核心矛盾就是这么来的单次 forward 明明只激活 top-k 个专家但显存必须把全部专家都装下。1.2 为什么偏偏是专家权重适合搬走显存不够的时候很多人第一反应是“把模型参数全搬到 CPU 再按需加载”。这个思路方向对但对象要挑对。我们来对比一下各类参数的 offload 可行性attention 权重、embedding 权重这些参数每个 token 都会用到搬走没有任何休息窗口加载回来反而多一次 PCIe 传输KV cache offload 是另一套做法但它按请求动态变化和权重 offload 不是一个性质而 expert 权重天生就是“按需激活”的——一次 forward 只选 1 到 2 个专家剩下几十个专家可能整批推理都没被选中。这是 MoE expert offload 的根本前提显存里总有一部分专家处于“闲置”状态。把它们放到 CPU 内存等 router 判断某个 token 需要某个专家时再切回 GPU逻辑完全通。这也是为什么我只建议搬专家权重而不是把整个模型一股脑丢 CPU。1.3 什么样的场景下 offload 真的划算不是所有显存不够的场景都该 offload这个判断要先做对。场景一长序列、小 batch 的在线推理服务尤其是逐 token 生成。这种场景 token 少、计算量小专家切换频繁offload 的搬移开销很难被计算掩盖。能用但需要比较精细的调度后面章节会讲怎么救。场景二离线批量推理、模型评测、批量 prompt 处理。这是 offload 最划算的场景因为 batch 足够大我们可以先统计这批数据到底需要哪些专家然后批量加载搬移成本被摊薄到大量 token 上。场景三训练。MoE 训练时每个 step 都要对选中的专家算梯度权重来回搬移不仅频繁还要保证梯度累积的正确性。训练场景我建议优先用 gradient checkpointing、ZeRO、序列并行这套组合拳实在到了显存绝境才考虑专家权重 offload而且通常是在推理微调阶段用。2. Megatron 中的 MoE 层与路由往 CPU 搬的落点在哪2.1 Megatron 的 MoE 层拆开看要在 Megatron 里做 offload首先得知道 MoE 层的代码结构长什么样。Megatron-Core 的 MoE 相关实现集中在megatron/core/transformer/moe.py核心组件有三个Router负责对输入的 hidden states 做 gating输出每个 token 被分配到哪些专家以及对应的概率也就是 top-k 选择。TokenDispatcher负责把 token 按照路由结果分发到不同的专家计算完后再把结果合并回原始 token 顺序内部对应token_permutation和token_combination两个阶段。experts一个 ModuleList每个元素就是一个专家 MLP通常是 SwiGLU 结构。forward 的数据流大致是输入 hidden states 交给 Router 算 top-k然后 TokenDispatcher 把 token 重新排列分发给被选中的专家每个专家只处理分给它的那部分 token最后 TokenDispatcher 把输出还原顺序。这个流程里offload 的切入点就在“Router 拿到选择结果之后、专家计算之前”。换句话说当你知道这一批 token 会用到哪些专家就可以立刻判断哪些专家需要从 CPU 拉回来。2.2 路由结果决定谁是“热专家”谁去 CPUMegatron 的 Router 在每次前向计算中会得到一个tokens_per_expert向量里面记录着每个专家被分配了多少 token。这个向量就是做冷热判定的第一手数据。实际操作中我给每个专家维护一个“最近 N 个 batch 的累计访问次数”。访问次数排在前 30% 的专家属于热专家常驻 GPU 显存永不参与搬移剩下 70% 的冷专家平时住在 CPU 内存。每来一个 batch先看 router 输出如果一个冷专家被选中了就在真正计算它之前把它从 CPU 拉回 GPU。这里有个很容易忽略的点冷热划分不是一劳永逸的。不同数据分布下专家访问频率会发生漂移所以我会每隔固定步数重新统计一次动态调整常驻 GPU 的专家集合。这个机制是必须的否则模型处理了一段与训练分布差异很大的数据热专家名单就失效了。2.3 叠加专家并行每个 rank 只搬自己的专家Megatron 的 MoE 层通常还会配合专家并行Expert Parallelism也就是通过 EP_SIZE 参数把专家列表切到多个 GPU rank 上。这个切分天然就是 offload 的边界每个 rank 只负责自己名下的那部分专家决定它们是留 GPU 还是去 CPU搬移路径也只在本地 CPU 内存和本地 GPU 之间完全不涉及跨节点通信。这个特性在工程实现上省了非常多事。你不需要去处理全局的专家放置表不需要一个中心节点去做“哪个专家在哪个机器上”的查询每个 rank 的 offload 逻辑完全自治。如果叠加了 tensor parallel 和 pipeline parallel情况会复杂一些MoE 层的处理在不同版本 Megatron 里差异不小建议直接对着自己分支的 moe.py 看不要盲信网上旧代码。但从 offload 的角度你只关注每个 rank 本地专家的 device 状态就够了。3. 从 GPU 到 CPU专家权重搬移的实现方案3.1 最朴素的懒加载能跑但别当真把问题缩小到一个具体动作权重怎么在 CPU 和 GPU 之间搬。最直接的方案是改Expert.forward在计算前判断权重在哪个设备上如果发现权重在 CPU 就拿回来def forward(self, x): if self.weight.device ! x.device: self.weight self.weight.to(x.device) self.bias self.bias.to(x.device) return F.linear(x, self.weight, self.bias)这段代码基于常见实践的补全正确性没问题但我强烈不建议把它作为常态方案。原因在于默认的.to()调用是在当前 CUDA stream 上同步执行的一旦遇到冷专家被选中整个 GPU 的计算都得停下来等 PCIe 传输完成。一个 352MB 的专家权重在 PCIe 4.0 下要十几毫秒这段时间 GPU 就是干等。这种方案只适合“显存快炸了、必须把一次推理跑完”的救急场景。3.2 独立 stream 与预取队列让搬移和计算重叠要让 offload 真正可用关键不是“搬得快”而是“别让搬移挡住计算”。CUDA 提供了 stream 机制允许不同操作在不同 stream 上并发执行。我们可以单独开一个prefetch_stream专门负责从 CPU 往 GPU 搬专家权重主计算流继续算当前的 token两个流互不阻塞。下面是我按社区里常用做法补全的最小实现逻辑你可以直接参考这个骨架prefetch_stream torch.cuda.Stream() def prefetch_experts(expert_indices, experts, cpu_weights, gpu_cache): if not expert_indices: return with torch.cuda.stream(prefetch_stream): for idx in expert_indices: if idx not in gpu_cache: gpu_cache[idx] cpu_weights[idx].to(cuda, non_blockingTrue) # 在真正使用这批专家前主计算流等待预取完成 torch.cuda.current_stream().wait_stream(prefetch_stream) # 然后才是 expert 的 forward output experts[idx](token_input)三个关键点第一拷贝必须用non_blockingTrue这样才能在独立 stream 上异步执行第二主计算流的wait_stream要放在“真正使用专家权重之前”而不是预取一发出就立刻等待第三gpu_cache很重要——如果同一个专家被多次选中没必要每次都在 GPU 上往返搬运直接复用缓存权重。3.3 微批次轮转与批量加载真正的工程化姿势光有 stream 重叠还不够因为在线场景里你无法预知下一个请求会命中哪些专家预取永远是“事后补救”。更稳的做法是“微批次轮转”把一个大 batch 切成若干个 micro-batch。第一个 micro-batch 先只跑 router不做实际的专家计算拿到这批数据完整的tokens_per_expert统计。然后看接下来的几个 micro-batch 需要哪些专家一次性批量预取。等真正开始逐个计算 micro-batch 时需要的专家早已躺在 GPU 缓存里了。这里还涉及一个预取窗口的设计一次预取接下来 1 个、2 个还是 5 个 micro-batch 的专家窗口越大命中率越高但 GPU 显存里要预留的缓存空间也越大。按我的经验窗口取 2~3 个 micro-batch 最平衡预取时机大概提前一个 micro-batch 的计算时间这样搬移时间和计算时间能完美重叠。4. 搬移的代价把显存、带宽、延迟的账算清4.1 显存账能省下来多少先说好处。给一张通用表前提是 hidden_size4096、intermediate_size14336、FP16 精度专家数量单专家占用全部专家显存搬走 70% 冷专家后释放8~352MB~2.8GB~2.0GB64~352MB~22.5GB~15.8GB128~352MB~45GB~31.5GB如果你的模型是 128 个专家一次 offload 能腾出 30GB 显存这个量级足够把 batch size 翻一倍或者把序列长度拉长。我自己测试时一个 64 专家的模型冷热比按 7:3 划分后显存占用从 50GB 左右降到 34GB效果非常直观。4.2 带宽账搬一次专家要多久显存省下来了代价是传输时间。业内大概按这个量级估算 PCIe 有效带宽PCIe 3.0 x16 实测约 10~13GB/sPCIe 4.0 x16 约 23~28GB/sPCIe 5.0 x16 可以到 45GB/s 以上。按这个带宽算一个 352MB 的专家权重从 CPU 搬到 GPUPCIe 4.0 约 13msPCIe 3.0 约 28ms。而相比之下GPU 计算一个 micro-batch 的专家 MLP 通常只要几毫秒。如果你的搬移没有和计算重叠一次冷专家切换就能拖慢一个数量级的吞吐。这也是为什么第 3 章讲的 stream 重叠和批量预取不是“锦上添花”的优化项而是 offload 方案能不能成立的基础。只要搬移能在计算的同时完成这 13ms 就不产生实际延迟。4.3 什么情况 offload 反而更慢有些场景下offload 会让性能崩得比纯 GPU 还惨。最容易踩坑的是这几种第一小 batch 在线流式推理。每个请求只带几十个 token计算量极小但专家切换又频繁搬移完全没有隐藏空间。第二长尾请求突然命中一个很久没用的冷专家没有预取到只能现场同步加载。第三多个 GPU rank 同时在同一个 CPU 内存池里抢带宽互相拖慢这种情况在 EP 并行加 offload 时会遇到。对策也很明确尽量把请求攒成 batch 处理、拉大预取窗口、给最近用过的专家做 LRU 缓存以及监控 prefetch 命中率——如果命中率长期低于 80%说明你的冷热划分或预取策略需要调整。5. 实操中的坑与调优经验5.1 拷贝阻塞主计算流最难查的卡顿我在第一版实现里就踩过这个坑GPU 利用率周期性掉到 0kernel 与 kernel 之间出现大段空白但 CPU 侧看程序并没有卡死。用torch.profiler一查时间线发现每次专家切换之后都有一个超长的 memcpy 和一大段空洞。根因是主计算流过早地等待了预取流。我在prefetch_experts被调用后立刻执行了wait_stream而这个等待发生在当前 batch 还没算完的时候等于把两个流硬生生同步了。解决方法是把wait_stream的位置往后挪放到 token_permutation 完成、真正进入专家 forward 之前。这样 attention、permutation 这些计算都能和预取并行。5.2 别在 pageable 内存上栽跟头系统默认的 CPU 张量分配在 pageable 内存上做 H2D 拷贝时驱动要先做一次 bounce buffer 中转性能差距很大。我的实测经验是同样的数据量pageable 内存的拷贝耗时比 pinned memory 慢 30% 到 50%。但别一上来就把所有专家权重都申请成 pinned。pinned memory 是稀缺资源申请太多可能导致系统内存分配失败。稳妥做法是维护一个固定大小的 pinned buffer 池所有专家权重都复制进这个池子里的预分配区域而不是每个专家单独 pin。池子大小按 CPU 可用内存的 50% 左右规划就差不多了。5.3 用完别急着丢给专家权重加个 LRU 缓存刚开始我的实现是“冷专家用完就释放显存、回归 CPU”后来发现频繁的加载和释放不仅带来反复的 PCIe 传输还会产生显存碎片最终可用的显存反而越来越少。更好的方式是 LRU 缓存GPU 显存里始终保持一个“最近使用过”的专家集合缓存上限设为预留显存的 70% 左右满了才把最久没用过的专家逐出到 CPU。这样大部分连续请求都能命中缓存真正做到零搬移。我加上 LRU 缓存之后prefetch 命中率从 65% 直接拉到 92%吞吐提升很明显。5.4 冷热专家的统计别做死要滚动更新另一个常见错误是用训练集统计出“热专家名单”然后固定不变。推理数据的分布会漂移尤其业务流量有周期性原来访问频率很高的专家可能一周后完全没人用。我的做法是启动时先跑 50~100 个预热请求统计专家访问次数按降序取累计占比前 70%~80% 的专家作为常驻 GPU 的热专家之后每处理 500 个请求滚动更新一次统计。这个方案在实验里比固定名单稳定得多而且实现成本很低只是在一个循环里维护计数器而已。6. 再往前一步量化、磁盘与预测式调度6.1 量化后再搬搬移量直接砍半如果觉得带宽还是紧张另一个很成熟的优化方向是冷专家在 CPU 上以 INT8 甚至 INT4 保存需要时加载回 GPU 后再反量化或者直接用对应的 INT8 GEMM kernel。好处非常直接——搬移体积缩小到原来的 1/4 甚至 1/8PCIe 压力骤降。代价是每次加载后多一次反量化处理这部分计算可以放进 prefetch stream不影响主计算流。我在一个 128 专家的模型上试过冷专家全部 INT8 存储后搬移耗时降到原来的 1/4 左右而端到端推理精度损失在一个可接受的范围内。量化方案的取舍主要取决于你对精度的敏感度但作为 offload 的加速手段大概是最划算的一招。6.2 CPU 内存也不够时让最冷的专家住进 SSD模型大到一定程度CPU 内存也装不下所有冷专家这时候可以把最冷的专家继续下沉到 NVMe SSD。加载路径从“CPU 内存 → GPU 显存”变成“SSD → CPU 内存 → GPU 显存”多一跳磁盘 IO。NVMe SSD 的顺序读带宽大概 3~7GB/s比 PCIe 还低一个数量级所以这个方案只建议给访问频率极低的专家用而且预取窗口必须拉得更大否则现场读盘会直接卡死。更细的做法是用 mmap 把磁盘上的专家权重映射进虚拟内存由操作系统管理“什么时候真正从磁盘调入”。这块我还在尝试中没有特别成熟的结论但如果你模型实在太大这是一条能跑通的路。6.3 预测式预取让搬移更聪明一点再说一个我自己在实验的方向目前已有一些效果把 router 的 logits 喂给一个轻量预测头让它预估下一个 micro-batch 大概会用到哪些专家从而把预取时机提前。这和单纯靠历史统计相比优势在于能捕捉到短时间内的模式变化。不过要泼一盆冷水这个方案对大多数团队来说性价比不高因为历史统计加足够大的预取窗口已经能解决 90% 的问题。只有当你发现预取命中率怎么调都上不去、而批量读取又不是瓶颈时才值得考虑这种更智能的调度。我个人在这块折腾下来的体会是专家权重搬 CPU 这件事核心不在“搬移”这个动作上抠细节而在“调度”。把搬移藏在计算后面、把冷热分类做好、把缓存保住CPU offload 完全可以撑起数倍于显存容量的模型。如果你现在正被 MoE 显存问题卡住先别急着加卡按文里的思路把专家权重搬到 CPU 试一遍大概率能救你一次。