MoE大模型显存不够？Megatron下专家权重CPU Offload实战指南

最近在折腾大规模MoE模型时，我几乎被显存问题整崩溃。单卡80GB看着很大，可一旦模型里挂了64个专家，光专家模块的权重就能把显存吃掉大半。Megatron这套框架在模型并行上确实做得很极致，但面对MoE的海量专家权重，它默认也会把所有参数都压在GPU里，结果就是batch size一开大，OOM立刻找上门。于是我开始研究怎么把不常用的专家权重搬到CPU内存里，按需搬回GPU计算，这也就是大家常说的CPU offload。

这篇文章我会完整记录一下，在Megatron体系里把MoE专家权重搬到CPU的实操思路。适合正在用Megatron-LM或Megatron-Core训练MoE大模型、被显存卡死的同学参考。如果你还没到显存不够的阶段，也可以先了解一下方案选型和踩坑点，后面大概率用得上。

1. MoE模型显存为什么吃紧：先搞清楚钱花在哪了

1.1 MoE架构的基础：共享层与专家层

MoE（Mixture of Experts）并不是什么新概念，但在大模型时代它变成了扩展参数量又不线性增加算力的重要手段。简单来说，一个Transformer层里的FFN（Feed-Forward Network）不再是一个，而是变成了很多个并行的FFN，这些FFN就是“专家”。每个token经过路由器（router）时，会被打分并分配给top-k个专家去计算，所以理论上模型参数量可以很大，但每个token实际只激活一小部分专家，计算成本可控。

Megatron里的MoE层通常会再拆一层：一部分是共享专家（shared expert），一部分是常规专家（regular experts）。共享专家所有token都会走，常规专家按top-k激活。这样的设计让模型既保留了一定的密集计算能力，又具备大规模稀疏扩展的空间。

这里要算一笔账。假设隐藏层维度H=4096，专家的FFN中间维度ffn_hidden=14336，专家数量E=64。每个常规专家一般有3个线性层（gate、up、down），参数量大约是3 * H * ffn_hidden，也就是约1.76亿参数。64个专家加在一起，就是约113亿参数。如果是BF16精度，那就是约226GB的内存占用；如果是FP32训练，直接翻倍到45GB以上，这还只是专家模块本身，没算attention、共享专家、优化器状态和激活值。

1.2 显存账本：参数、梯度、优化器状态

训练一个大模型时，显存花销远不止模型参数这一项。推理只需要存参数，训练还需要存梯度和优化器状态。以Adam优化器为例，每个参数通常要额外存一阶动量m和二阶动量v，再加上常用的一些额外状态，每多一个参数，显存就要多出至少8到12个字节。

我们可以按单个专家来估算。每个专家参数量约1.76亿，如果每个参数需要12字节（参数本身2字节BF16 + 梯度2字节 + m 4字节 + v 4字节，具体取决于计算方式），那单个专家就需要约21GB。注意这只是一个专家。如果有64个专家，总需求会非常恐怖。GPU显存的增长速度永远赶不上这种量级的需求，这也是为什么大家一开始就想尽办法把权重往外搬。

Megatron提供了张量并行（TP）、流水线并行（PP）、专家并行（EP）等策略，能把专家分散到不同GPU上。比如专家并行度EP=8，64个专家会平均分到8张卡，每张卡只持有8个专家，显存压力确实小很多。可问题是，如果单卡显存本身就有限，或者EP无法继续扩展，比如只有4张卡却想放64个专家，那单卡仍然会被塞爆。CPU offload就是在这些并行策略都撑不住的时候，最后一道救命稻草。

1.3 为什么只搬专家权重，而不是全部

很多人会问：既然要offload，为什么不把整个模型都搬到CPU？答案其实很现实：CPU到GPU的传输带宽有限，而且attention等密集计算部分对延迟非常敏感，如果这些权重也走CPU，训练速度会直线下降。专家权重是体量最大但访问密度最低的部分，每个token只会用到top-k个专家，也就是说绝大多数时候，大部分专家只是安安静静躺在显存里占地方。

MoE的稀疏激活特性，决定了专家权重是天然的offload候选者。把不常用的冷专家放到CPU内存里，等router真正选中它时再临时搬到GPU计算，算完再放回去。这样一来，GPU显存只需要容纳一个或少数几个活跃专家就够了，而不是同时容纳全部64个专家。CPU内存通常比GPU显存大得多，几百GB甚至上TB都很常见，给专家权重找个“便宜仓库”是可行的。

不过CPU offload不是白拿的，它用通信带宽换显存，所以并不是所有场景都适合。这个取舍我在后面的章节里会详细展开。

2. CPU Offload的几种常见实现路线

2.1 全量offload还是按需offload

把专家权重搬到CPU，听起来很简单，但落地时有两个路线：全量offload和按需offload。

全量offload是指所有专家参数一开始就放在CPU上，forward计算时把所有需要的专家统一搬到GPU，算完再全部搬回去。这种做法的实现最简单，代码逻辑很直接，但问题也很突出：如果几个专家同时被token选中，且每个专家的权重都不小，一次性搬运到GPU时的显存峰值依然很高。而且每一步搬运的数据量很大，通信等待时间会被拉长。

按需offload是指把共享专家、高频活跃专家留在GPU上，只把不常访问的冷专家放到CPU。router计算出当前batch的top-k专家选择后，再针对性地把需要的专家从CPU搬到GPU。这种方式更符合MoE稀疏激活的实际情况，显存占用更少，通信开销也更可控，但实现复杂度高一些，需要自己管理专家热度统计和搬运调度。

我实际做的时候选的是按需offload。原因很简单：MoE模型里不同专家的访问频率是有偏向的，某些专家可能经常被激活，把它们放在CPU里只会白白增加通信压力。把热专家放在GPU，把冷专家放到CPU，能在显存和速度之间取得更好的平衡。

2.2 Megatron原生具备的能力与局限

Megatron-LM和后来的Megatron-Core对MoE的支持已经比较成熟。在Megatron-Core里，可以通过--num-experts、--moe-expert-model-parallel-size、--moe-grouped-gemm等参数控制MoE层的规模和并行方式。--moe-expert-model-parallel-size决定专家并行度，也就是把专家平均分到多少张卡上。这些特性在处理大规模MoE模型时很有用，但有一个尴尬的地方：Megatron原生并没有提供一键把专家权重放到CPU的功能。

这不是说Megatron做不到，而是它把主要精力放在了模型并行和集合通信上，CPU offload这种偏“硬件适配”的事情，默认没做成通用开关。导致的结果就是，如果你只想通过改启动参数来搬权重，基本是做不到的。要么自己改MoE层的实现，要么借助DeepSpeed这类第三方库。

这里简单提一下DeepSpeed的ZeRO-Offload。它在训练阶段可以把参数、梯度和优化器状态搬到CPU，且对普通Transformer层支持得不错。但问题在于，MoE + Megatron这种组合本身就涉及大量的all-to-all通信和自定义并行策略，强行套DeepSpeed的Offload逻辑，往往会出现兼容性问题，比如通信原语冲突、参数切分方式不一致等。我一开始也试过，后面还是决定自己写一个轻量的offload层。

2.3 我为什么选“自定义Expert层”而不是照搬DeepSpeed

照搬DeepSpeed能省不少事，但代价是失去对Megatron并行策略的精细控制。MoE在Megatron里的并行方式很有意思：专家并行和张量并行可以配合使用，每个rank可能持有模型的不同分片。DeepSpeed的ZeRO-Offload是按全局参数来切分的，它并不知道Megatron的MoE层里谁该被切分、谁该被offload。

所以我最后选择自己动手，写一个支持CPU offload的Expert模块。这样做有几个好处：一是继续保持Megatron原有的TP、PP、EP设置，不需要改框架核心；二是可以精确控制哪些专家留在GPU、哪些搬到CPU；三是offload逻辑可以完全自定义，比如加入预取（prefetch）策略，根据router的统计结果提前把未来几步可能用到的专家搬到GPU缓冲区。

缺点是开发量相对大，需要自己对PyTorch的存储机制、CUDA流和异步传输有比较清楚的理解。不过一旦跑通，这套逻辑就能作为通用模块复用在其他模型上。

3. 实操：MoE专家权重搬到CPU的完整流程

3.1 环境准备与版本选择

先说我用的环境，给大家一个参考：PyTorch 2.1+，CUDA 12.x，Megatron-Core 0.5以上，Python 3.10。如果你还在用老版本的Megatron-LM，建议先升级到Megatron-Core，因为老版本对MoE的支持确实不够顺手，很多配置项都不存在。

CPU offload依赖几个基础能力，这些都要在环境里确认好：

• torch.cuda.Stream，用来做异步拷贝和计算重叠。
• torch.Tensor.pin_memory()，把CPU张量固定在物理内存里，这样H2D拷贝可以走更快的内存通道。
• non_blocking=True，在拷贝到GPU时使用异步方式，避免同步等待。
• 确认CPU内存足够大。建议至少比专家权重总量多出20%的余量，否则系统会开始swap，性能直接崩掉。

我的建议是先用一块卡做小规模验证，比如8个专家、小hidden size，把整个offload逻辑跑通，再放到真正的模型上。不要在64个专家的大模型上直接调bug，调试成本太高。

3.2 核心思路：把专家参数从GPU Parameter改成CPU pinned memory

最核心的思路，是把专家层的nn.Parameter从一开始就注册在CPU上，并且用pin_memory()把它钉住。这样参数本身不占GPU显存，只在需要计算时临时拷贝到GPU的缓冲区里。

下面是一个简化的CPU offload专家模块示例：

import torch import torch.nn as nn class CPUOffloadExpert(nn.Module): def __init__(self, hidden_size, ffn_hidden, dtype=torch.bfloat16): super().__init__() self.hidden_size = hidden_size self.ffn_hidden = ffn_hidden # 参数直接创建在CPU上，并锁定物理内存 self.w_gate = nn.Parameter( torch.empty(ffn_hidden, hidden_size, dtype=dtype).pin_memory() ) self.w_up = nn.Parameter( torch.empty(ffn_hidden, hidden_size, dtype=dtype).pin_memory() ) self.w_down = nn.Parameter( torch.empty(hidden_size, ffn_hidden, dtype=dtype).pin_memory() ) # 初始化：这里用简单的kaiming方式，实际请按模型配置走 with torch.no_grad(): self.w_gate.normal_(0, 0.02) self.w_up.normal_(0, 0.02) self.w_down.normal_(0, 0.02) # 在GPU上准备一块可复用的临时缓冲区，避免反复申请显存 self.register_buffer( "_gpu_gate", torch.empty(ffn_hidden, hidden_size, dtype=dtype, device="cuda"), persistent=False, ) self.register_buffer( "_gpu_up", torch.empty(ffn_hidden, hidden_size, dtype=dtype, device="cuda"), persistent=False, ) self.register_buffer( "_gpu_down", torch.empty(hidden_size, ffn_hidden, dtype=dtype, device="cuda"), persistent=False, ) def forward(self, x, copy_stream): # 等待上一个使用该缓冲区的stream结束 copy_stream.wait_stream(torch.cuda.current_stream()) # 把CPU参数异步拷贝到GPU缓冲区 with torch.cuda.stream(copy_stream): self._gpu_gate.copy_(self.w_gate, non_blocking=True) self._gpu_up.copy_(self.w_up, non_blocking=True) self._gpu_down.copy_(self.w_down, non_blocking=True) # 让当前计算流等待拷贝流完成 torch.cuda.current_stream().wait_stream(copy_stream) # 在GPU缓冲区上执行专家计算 h = torch.nn.functional.silu( torch.matmul(x, self._gpu_gate.t()) * torch.matmul(x, self._gpu_up.t()) ) y = torch.matmul(h, self._gpu_down.t()) return y

这段代码的核心是：权重始终在CPU侧，GPU侧只维护一块临时缓冲区，计算时用copy_把权重搬过来。copy_stream由外部传入，可以让不同专家的拷贝走不同的CUDA流，这样它们之间不会互相阻塞。

需要注意，self._gpu_gate那三个buffer不是模型参数，不会参与梯度计算。梯度会通过临时缓冲区回流到CPU参数上吗？实际上，由于计算图的叶子节点是CPU参数，反向传播时梯度会沿着copy_操作从GPU临时张量传回CPU参数，所以优化器可以直接在CPU上更新参数。这正好利用了PyTorch自动求导的支持。

但这里有个非常关键的细节：如果每个expert都有自己的固定GPU缓冲区，64个专家就有64套缓冲区，显存照样吃不住。所以更合理的做法是只维护一个GPU侧的“专家暂存区”，也就是同时只驻留一个或少数几个专家的权重，计算完立刻释放或覆盖。

3.3 搬运流程与调度逻辑

真正放到MoE层里，事情会复杂一些，因为有一个router在决定哪些专家会被激活。我的做法是分三步走：

第一步，router先完成token到专家的分配。这一步可以沿用Megatron里的token dispatch逻辑，先不用管CPU或GPU，一切照旧。

第二步，根据router分配结果，把当前step需要使用的专家列表取出来。对这些专家做依次搬运和计算。每次只搬一个专家到共享缓冲区，算完这个专家的结果就把它在缓冲区里标记为可覆盖，再搬下一个。

第三步，所有专家计算完成后，将结果按照router的索引重新组合，得到MoE层的输出。

这样做的好处是，GPU缓冲区只需要容纳单个专家的权重，显存峰值被压到极低。代价是如果某个step激活了大量专家，搬运和计算会变成串行，速度受影响。为了缓解这个问题，可以把缓冲区扩大到两个或四个专家的大小，让一部分拷贝和一部分计算重叠起来。

我实际用的是两个缓冲区的double buffer方案：一个缓冲区正在被计算流使用，另一个缓冲区正在被拷贝流填充。这样等计算完第一个缓冲区时，第二个大概率已经拷贝完成，直接开始计算，不需要干等。实现上并不复杂，就是把copy_stream和计算流都放在一个循环里交替等待。

3.4 通信带宽估算：什么情况下划算

很多人会忽略一个关键问题：CPU offload到底值不值，要看你CPU和GPU之间的传输带宽。我以PCIe 4.0 x16来算，单向理论带宽约32GB/s，实际能跑到25GB/s就算不错。一个专家的三个权重，参数总量约1.76亿个BF16，也就是约3.5GB的数据。H2D拷贝一次大概需要140ms左右，D2H再搬回去又要140ms。如果每个step都要搬几个专家，这开销很快就把训练时间拖垮了。

但实际并没有这么悲观，因为你不需要一个step把所有专家都搬一遍。MoE的top-k通常在2到8之间，也就是说一个token只会激活少数专家。如果用的是按需offload，每个step只需要把被选中的那几个专家搬上GPU，比如4个专家，那就是约14GB的数据，耗时约560ms。这个数字虽然不小，但如果和GPU上的计算重叠起来，训练吞吐的损失可能控制在20%到30%以内，而显存占用能降下来好几倍。

如果你的机器支持NVLink-C2C或者PCIe 5.0，带宽会更高，offload的性价也会更好。反过来，如果你的机器是老旧的PCIe 3.0，单向带宽只有约12GB/s，搬一个专家就要接近300ms，这种场景下我建议慎重，不如尝试缩小专家数量或降低专家并行度。动手之前，先跑一下nvidia-smi topo -m看看你的GPU和CPU之间的连接方式，再决定要不要上offload。

3.5 性能调优参数与启动项

我自己实现时新增了几个启动参数，方便在训练脚本里控制offload行为：

python pretrain_gpt.py \ --num-experts 64 \ --moe-expert-model-parallel-size 8 \ --enable-cpu-offload-expert \ --cpu-offload-prefetch-num 4 \ --cpu-offload-buffer-size 2 \ --cpu-offload-pinned-memory

这些参数的含义分别是：

• --enable-cpu-offload-expert：开关，打开后专家权重创建在CPU pinned memory上。
• --cpu-offload-prefetch-num：预取专家数，根据router历史统计，提前把未来几步可能用到的专家搬到GPU缓冲区。
• --cpu-offload-buffer-size：GPU侧可同时驻留的专家缓冲区个数。设置成2，就是double buffer；设置成4，可以进一步重叠拷贝与计算，但显存峰值会相应提高。
• --cpu-offload-pinned-memory：启用pinned memory，建议总是开启。

在分布式多卡场景下，这些参数需要配合数据并行组和专家并行组一起使用。每个rank只负责自己那部分专家，offload也是局部的。因为每个rank的本地专家数量变少了，单个rank的CPU内存占用也随之降低，多卡之间不会重复存储同一批专家权重。这一点需要在实现时注意，别把全局专家都复制到每张卡的CPU上。

4. 常见问题与排查技巧实录

4.1 显存没有减下去

这个问题我一开始就踩过。参数确实放到CPU上了，但forward时又把同一批专家权重全部拷贝回了GPU，而且每个专家都申请了自己的临时缓冲区，导致显存峰值一点没降。排查的时候发现，问题出在我为每个专家都单独实例化了GPU缓冲区，64个专家等于重新开了一整套GPU内存。

解决办法是共享缓冲区。全局只维护2到4个“专家槽位”，每个槽位在任意时刻只保存一个专家的权重。被router选中的专家，按顺序复制到空闲槽位中进行计算。算完后槽位立即释放，供下一个专家使用。这样显存里面的专家权重总量永远不会超过缓冲区数量乘以单个专家大小。

4.2 CPU通信成为瓶颈

如果你发现训练速度下降明显，先别急着改代码，看一下是不是通信已经饱和。我的排查方法很简单：在训练循环里加一个计时器，分别统计专家计算耗时、H2D拷贝耗时、D2H拷贝耗时、等待耗时。如果拷贝等待占比超过40%，说明H2D/D2H已经成为瓶颈。

这种情况下有几个调整方向。一是减少每次搬运的数据量，也就是让专家切分得更细，比如把单个专家再按中间维度切块，搬运一块算一块。二是增加预取，让搬运提前发生。三是减少同步点，把torch.cuda.synchronize()的调用次数降到最低，让多个流充分重叠。四是升级硬件，比如从PCIe 3.0切换到支持更高带宽的平台。最后这一点不是开玩笑，机器硬件直接决定了offload方案的上限。

4.3 训练速度下降太快

训练速度暴跌，除了通信瓶颈之外，还有一个很常见的原因是混合精度和master weight的处理出了问题。CPU offload之后，如果优化器状态在CPU侧，而模型权重在GPU临时缓冲区上，梯度回传时如果不同步，容易出现数值抖动；如果每个step都强制同步，又会把流水线彻底打断。

我的建议是：不要把梯度每次都用同步方式回传CPU。用异步拷贝，让梯度在backward结束时自动排入回传队列，由独立stream执行D2H。优化器step则在CPU侧做，参数更新后再标记对应的GPU缓冲区为脏数据，下次需要时再重新搬运。这样虽然逻辑复杂一点，但训练吞吐能得到明显恢复。

4.4 与Megatron并行策略冲突

在Megatron里，MoE层通常是张量并行和专家并行混着来的。专家并行可以把不同专家放到不同rank上，每个rank只处理自己那部分专家。offload必须清楚这些专家是“本地专家的全局副本”还是“切分后的分片”。

如果你把张量并行也用在了专家层，那一个专家的权重可能被切成TP份，每个rank只持有其中一份。offload时只能offload自己这份分片，不能试图把整个专家搬走，否则通信结果会错乱。检查--moe-grouped-gemm和--disable-moe-tp这两个参数，确认你的专家层到底走的是哪种并行方式，再决定offload粒度和索引映射。

另外，如果开启了activation checkpointing，expert层的输入会被重新计算，输出也可能被重算。offload版本下，专家权重搬运逻辑需要保证是幂等的，否则重算时的结果会和原始计算不一致，导致loss波动。

4.5 问题速查表

为了方便排查，我整理了一个表格，基本覆盖了常见问题：

4.6 一个容易被忽略的坑：CPU内存分配碎化

最后说一个不太容易注意到的点。CPU offload运行久了，CPU内存会出现碎化，导致pinned memory申请失败，训练突然崩溃。我的做法是在训练启动前一次性把可能用到的CPU专家权重都分配好，整个训练过程中不再频繁申请和释放CPU内存，只做参数内容的覆盖更新。这样虽然初始内存占用高一些，但稳定性会好很多。

5. 一些个人体会和建议

CPU offload不是银弹，它是在“显存实在不够”的前提下，用通信带宽换存储容量的一种妥协方案。如果条件允许，优先还是考虑增加GPU数量、降低专家并行度、使用更高效的专家实现，这些都比offload对训练速度的影响小。但如果你已经站在OOM边缘，那我上面这套按需offload + 共享缓冲区 + CUDA流重叠的方案，算是目前实践下来最平衡的做法。

从工程实现角度，我最大的体会是：别把offload理解成简单的.to('cpu')然后再.to('cuda')。直接这样做，一次同步拷贝就会打断整个流水线，训练效率低到没法接受。真正有效的offload一定是异步的、分块的、知道哪些权重该在什么时候驻留在哪里的。把这些调度逻辑写清楚，offload才能从“应急方案”变成“可选配置”。

最后再分享一个小技巧：在你把专家权重搬到CPU之前，先用torch.profiler或者简单的torch.cuda.Event打点统计一下，看看当前模型里到底哪些层在等什么。很多时候你会发现，真正拖慢训练的并不是专家权重的搬运，而是因为搬运导致其他通信原语的排队。先把这张调度图梳理清楚，再动手写offload代码，能省掉好几天的调优时间。