MoE模型的高效微调

Wed, 21 Jan 2026 11:00:00 +0800

2026年了，MoE大模型微调竟然还没有彻底极其方便的通解。这合理吗？

为什么MoE大模型多卡分布式LoRA微调是困难的？

📌 重要

一句话总结

MoE模型在LoRA微调中产生巨量低秩矩阵，叠加低效的MoE实现，导致多卡梯度计算困难，最终容易卡死。

一开始我也很惊奇，都2026年了，Qwen3模型问世已半年有余，**难道业界就没有MoE做LoRA微调的需求么？**怎么会一直存在这个问题呢？

后来多方查阅资料，这个问题还真没那么好解决：

MoE架构的模型中存在数量极其庞大的小专家FFN，每个专家（Gated MLP）内有三个Linear模块；
在LoRA微调时，LoRA会为每个Linear模块添加两个低秩矩阵(A/B)；
这导致模型中参与计算的模块数量极其庞大，CUDA计算图相较Dense模型极其复杂；
在 HF transformers 4.* 版本中，MoE的专家路由与FFN选择甚至使用Python for循环实现；
低效的专家选择和计算加剧了多卡分布式计算的困难，也容易导致多卡不同rank数据不对齐，reduce失败；
最终导致训练完全卡死，表现为显卡占用率100%却没有计算量，NCCL通信超时；
所有基于transformers库搭建的框架（LLaMA-Factory等）都面临相同问题。

包括我在内，很多人惯性认为，现在的计算能力足以满足更大规模的模型训练，小小MoE引入的计算开销应当不足为惧才对？

让我们看看例子，以Qwen3-30B-A3B为例：

模型中每一层有128个专家
模型总共48层，因此整个模型就有6000多个专家FFN模块
而每个FFN模块都是一个Gated MLP模块，内含3个Linear矩阵
此时，单论专家部分，模型中就总共有18000+ Linear矩阵了，还没考虑Attention等部件
使用LoRA微调时，需要对每个Linear增加两个矩阵(A/B)，再次翻倍变成36000+矩阵参与反向传播

作为对比，Qwen3-32B共64层，每层包括Attention和Norm在内也不过11个Linear矩阵，整个模型应用LoRA后的可更新Linear数控制在2000以内，还不到MoE模型的零头。

解决方案

分析了问题原因，解决方案呼之欲出，无非两种：

Sensei，我不做LoRA微调啦！泷泽LoRA哒！
你去把 Huggingface 师徒干掉。我？

第一种方案是研究论文的主要手段，研究一些小模型尚且可以接受。

但眼下是大模型时代！Qwen3发布的MoE最小也有30B参数量，全参数微调需要的VRAM奔着360G以上去了，只有财大气粗的团队才能轻松实现吧。

第二种方案说得好，既然transformers库的MoE实现有问题，那我们不用不就行了？

很好，现在请开始手搓一个Fused MoE Kernel吧。这就是传说中的从入门到入土吗？

等一下，为什么要自己手搓呢？还好，有人帮我们实现过了！

Megatron + Swift

现在，有请NVIDIA开发的Megatron库，以及Modelscope开发的ms-swift库。

上面两个库是什么，还请自行搜索。总而言之，借助这两个库，我们可以在transformers v5真正解决MoE计算效率问题前，高效地训练MoE模型LoRA。

现在让我们直接开始准备吧！

环境配置

这一套方案需要使用Python 3.11，原始配环境教程在这里。

相较于原始教程，我做了一点点修改，总结了一些踩坑的改进，可以和原文档一起阅读。

假设此时系统已配置好CUDA Toolkit，也就是nvcc那一套东西。

另外，我现在更偏好使用uv作为pip的替代，使用pip时请自行调整。

基础环境

准备以下requirements.txt：

ms-swift==3.12
torch==2.8.0
transformers<4.58
trl<0.25
peft<0.19
deepspeed<0.18
pybind11

需要注意，这一套环境中，PyTorch 2.9.0或更高版本似乎存在一些问题，所以建议使用2.8.0。

接下来，根据机器上NVIDIA驱动的CUDA版本，选择好希望使用的CUDA版本号，然后开装！这里我们使用CUDA 12.6版本的PyTorch。

1
2

uv venv -p 3.11
uv pip install -r requirements.txt --torch-backend cu126

外部依赖项

接下来，我们需要获取一些外部依赖项，这些依赖需要手动从Github上拉取后本地编译安装。

# Create a folder, or NOT as you wish.
mkdir deps

# NVIDIA APEX
git clone https://github.com/NVIDIA/apex deps/apex

# NVIDIA Megatron
git clone --branch core_r0.15.0 https://github.com/NVIDIA/Megatron-LM.git deps/Megatron-LM

本地编译依赖

现在，我们需要开始一项项编译依赖。首先是transformer_engine：

1
2
3

source .venv/bin/activate

SITE_PACKAGES=$(uv run python -c "import site; print(site.getsitepackages()[0])") && echo $SITE_PACKAGES && CUDNN_PATH=$SITE_PACKAGES/nvidia/cudnn CPLUS_INCLUDE_PATH="$SITE_PACKAGES/nvidia/cudnn/include:$SITE_PACKAGES/nvidia/nccl/include" uv pip install --no-build-isolation transformer_engine[pytorch]

上面这一串看起来怪吓人的，其实就是把Python虚拟环境中的include文件夹暴露给编译器，避免重复下载CUDNN或者NCCL等库的麻烦。参考这个Issue。

接下来是APEX：

`1`	`uv pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings="--build-option=--cpp_ext" --config-settings="--build-option=--cuda_ext" ./deps/apex/`

MS-Swift教程中的命令不完全兼容uv指令，需要像上面的例子一样，--config-settings后用等于号=连接参数字符串。

然后是Megatron：

`1`	`uv pip install ./deps/Megatron-LM/`

最后是Flash Attention：

`1`	`uv pip install "flash-attn==2.8.3" --no-build-isolation`

如果没有出错，那么环境就配好啦。

参数调优

请参考Megatron Core MoE

坑

数据加载有问题

上述依赖会在环境中使用datasets==3.6.0，如果你生成的数据集文件使用了不同版本的datasets库，会容易出错。

解决方法很简单：在生成数据集文件时，就使用datasets==3.6.0。

CUDA on 个人技术随想