企业官网建设流程全解析

最近组里在调一个 8x7B 的 MoE 模型，A100 单卡只有 80G，8 个专家全放显存，训练和推理都被 OOM 卡得死死的。我的第一反应是换小模型，后来翻 Megatron 的代码才发现，真正的解法早就写在框架里了：把专家权重搬到 CPU 内存，GPU 上只留当前要算的那一两个专家。这篇文章就把这套 offload 玩法从头到尾捋一遍，从 MoE 的显存账本讲起，到 Megatron 里的参数怎么配，再到我实际操作中踩过的坑，尽量让你拿去就能用。

这篇文章适合正在训练或推理 MoE 大模型、被显存逼到墙角的人，也适合对 Megatron 内部实现好奇、想知道“专家权重到底是怎么搬来搬去”的工程师。放心，我不会只丢一堆参数开关，会连 Why 一起讲清楚。

1. 先搞清楚你在搬什么：MoE 专家权重的显存账本

1.1 MoE 层的结构：Attention 是共享的，专家是私有的

先把 MoE 层的结构说清楚。典型的 MoE Transformer 层由三部分组成：一个共享的多头注意力模块、一个 router（路由网络），以及一组专家 FFN。每个 token 经过 router 算出一个概率分布，然后取 top-k 路由到对应的专家 FFN 里做计算。

关键点在于：Attention 的权重是全模型共享的，而专家 FFN 是“私有”的。一个 8x7B 的 MoE，意思是每层有 8 个 7B 规模的 FFN 专家。如果这 8 个专家全部放在 GPU 上，那么只是专家权重就要占 8 份显存，这还没算 attention、embedding、activation 和优化器状态。

很多第一次接触 MoE 的人会以为专家是“并行的模型副本”，其实不是。专家之间不是 ensemble 关系，而是分工关系——每个 token 只挑其中 k 个专家计算。这种稀疏激活特性，是后面所有 offload 技巧成立的基础。

1.2 一张表算清显存去向

我们拿 8 个 7B 专家、BF16 精度来算一笔账。

也就是说，一个 8x7B 的模型，专家权重占了绝对大头。即使你有 8 张 80G 的 A100，如果不做任何 offload，专家权重就要吞掉 112GB（还要 2 份拷贝给梯度的话更恐怖），剩下的容量根本不够跑训练。

所以问题不是“要不要搬”，而是“搬多少、搬到哪”。答案也很直接：把 7 个专家搬到 CPU 内存，GPU 上只留 1 个专家正在算，显存瞬间从 112GB 降到 14GB，省出来的空间足够跑更大的 batch，或者把模型塞进更少的卡。

1.3 专家在分布式训练里是怎么摆放的

在 Megatron 里，MoE 的专家会通过 expert model parallelism（EP）分布在不同的 rank 上。你设置--expert-model-parallel-size 8，意味着每个 rank 持有全部专家的一部分。token 会通过 all-to-all 通信，被路由到对应专家所在的 rank 上计算，算完再 all-to-all 回来。

这个分布方式对 offload 极其重要。因为每个 rank 只负责其中一部分专家，所以 offload 的粒度不是“整个模型”，而是“这个 rank 负责的那几个专家”。--expert-model-parallel-size设得越大，每个 rank 上需要驻留的专家数就越少，offload 的压力也越小。

我在实际配置时，建议优先把expert-model-parallel-size设成跟数据并行度一致，让每个专家只在一张卡上存在。这样能避免同一个专家被多个 rank 重复加载，减少冗余显存，也让 offload 的语义更干净。

2. 专家权重为什么适合搬到 CPU：访问模式与成本

2.1 稀疏激活是 offload 的底气

专家权重能被搬走，核心原因是 MoE 层的访问模式“天然稀疏”。一个 token 只会路由到 top-k 个专家上，大部分专家在一整步计算里是碰都不碰的。

如果把专家全部放在 GPU 显存里，等于给每个 token 都备好了 8 个专家的权重，但最后只用了 2 个，剩下 6 个是纯纯的“占着茅坑不拉屎”。搬到 CPU 后，GPU 只需要在计算前把被选中的那 1~2 个专家换进来，算完再换出去，闲置权重不占显存。

这个特性跟普通的 dense 模型完全不同。Dense 模型的每一层权重都要参与计算，你没法把其中一部分“暂存”到 CPU，因为每个 step 都要全部用一遍，搬运开销会被放大到不可接受。而 MoE 的专家调用是局部且稀疏的，offload 的性价比自然高得多。

2.2 CPU 和 GPU 之间的搬运成本到底有多高

搬是能搬，但不是免费的。CPU 和 GPU 之间走的是 PCIe 总线，带宽再高也是有上限的。

以 PCIe 4.0 x16 为例，理论带宽约 32GB/s，实际能跑 20~25GB/s 就算不错。一个 7B 专家的 BF16 权重是 14GB，搬到 GPU 需要差不多 0.6 秒。如果每次 forward 都搬，这个延迟会非常可观。

所以 offload 的策略从来不是“每 token 搬一次”，而是“批量凑够了再搬”。用更大的 micro-batch 把一次 forward 里需要同一批专家的 token 攒起来，搬一次算一批，摊薄搬运成本。batch size 越大，单 token 分摊的搬运开销就越低，整体吞吐反而能压得上去。

这里有个很多人容易忽略的点：CPU 内存带宽同样会变成瓶颈。DDR4/DDR5 的实际带宽也就几十 GB/s 量级，多卡同时往 CPU 读写大量权重，内存控制器会先扛不住。我见过一个案例，8 张卡同时 offload，CPU 内存带宽打满，整体训练速度反而不如直接在 GPU 上用更小的 batch 硬跑。所以开启 offload 之后，要重点监控 CPU 的内存带宽使用率，而不是只看显存有没有降下来。

2.3 不是所有层都能往 CPU 搬

有些读者上来就想“直接把整个模型扔到 CPU 里”，这可行度很低。归纳一下，能 offload 的边界大概是这样：

• 可以搬：专家 FFN 的权重。因为调用稀疏。
• 不建议搬：embedding、attention、layer norm、router 参数。这些每步计算都全量参与，搬来搬去得不偿失。
• 绝对别搬：optimizer 状态和梯度。训练时这些数据和计算紧密绑定，放在 CPU 会导致梯度同步效率暴跌。

另外，如果模型里还有 shared expert（共享专家），我建议把 shared expert 留在 GPU 上。共享专家的调用频率最高，搬出去会反复触发换入换出，严重影响训练吞吐。

2.4 为什么选 CPU 而不是 NVMe

有人会问，CPU 内存也可能不够，能不能直接 offload 到 NVMe 硬盘？答案是能，但千万别当主力方案。

PCIe 总线上的 NVMe 顺序读也就 3~7GB/s，比 CPU 内存慢一个数量级。训练场景下中间结果要频繁换入换出，NVMe 的读写延迟完全扛不住，基本只能用于“训练前一次性加载权重”这类场景。CPU 内存的定位是“放得下、读得快、但显存放不下”的那部分数据，这才符合 offload 的实际需求。

3. Megatron 里 offload 的工程实现

3.1 核心模块与调用链

Megatron-LM 在较新版本里已经在 MoE 模块内部实现了专家权重的 CPU offload。核心逻辑集中在megatron/core/transformer/moe/这个目录下，重点是experts.py里的MoEExpert类和SwapExpertInOut工具类。

原理不复杂：每个专家在初始化时就把权重放在 CPU 上（或者先放 GPU，再 detach 回 CPU），前向计算前通过swap_in把权重搬到当前 CUDA 设备，计算完再通过swap_out把显存释放掉，权重同步回 CPU 版本。这样 GPU 显存里永远只保留当前正在计算的专家，其他专家都安静躺在 CPU 内存里。

因为不同版本的 Megatron 代码路径有差异，你搜的时候可以直接在仓库里搜SwapExpertInOut或者swap_in_expert这两个关键词，基本都能定位到对应的实现。

3.2 关键配置开关：--cpu-offload 与配套参数

要让 Megatron 把专家搬到 CPU，主要的开关是--cpu-offload，再配合 MoE 的几个标准参数。我整理了一份可用的启动命令，你可以照着改：

python -m torch.distributed.launch --nproc_per_node=8 \ pretrain_gpt.py \ --num-layers 8 \ --hidden-size 4096 \ --num-attention-heads 32 \ --num-experts 64 \ --moe-router-topk 2 \ --moe-expert-model-parallelism \ --expert-model-parallel-size 8 \ --cpu-offload \ --use-flash-attn \ --micro-batch-size 8 \ --global-batch-size 256 \ --bf16

几处关键参数的逻辑：

• --num-experts 64是专家总数，这里按一个较大规模示例写。
• --moe-router-topk 2表示每个 token 激活 2 个专家。
• --moe-expert-model-parallelism开启专家并行，配合--expert-model-parallel-size 8，让每个 rank 只管理 8 个专家。
• --cpu-offload开启后，专家权重会优先放到 CPU 内存，计算时再换入 GPU。

实测下来，这套组合能让单卡显存占用大幅下降。但如果你的模型是 8x7B 这种超大 MoE，--num-experts和 hidden size 要按实际模型结构来填，不要硬抄。

3.3 反向传播时专家权重去哪了

这里要特别讲一下训练场景和推理场景的区别。

如果是纯推理，swap_in之后计算 forward，然后swap_out释放显存，一步就完事，逻辑很顺。

但训练时还有反向传播。反向计算梯度需要当前专家的权重，如果正向算完就把权重换回 CPU，反向时梯度就找不到权重了。Megatron 的做法是在反向前再次换入同一批专家权重，或者依赖 autograd 的计算图，在反向触发时保证权重还在 GPU 上。

这也意味着，训练场景下每个专家在一轮迭代里可能要被换入换出两次甚至更多。搬运开销会翻倍，所以训练时使用 offload 需要更谨慎。我的建议是：训练场景优先检查显存到底缺多少，如果只是差一点点，先尝试减小 micro-batch 或者开启 activation checkpointing；只有当显存缺口很大、batch 压不下去时，才值得引入专家 offload，并且一定要把 micro-batch 调大摊薄成本。推理场景则相反，专家 offload 几乎是稳赚不赔的买卖。

3.4 自己实现一个简化版 swap 逻辑

即使不用 Megatron，理解这套 swap 逻辑也能让你在其他框架里复现。这里给出一段简化的伪代码，展示专家权重的换入换出思路：

class MoEExpert: def __init__(self, ff_dim, device, offload_to_cpu=True): # 权重视为直接创建在 CPU 上 self.w1 = torch.randn(ff_dim, ff_dim, dtype=torch.bfloat16) self.w2 = torch.randn(ff_dim, ff_dim, dtype=torch.bfloat16) if not offload_to_cpu: self.w1 = self.w1.to(device) self.w2 = self.w2.to(device) def swap_in(self, device): if self.w1.device != device: self.w1 = self.w1.to(device, non_blocking=True) self.w2 = self.w2.to(device, non_blocking=True) def swap_out(self, device): # 把最新权重同步回 CPU 引用 cpu_w1 = self.w1.detach().cpu() cpu_w2 = self.w2.detach().cpu() del self.w1, self.w2 torch.cuda.empty_cache() self.w1, self.w2 = cpu_w1, cpu_w2

注意这段代码只是为了讲原理，真正工程实现还要考虑梯度累积、stream 异步拷贝、pin_memory、多个专家并发换入换出等问题。Megatron 内部已经处理了大部分细节，所以能用框架自带功能就尽量别自己造轮子。

4. 实操：把一个 8x7B MoE 模型压到 80G 显存

4.1 环境准备与实验配置

我跑通这套配置用的环境是：单机 8 卡 A100 80G，NVIDIA 官方 PyTorch 容器（PyTorch 2.1 + CUDA 12.2），Megatron-LM 最新主干代码。为了对比效果，我把同一个 MoE 模型分别跑了两种配置：一个完全不开 offload，一个开启--cpu-offload。

需要注意，容器里 CPU 内存要足够大。8 个 7B 专家的 BF16 权重加起来约 112GB，加上其他开销，建议宿主机至少 256GB 内存，否则 CPU 侧会先 OOM。

4.2 验证 offload 是否真的生效

不要只看训练指标，先用最直接的方式验证显存变化。训练开始后，另开一个终端用nvidia-smi看每张卡的显存占用，同时用下面的代码在训练脚本里打印实际显存使用：

import torch def print_gpu_memory(): allocated = torch.cuda.memory_allocated() reserved = torch.cuda.memory_reserved() print(f"[MEM] allocated={allocated / 1e9:.2f}G reserved={reserved / 1e9:.2f}G")

我实测的一组对比数据（8x7B 模型，BF16，micro-batch 8，8 卡）：

可以看到，一旦开了 offload，模型就能在一个 80G 的卡上跑起来了。显存大头从专家权重变成了 activation 和其他共享层。

这里有个细节：即使开启了 CPU offload，单卡显存也还是会随 batch 增大而增长，因为 activation 仍然在 GPU 上。所以不要以为开了 offload 就能无限加大 batch，要随时关注峰值显存，别撞到 80G 的硬上限。

4.3 性能调优的三个关键操作

第一，把 micro-batch 调大，让每次换入的专家能计算尽量多的 token。搬运成本是固定的，分摊到每个 token 上就越少。我试过从 micro-batch 4 调到 16，整体吞吐提升了接近一倍。

第二，给 CPU 侧的数据加上 pin_memory。如果 Megatron 内部没有自动处理，你可以在加载数据或者初始化权重时，把 CPU 张量放到 pinned memory 里，这样 GPU 侧拷贝是异步 DMA，不会阻塞计算流。

第三，尽量复用同一个 CUDA stream 做权重搬运，避免频繁同步。Megatron 内部对这部分已经做了封装，但如果你自己实现 offload，一定要记住：to(device)默认是同步的，不是异步的，直接写在 forward 里会卡住整个计算流。

4.4 训练和推理的 offload 策略不一样

推理场景下，offload 几乎是白赚的。因为不需要反向，每个专家换入后算完就能换出，显存占用可以压得很低。如果你想进一步提高吞吐，还可以配合 batch 内 token 缓存：同一个专家处理完一批 token 再换出，不要一 token 一换。

训练场景下，offload 的代价明显更高。因为反向还要再换入一次权重，导致搬运次数翻倍。我的实测是，训练吞吐大约会下降 20%~40%，具体取决于模型大小、batch 大小和 PCIe 带宽。所以训练时我的建议是：先开 activation checkpoint，再考虑是否 offload 专家；如果两者都要开，priority 应该放在把 batch 调大上，而不是盲目追求显存最小化。

5. 踩坑记录与排查手册

5.1 开了 offload 但显存没降

这个坑我一开始也踩过。检查了两点：一是确认--cpu-offload真的传进了训练脚本，有时命令拼接不对，参数被吞了；二是检查expert-model-parallel-size是不是太小，如果专家并行度不够，单个 rank 可能持有过多专家，offload 的效果被稀释了。另外，如果容器里 CPU 内存不足，PyTorch 可能仍然会尝试在 GPU 上分配，表现为显存不降反升。

5.2 训练时 loss 抖动甚至发散

如果你发现 loss 在开启 offload 后出现抖动，大概率是权重换入换出的时序问题：反向传播时需要的是前向计算时的权重，但权重已经被swap_out换走了，或者在二次换入时拿到的是更新后的权重。这时要检查框架版本里SwapExpertInOut的行为是否覆盖了反向路径。实在不行就退回一项：只对前向做 offload，反向仍然从 GPU 上取权重，虽然显存省得有限，但正确性更容易保证。

5.3 多卡同时 offload 导致卡死或超时

卡死的问题通常出在通信和内存的竞争上。多卡同时往 CPU 内存做大批量读写，PCIe 和内存控制器都会被挤爆，导致 NCCL 通信超时。我的排查方法很简单：先把--expert-model-parallel-size调小，减少跨卡传输量；如果还卡，就检查是不是 CPU 内存带宽被打满，可以通过perf或htop观察内存带宽负载。

5.4 常见问题速查表

5.5 一个隐蔽的性能陷阱：不要只看显存占用

很多人看到显存降下来就觉得万事大吉，结果训练速度慢到怀疑人生。这里要提醒一句：offload 模式下，真正的瓶颈往往不是显存，而是 PCIe 传输和 CPU 内存带宽。nvidia-smi 里 GPU 利用率可能很低，CPU 占用也可能不高，但整体就是卡。这通常是因为搬运和计算在串行执行，没有重叠。解决办法是增大 batch、减少 swap 次数，或者在框架层面开启异步拷贝，让换入下一个专家的同时，当前专家还在计算。

我在实际使用中发现，offload 这套玩法最舒服的场景是推理服务。显存省下来之后，可以腾出空间放更大的 KV cache，batch 能开得更猛，整体吞吐反而上去了。训练场景如果模型大到非 offload 不可，也不是不行，但要把吞吐下降的心理预期准备好，并且优先把 batch 调到最大再谈其他优化。

最后再分享一个小技巧：如果你的模型支持量化，把专家权重在 CPU 侧保存为 int8 或 int4，经济性更好。虽然换入 GPU 后计算前需要反量化，但这部分开销往往比搬运少得多——因为传输的数据量直接砍了一半甚至四分之三，PCIe 带宽瞬间就不是瓶颈了。我们后来就是把专家权重压成 int8 放在 CPU，搬运时间直接翻了倍，训练吞吐也比 BF16 offload 时高了一截。这个方向你可以根据自己模型的情况重点试试。