Emerging Techniques Moe Training 训练

# MoE 训练指南 基于 DeepSpeed 官方文档和生产实践的完整训练指南。 ## 目录 - DeepSpeed MoE 设置 - 训练配置 - PR-MoE (金字塔残差 MoE) - Mixture-of-Students (MoS) - 超参数调优 - 生产环境训练 ## DeepSpeed MoE 设置 **来源**：DeepSpeed MoE 教程 (https://www.deepspeed.ai/tutorials/mixture-of-experts-nlg/) ### 要求 bash # 安装 DeepSpeed v0.6.0 或更高版本 pip install deepspeed>=0.6.0 # 克隆 Megatron-DeepSpeed git clone https://github.com/microsoft/Megatron-DeepSpeed cd Megatron-DeepSpeed pip install -r requirements.txt ### 基础 MoE 配置 { "train_batch_size": 256, "gradient_accumulation_steps": 1, "fp16": { "enabled": true, "loss_scale": 0, "initial_scale_power": 16 }, "moe": { "enabled": true, "num_experts": 128, "expert_parallel_size": 8, "moe_loss_coeff": 0.01, "train_capacity_factor": 1.25, "eval_capacity_factor": 2.0, "min_capacity": 4, "drop_tokens": true }, "zero_optimization": { "stage": 1 } } ## 训练参数 ### 核心 MoE 参数 **来自 DeepSpeed 文档：** 1. **`--num-experts`** - 每个 MoE 层的专家数量 - 推荐：128 个专家 - 范围：8-256，取决于规模 2. **`--moe-expert-parallel-size`** - 专家并行度 - 将专家分布在 GPU 上 - 示例：128 个专家 / 8 个 GPU = 每个 GPU 16 个专家 3. **`--moe-loss-coeff`** - MoE 辅助损失系数 - 推荐：0.01 - 控制负载均衡强度 4. **`--moe-train-capacity-factor`** - 训练容量乘数 - 默认：1.25 - 公式：capacity = (tokens/num_experts) × capacity_factor 5. **`--moe-eval-capacity-factor`** - 评估容量乘数 - 默认：2.0（评估时不丢弃 Token） 6. **`--moe-min-capacity`** - 最小专家容量 - 默认：4 - 确保每个专家处理最少 Token 数 7. **`--disable-moe-token-dropping`** - 移除专家容量限制 - 处理所有 Token（不丢弃） - 可能会增加内存使用 ### 训练脚本示例 bash #!/bin/bash deepspeed --num_gpus 8 pretrain_gpt_moe.py --tensor-model-parallel-size 1 --pipeline-model-parallel-size 1 --num-layers 24 --hidden-size 1024 --num-attention-heads 16 --seq-length 2048 --max-position-embeddings 2048 --micro-batch-size 4 --glo