最近在折腾大规模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架构的基础共享层与专家层MoEMixture of Experts并不是什么新概念但在大模型时代它变成了扩展参数量又不线性增加算力的重要手段。简单来说一个Transformer层里的FFNFeed-Forward Network不再是一个而是变成了很多个并行的FFN这些FFN就是“专家”。每个token经过路由器router时会被打分并分配给top-k个专家去计算所以理论上模型参数量可以很大但每个token实际只激活一小部分专家计算成本可控。Megatron里的MoE层通常会再拆一层一部分是共享专家shared expert一部分是常规专家regular experts。共享专家所有token都会走常规专家按top-k激活。这样的设计让模型既保留了一定的密集计算能力又具备大规模稀疏扩展的空间。这里要算一笔账。假设隐藏层维度H4096专家的FFN中间维度ffn_hidden14336专家数量E64。每个常规专家一般有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上。比如专家并行度EP864个专家会平均分到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.1CUDA 12.xMegatron-Core 0.5以上Python 3.10。如果你还在用老版本的Megatron-LM建议先升级到Megatron-Core因为老版本对MoE的支持确实不够顺手很多配置项都不存在。CPU offload依赖几个基础能力这些都要在环境里确认好torch.cuda.Stream用来做异步拷贝和计算重叠。torch.Tensor.pin_memory()把CPU张量固定在物理内存里这样H2D拷贝可以走更快的内存通道。non_blockingTrue在拷贝到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, dtypetorch.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, dtypedtype).pin_memory() ) self.w_up nn.Parameter( torch.empty(ffn_hidden, hidden_size, dtypedtype).pin_memory() ) self.w_down nn.Parameter( torch.empty(hidden_size, ffn_hidden, dtypedtype).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, dtypedtype, devicecuda), persistentFalse, ) self.register_buffer( _gpu_up, torch.empty(ffn_hidden, hidden_size, dtypedtype, devicecuda), persistentFalse, ) self.register_buffer( _gpu_down, torch.empty(hidden_size, ffn_hidden, dtypedtype, devicecuda), persistentFalse, ) 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_blockingTrue) self._gpu_up.copy_(self.w_up, non_blockingTrue) self._gpu_down.copy_(self.w_down, non_blockingTrue) # 让当前计算流等待拷贝流完成 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-sizeGPU侧可同时驻留的专家缓冲区个数。设置成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 checkpointingexpert层的输入会被重新计算输出也可能被重算。offload版本下专家权重搬运逻辑需要保证是幂等的否则重算时的结果会和原始计算不一致导致loss波动。4.5 问题速查表为了方便排查我整理了一个表格基本覆盖了常见问题现象可能原因解决方向显存没有明显下降每个专家都分配了独立GPU缓冲区改用共享缓冲区限制同时驻留的专家数训练速度暴跌H2D/D2H通信串行等待增加CUDA流使用double buffer开启prefetchCPU内存持续上涨参数重复存储在CPU侧检查是否所有rank都复制了全量专家改为按EP切分多卡all-to-all通信OOMtoken dispatch临时张量过大调整top-k和专家并行度或者分批dispatchBF16混合精度下loss不稳master weight和梯度更新时序不一致在CPU侧统一维护master weight异步回写梯度系统卡死无响应CPU内存不够开始swap到磁盘减少offload专家数优先把热专家留在GPU4.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代码能省掉好几天的调优时间。