一、MoE架构与Dense架构对比1.1MoE架构概述MOE 架构的基本思想是在传统 Transformer 模型中将每个前馈网络FFN层替换为一个 MOE 层。一个 MOE 层通常由两个关键部分组成专家网络这些层代替了传统 Transformer 模型中的前馈网络 (FFN) 层。MoE 层包含若干专家(例如4个)每个专家本身是一个独立的神经网络。在实际应用中这些专家通常是前馈网络 (FFN)但它们也可以是更复杂的网络结构。门控网络该模块负责根据输入token的特征动态选择激活哪些专家: 在下图中“More”这个令牌可能被发送到第二个专家而“Parameters”这个令牌被发送到第一个专家。1.2MoE技术基础1.2.1专家网络角色在 MOE 模型中每个专家网络只在接收到相应路由时参与计算。初期由于门控网络参数随机各专家接收到的数据分布比较均匀但随着训练进行局部梯度更新使得某些专家逐渐专注于处理特定类型的输入数据。这种自发专化现象使得模型整体具备了多样化的表示能力。尽管MoE架构有着很大潜力但现有的MoE架构可能存在知识混杂Knowledge Hybridity和知识冗余Knowledge Redundancy的问题限制了专家的专业化具体如下知识混杂分配给特定专家的token可能会涵盖不同知识知识混合问题出现在专家数量有限的场景下容易阻碍专家们在特定领域进行深入。知识冗余当MoE模型中的不同专家学习到相似的知识时就会出现知识冗余这与模型设计初衷相违背。解决方案引入共享专家和细粒度/垂类专家。1.2.2 门控网络机制门控网络一开始的路由决策可能是近似随机的但随着专家逐步积累专长门控网络也会调整其路由策略。如果某个专家因早期获得较多特定类型数据而表现出色门控网络便倾向于将更多此类数据路由给它。专家因接收到较多特定数据而“专长”而门控网络根据反馈不断更新参数使得路由更加精准。路由编码机制基于softmax具体分为以下两种软编码路由每个输入token按概率权重分配给所有专家输出是各专家处理结果的加权融合硬编码路由每个输入token仅分配给得分最高的k个专家如top-2其他专家完全不参与该token的处理。1.3 MoE 与 Dense 架构深度对比Dense架构稠密模型MoE架构稀疏模型技术影响激活机制全激活所有参数参与每次前向计算稀疏激活每层仅Top-k专家通常k1~2参与计算MoE实现条件计算显著降低计算复杂度计算效率计算量随模型规模线性增长O(L·D²)D为隐藏维度计算量由激活参数决定O(L·k·E)E为专家规模同等性能下MoE可减少3~10倍计算量显存占用与模型参数量成正比参数量 存储需求存储-计算解耦需加载所有专家权重激活量仅占总量10~25%MoE需更大显存存储参数但计算时仅激活部分推理延迟特性延迟稳定可预测与序列长度呈二次关系延迟等于小模型但权重加载可能成瓶颈MoE实际延迟受内存带宽制约明显显存带宽需求中等权重加载一次重复使用极高需频繁加载不同专家权重Batch Size小时IO占比高MoE对高带宽内存HBM依赖性强训练收敛动态收敛曲线平滑稳定需充足训练计算量前期加速明显专家快速专长化但需负载均衡机制MoE在相同计算预算下可达到更低Loss可扩展上限受限于单卡显存与计算单元受限于总显存容量与通信带宽MoE通过稀疏性突破单设备算力限制二、基于Megatron的MoE模型训练技术方案2.1 训练框架2.1.1 Megatron框架Megatron框架是NVIDIA开发的高效分布式训练框架专为大规模Transformer模型设计。结合MoEMixture of Experts架构可实现参数量超大规模而计算量可控的高效训练。本文档基于Megatron-swift开源训练框架进行MoE模型训练的技术要点、关键参数及实践。2.1.2 Megatron-swift开源镜像环境镜像环境modelscope-registry.cn-hangzhou.cr.aliyuncs.com/modelscope-repo/modelscope:ubuntu22.04-cuda12.8.1-py311-torch2.8.0-vllm0.11.0-modelscope1.31.0-swift3.10.1关键依赖库Megatron-LM、transformer_engine、mcore-bridge2.2 分布式并行训练策略MoE模型训练通常需要3D并行策略在基于deepspeed的开源训练框架llama-factory、ms-swift的训练均出现训练卡住的bug张量并行 (TP): 张量切分方式分为按行进行切分和按列进行切分分别对应行并行Row Parallelism与列并行Column Parallelism。交叉进行维持通用矩阵的矩阵乘法all-reduce作为集体通讯从而将transformer层拆分到多个GPU。tensor_model_parallel_size模型张量并行度expert_model_parallel_size专家内部张量并行度只对专家切分上下游模型层张量不拆分微批次流水线并行 (PP):通过将传入的小批次minibatch分块为微批次microbatch并人为创建流水线来解决 GPU 空闲问题从而允许不同的 GPU 同时参与计算过程可以显著提升流水线并行设备利用率减小设备空闲状态的时间。pipeline_model_parallel_size : 流水线并行度micro_batch_size每个流水线的批次序列并行SP):将输入序列在不同GPU上分片处理只适用于长序列训练。sequence_parallel:序列并行度2.3 全局批次计算数据并行度计算数据并行大小 (DP) 总GPU数 / (TP × PP × SP)。需注意TP × PP × SP不能超过gpu总数另外PP的通信次数最短参数尽可能大有利于训练加速。global_batch_size/wbs全局批次大小等价于micro_batch_size*数据并行大小*梯度累加步数。2.4 mcore模型权重格式Mcore-Bridge兼容Dense/MoE/多模态等多种模型架构将safetensors 格式的模型权重转换为mcore格式让Megatron训练像transformers一样简单易用。2.5 MoE特定训练优化技术2.5.1 负载均衡训练MoE训练中的负载不均衡是训练动力学的必然结果。即使初始差异极小在正反馈循环的作用下经过数十个训练批次的迭代这些微小差异会指数级放大最终导致严重的专家利用不均甚至模型尺寸退化。为解决负载集中问题我们在总损失函数中引入负载均衡约束Ltotal Lmain λ⋅Laux其中Lmain是主任务损失如交叉熵Laux是辅助损失专门惩罚负载不均衡λ是平衡权重控制均衡约束的强度通常设为0.0010.0001或者0辅助损失函数为其中f为各专家实际处理的token个数占比P为各专家被路由选择的概率均值通过门控函数计算均衡时最小值为1.0。moe_aux_loss_coeff: 默认为0不使用aux_loss。通常情况下该值设置的越大训练效果越差但MoE负载越均衡请根据实验效果选择合适的值。moe_expert_capacity_factor: 专家容量因子超出专家容量的 token 会基于其被选中的概率被丢弃。可以令训练负载均匀提升训练速度避免专家过载sft微调中不建议使用会有信息熵的损失。如果是lora训练需留意lora参数配置target_modules可选all-linear all-router 。目标模块all-linear: 所有线性层all-router: MoE路由器lora_rankLoRA秩注意要避免r过大导致路由层lora块参数比模型路由层参数还要大的情况。2.5.2 计算-通信重叠优化overlap_grad_reduce: 重叠梯度reduce操作减少DP通信等待建议开启。overlap_param_gather: 重叠参数all-gather操作减少优化器通信成本建议开启。tp_comm_overlap: 重叠TP通信与GEMM计算建议开启。moe_permute_fusion启用token重排融合moe_grouped_gemm使用分组GEMM操作通过TransformerEngine优化多个专家同时计算moe_shared_expert_overlap共享专家计算与通信重叠2.5.3 其他MoE参数默认配置以下MoE参数默认与模型config配置一致即可num_experts: 专家数量moe_layer_freq: MoE层与Dense层的分布频率moe_router_topk: 每个token选择的专家数moe_shared_expert_intermediate_size: 共享专家大小moe_router_dtype: 路由计算数据类型moe_router_dtype fp32: 保持数值稳定性尤其在专家数多时训练案例nl2sql-dsl的训练案例大家可尝试修改模型和数据相关路径进行实践。训练脚本路径/AIOT-vePFS/sf_01441980/sf01441980/ms-swift/examples/train/megatron/moe/qwen3_moe.shPYTORCH_CUDA_ALLOC_CONFexpandable_segments:True\NPROC_PER_NODE8\CUDA_VISIBLE_DEVICES0,1,2,3,4,5,6,7\megatron sft\--model_type qwen3_nothinking\--model /AIOT-vePFS/sf_ModelZoo/Qwen3-30-A3B-2507\--load_safetensorstrue\--save_safetensorstrue\--merge_loratrue\--train_type lora\--target_modulesall-linear\all-router\--lora_rank64\--lora_alpha256\--dataset/AIOT-vePFS/sf_01441980/sf01441980/ms-swift/data/dsl/1218/dsl0_online_1027-1120.jsonl\/AIOT-vePFS/sf_01441980/sf01441980/ms-swift/data/dsl/1218/dsl0_v2_1120.jsonl\/AIOT-vePFS/sf_01441980/sf01441980/ms-swift/data/dsl/1218/dsl0_online_1027-1120.jsonl\/AIOT-vePFS/sf_01441980/sf01441980/ms-swift/data/dsl/1218/time_traindata.jsonl\/AIOT-vePFS/sf_01441980/sf01441980/ms-swift/data/dsl/1218/time_traindata.jsonl\/AIOT-vePFS/sf_01441980/sf01441980/ms-swift/data/dsl/1218/数据小精灵-20251103-1128eval-onlineresult-gptcheck_updated.jsonl\/AIOT-vePFS/sf_01441980/sf01441980/ms-swift/data/dsl/1218/数据小精灵-20251103-1128eval-onlineresult-gptcheck_updated.jsonl\/AIOT-vePFS/sf_01441980/sf01441980/ms-swift/data/dsl/1218/ratio_traindata_v2.jsonl\--pipeline_model_parallel_size2\--expert_model_parallel_size4\--tensor_model_parallel_size4\--moe_permute_fusiontrue\--moe_grouped_gemmtrue\--moe_shared_expert_overlaptrue\--moe_aux_loss_coeff0\--moe_shared_expert_overlaptrue\--micro_batch_size4\--global_batch_size256\--packingfalse\--log_interval1\--recompute_granularity full\--recompute_method uniform\--recompute_num_layers1\--max_epochs3\--finetunetrue\--cross_entropy_loss_fusiontrue\--lr 1e-6\--lr_decay_style cosine\--lr_warmup_fraction0.05\--save /AIOT-vePFS/sf_LLM/Fengyu-LLM/models/dsl/dsl0-Qwen3-30-A3B-2507-1222-64r-aux-0\--save_interval2000\--max_length8192\--num_workers8\--dataset_num_proc8\--no_save_optimtrue\--no_save_rngtrue\--sequence_paralleltrue\--attention_backend flash