Moe Training

orchestra-research/ai-research-skills/19-emerging-techniques/moe-training

作者 orchestra-research773a52944ba4MIT13K 个星标收录于 2026年10月8日更新于 2026年10月8日仓库3个月前更新

Train Mixture of Experts (MoE) models using DeepSpeed or HuggingFace. Use when training large-scale models with limited compute (5× cost reduction vs dense models), implementing sparse architectures like Mixtral 8x7B or DeepSeek-V3, or scaling model capacity without proportional compute increase. Covers MoE architectures, routing mechanisms, load balancing, expert parallelism, and inference optimization.

仅含说明AI & Agents
AI 生成的概览

指导使用 DeepSpeed 或 HuggingFace 训练与推理混合专家模型，涵盖路由、负载均衡与专家并行。

功能
说明如何构建和训练稀疏混合专家模型，包括 MoE 层代码、top-1、top-2 与专家选择路由、辅助负载均衡损失与路由 z-loss，以及专家并行配置。提供 DeepSpeed 与 HuggingFace 的安装命令、训练脚本和配置文件，并给出专家数量、容量因子、学习率与推理优化的建议。还指向架构、训练与推理的参考文件。
适用场景
适用于在算力有限时训练大模型、实现 Mixtral 8x7B 或 DeepSeek-V3 等稀疏架构，或在不按比例增加算力的前提下扩展模型容量。也适合调优 MoE 训练中的路由、负载均衡或专家并行。
运行要求
需要安装 deepspeed、transformers、torch 和 accelerate 的 Python 环境；专家并行通常需要多 GPU 硬件。不附带脚本，仅为说明文档与代码示例。

MoE Training: Mixture of Experts

When to Use This Skill

Use MoE Training when you need to:

  • Train larger models with limited compute (5× cost reduction vs dense models)
  • Scale model capacity without proportional compute increase
  • Achieve better performance per compute budget than dense models
  • Specialize experts for different domains/tasks/languages
  • Reduce inference latency with sparse activation (only 13B/47B params active in Mixtral)
  • Implement SOTA models like Mixtral 8x7B, DeepSeek-V3, Switch Transformers

Notable MoE Models: Mixtral 8x7B (Mistral AI), DeepSeek-V3, Switch Transformers (Google), GLaM (Google), NLLB-MoE (Meta)

Installation

bash
# DeepSpeed with MoE supportpip install deepspeed>=0.6.0
# Megatron-DeepSpeed for large-scale traininggit clone https://github.com/microsoft/Megatron-DeepSpeedcd Megatron-DeepSpeedpip install -r requirements.txt
# Alternative: HuggingFace Transformerspip install transformers accelerate

Quick Start

Basic MoE Architecture

python
import torchimport torch.nn as nn
class MoELayer(nn.Module):    """Sparse Mixture of Experts layer."""
    def __init__(self, hidden_size, num_experts=8, top_k=2):        super().__init__()        self.num_experts = num_experts        self.top_k = top_k
        # Expert networks (FFN)        self.experts = nn.ModuleList([            nn.Sequential(                nn.Linear(hidden_size, 4 * hidden_size),                nn.GELU(),                nn.Linear(4 * hidden_size, hidden_size)            )            for _ in range(num_experts)        ])
        # Gating network (router)        self.gate = nn.Linear(hidden_size, num_experts)
    def forward(self, x):        # x shape: (batch_size, seq_len, hidden_size)        batch_size, seq_len, hidden_size = x.shape
        # Flatten for routing        x_flat = x.view(-1, hidden_size)  # (batch_size * seq_len, hidden_size)
        # Compute gate scores        gate_logits = self.gate(x_flat)  # (batch_size * seq_len, num_experts)
        # Top-k routing        gate_scores = torch.softmax(gate_logits, dim=-1)        topk_scores, topk_indices = torch.topk(gate_scores, self.top_k, dim=-1)
        # Normalize top-k scores        topk_scores = topk_scores / topk_scores.sum(dim=-1, keepdim=True)
        # Dispatch and combine expert outputs        output = torch.zeros_like(x_flat)
        for i in range(self.top_k):            expert_idx = topk_indices[:, i]            expert_scores = topk_scores[:, i].unsqueeze(-1)
            # Route tokens to experts            for expert_id in range(self.num_experts):                mask = (expert_idx == expert_id)                if mask.any():                    expert_input = x_flat[mask]                    expert_output = self.experts[expert_id](expert_input)                    output[mask] += expert_scores[mask] * expert_output
        # Reshape back        return output.view(batch_size, seq_len, hidden_size)

DeepSpeed MoE Training

bash
# Training script with MoEdeepspeed pretrain_gpt_moe.py \  --num-layers 24 \  --hidden-size 1024 \  --num-attention-heads 16 \  --seq-length 2048 \  --max-position-embeddings 2048 \  --micro-batch-size 4 \  --global-batch-size 256 \  --train-iters 500000 \  --lr 0.0001 \  --min-lr 0.00001 \  --lr-decay-style cosine \  --num-experts 128 \  --moe-expert-parallel-size 4 \  --moe-loss-coeff 0.01 \  --moe-train-capacity-factor 1.25 \  --moe-eval-capacity-factor 2.0 \  --fp16 \  --deepspeed_config ds_config.json

Core Concepts

1. MoE Architecture

Key Components:

  • Experts: Multiple specialized FFN networks (typically 8-128)
  • Router/Gate: Learned network that selects which experts to use
  • Top-k Routing: Activate only k experts per token (k=1 or k=2)
  • Load Balancing: Ensure even expert utilization
Input Token    ↓Router (Gate Network)    ↓Top-k Expert Selection (e.g., 2 out of 8)    ↓Expert 1 (weight: 0.6) + Expert 5 (weight: 0.4)    ↓Weighted Combination    ↓Output

2. Routing Mechanisms

Top-1 Routing (Switch Transformer):

python
# Simplest routing: one expert per tokengate_logits = router(x)  # (batch, seq_len, num_experts)expert_idx = torch.argmax(gate_logits, dim=-1)  # Hard routing

Top-2 Routing (Mixtral):

python
# Top-2: two experts per tokengate_scores = torch.softmax(router(x), dim=-1)top2_scores, top2_indices = torch.topk(gate_scores, k=2, dim=-1)
# Normalize scorestop2_scores = top2_scores / top2_scores.sum(dim=-1, keepdim=True)
# Combine expert outputsoutput = (top2_scores[:, :, 0:1] * expert_outputs[top2_indices[:, :, 0]] +          top2_scores[:, :, 1:2] * expert_outputs[top2_indices[:, :, 1]])

Expert Choice Routing:

python
# Experts choose top-k tokens (instead of tokens choosing experts)# Guarantees perfect load balancingexpert_scores = router(x).transpose(-1, -2)  # (batch, num_experts, seq_len)topk_tokens = torch.topk(expert_scores, k=capacity_per_expert, dim=-1)

3. Load Balancing

Auxiliary Loss:

python
def load_balancing_loss(gate_logits, expert_indices, num_experts):    """Encourage uniform expert usage."""    # Fraction of tokens routed to each expert    expert_counts = torch.bincount(expert_indices.flatten(), minlength=num_experts)    expert_fraction = expert_counts.float() / expert_indices.numel()
    # Gate probability for each expert (average across tokens)    gate_probs = torch.softmax(gate_logits, dim=-1).mean(dim=0)
    # Auxiliary loss: encourage alignment    aux_loss = num_experts * (expert_fraction * gate_probs).sum()
    return aux_loss
# Add to main losstotal_loss = language_model_loss + 0.01 * load_balancing_loss(...)

Router Z-Loss (Stability):

python
def router_z_loss(logits):    """Encourage router to have lower entropy (more decisive)."""    z_loss = torch.logsumexp(logits, dim=-1).pow(2).mean()    return z_loss
total_loss = lm_loss + 0.01 * aux_loss + 0.001 * router_z_loss(gate_logits)

4. Expert Parallelism

python
# DeepSpeed configuration{  "train_batch_size": 256,  "fp16": {"enabled": true},  "moe": {    "enabled": true,    "num_experts": 128,    "expert_parallel_size": 8,  # Distribute 128 experts across 8 GPUs    "capacity_factor": 1.25,    # Expert capacity = tokens_per_batch * capacity_factor / num_experts    "drop_tokens": true,        # Drop tokens exceeding capacity    "use_residual": false  }}

Training Configuration

DeepSpeed MoE Config

json
{  "train_batch_size": 256,  "gradient_accumulation_steps": 1,  "optimizer": {    "type": "Adam",    "params": {      "lr": 0.0001,      "betas": [0.9, 0.999],      "eps": 1e-8    }  },  "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,    "use_residual": false,    "use_tutel": false  },  "zero_optimization": {    "stage": 1  }}

Training Script

bash
#!/bin/bash
# Mixtral-style MoE trainingdeepspeed --num_gpus 8 pretrain_moe.py \  --model-parallel-size 1 \  --num-layers 32 \  --hidden-size 4096 \  --num-attention-heads 32 \  --seq-length 2048 \  --max-position-embeddings 4096 \  --micro-batch-size 2 \  --global-batch-size 256 \  --train-iters 500000 \  --save-interval 5000 \  --eval-interval 1000 \  --eval-iters 100 \  --lr 0.0001 \  --min-lr 0.00001 \  --lr-decay-style cosine \  --lr-warmup-iters 2000 \  --clip-grad 1.0 \  --weight-decay 0.1 \  --num-experts 8 \  --moe-expert-parallel-size 4 \  --moe-loss-coeff 0.01 \  --moe-train-capacity-factor 1.25 \  --moe-eval-capacity-factor 2.0 \  --disable-moe-token-dropping \  --fp16 \  --deepspeed \  --deepspeed_config ds_config_moe.json \  --data-path /path/to/data \  --vocab-file /path/to/vocab.json \  --merge-file /path/to/merges.txt

Advanced Patterns

Mixtral 8x7B Architecture

python
class MixtralMoEBlock(nn.Module):    """Mixtral-style MoE block with 8 experts, top-2 routing."""
    def __init__(self, config):        super().__init__()        self.hidden_dim = config.hidden_size        self.ffn_dim = config.intermediate_size        self.num_experts = config.num_local_experts  # 8        self.top_k = config.num_experts_per_tok       # 2
        # 8 expert FFNs        self.experts = nn.ModuleList([            nn.Sequential(                nn.Linear(self.hidden_dim, self.ffn_dim, bias=False),                nn.SiLU(),                nn.Linear(self.ffn_dim, self.hidden_dim, bias=False)            )            for _ in range(self.num_experts)        ])
        # Router        self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False)
    def forward(self, hidden_states):        batch_size, sequence_length, hidden_dim = hidden_states.shape
        # Flatten        hidden_states = hidden_states.view(-1, hidden_dim)
        # Router logits        router_logits = self.gate(hidden_states)  # (batch * seq_len, num_experts)
        # Softmax and top-2        routing_weights = torch.softmax(router_logits, dim=1)        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
        # Normalize routing weights        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
        # Initialize output        final_hidden_states = torch.zeros_like(hidden_states)
        # Route to experts        for expert_idx in range(self.num_experts):            expert_layer = self.experts[expert_idx]            idx, top_x = torch.where(selected_experts == expert_idx)
            if idx.shape[0] == 0:                continue
            # Current expert tokens            current_hidden_states = hidden_states[idx]
            # Expert forward            current_hidden_states = expert_layer(current_hidden_states)
            # Weighted by routing scores            current_hidden_states *= routing_weights[idx, top_x, None]
            # Accumulate            final_hidden_states.index_add_(0, idx, current_hidden_states)
        # Reshape        return final_hidden_states.view(batch_size, sequence_length, hidden_dim)

PR-MoE (Pyramid-Residual-MoE)

bash
# DeepSpeed PR-MoE: 3x better parameter efficiencydeepspeed pretrain_gpt_moe.py \  --num-layers 24 \  --hidden-size 1024 \  --num-attention-heads 16 \  --num-experts "[128, 64, 32, 16]" \  --mlp-type residual \  --moe-expert-parallel-size 4 \  --moe-loss-coeff 0.01 \  --fp16

Best Practices

1. Expert Count Selection

python
# Rule of thumb: More experts = more capacity, but diminishing returns# Typical configurations:# - Small models (1B-7B): 8-16 experts# - Medium models (7B-30B): 8-64 experts# - Large models (30B+): 64-256 experts
# Example: Mixtral 8x7B# Total params: 47B (8 experts × 7B each)# Active params: 13B (2 experts × 7B, top-2 routing)# Efficiency: 47B capacity with 13B compute

2. Capacity Factor Tuning

python
# Capacity = (tokens_per_batch / num_experts) * capacity_factor
# Training: Lower capacity (faster, drops some tokens)train_capacity_factor = 1.25  # 25% buffer
# Evaluation: Higher capacity (no dropping)eval_capacity_factor = 2.0    # 100% buffer
# Formula:expert_capacity = int((seq_len * batch_size / num_experts) * capacity_factor)

3. Learning Rate Guidelines

python
# MoE models need lower LR than dense models# - Dense model: lr = 6e-4# - MoE model: lr = 1e-4 (3-6× lower)
# Also extend decay scheduledense_lr_decay_iters = 300000moe_lr_decay_iters = 500000  # 1.5-2× longer

4. Loss Coefficient Tuning

python
# Start with standard valuesmoe_loss_coeff = 0.01    # Auxiliary loss (load balancing)router_z_loss_coeff = 0.001  # Router entropy (stability)
# If load imbalance persists, increase aux lossif max_expert_usage / min_expert_usage > 2.0:    moe_loss_coeff = 0.1  # Stronger load balancing
# If training unstable, increase z-lossif grad_norm > 10.0:    router_z_loss_coeff = 0.01

5. Avoid Common Pitfalls

python
# ❌ Bad: Using same LR as dense modeloptimizer = Adam(model.parameters(), lr=6e-4)
# ✅ Good: Lower LR for MoEoptimizer = Adam([    {'params': model.non_moe_params, 'lr': 6e-4},    {'params': model.moe_params, 'lr': 1e-4}])
# ❌ Bad: No load balancingloss = lm_loss
# ✅ Good: Add auxiliary lossloss = lm_loss + 0.01 * aux_loss + 0.001 * z_loss
# ❌ Bad: Too many experts for small datasetnum_experts = 128  # Overfitting risk
# ✅ Good: Match experts to data diversitynum_experts = 8  # Better for small datasets

Inference Optimization

Sparse Inference

python
# Only activate top-k experts (huge memory savings)@torch.no_grad()def moe_inference(x, model, top_k=2):    """Sparse MoE inference: only load k experts."""    # Router    gate_logits = model.gate(x)    topk_scores, topk_indices = torch.topk(        torch.softmax(gate_logits, dim=-1),        k=top_k,        dim=-1    )
    # Load and run only top-k experts    output = torch.zeros_like(x)    for i in range(top_k):        expert_idx = topk_indices[:, i]        # Load expert from disk/offload if needed        expert = model.load_expert(expert_idx)        output += topk_scores[:, i:i+1] * expert(x)
    return output

Resources

See Also

  • references/architectures.md - MoE model architectures (Mixtral, Switch, DeepSeek-V3)
  • references/training.md - Advanced training techniques and optimization
  • references/inference.md - Production deployment and serving patterns

来源与署名

来源：orchestra-research/ai-research-skills位于19-emerging-techniques/moe-training提交773a529

许可证： MIT

内容归原作者所有。SourceWeft 从公开仓库中收录这些内容。

举报或申请下架

更多来自 orchestra-research/ai-research-skills 的技能

Speculative Decoding

orchestra-research

指导使用推测解码、Medusa 多头预测和前瞻解码来加速大语言模型推理。

AI & Agents13K3个月前更新

Systems Paper Writing

orchestra-research

为 OSDI、SOSP、ASPLOS、NSDI 等系统会议论文写作提供段落级结构与篇幅分配指导。

Writing & Content13K3个月前更新

Presenting Conference Talks

orchestra-research

Generates conference presentation slides (Beamer LaTeX PDF and editable PPTX) from a compiled paper with speaker notes and talk script. Use when preparing oral talks, spotlight presentations, or invited talks for ML and systems conferences.

待分类13K3个月前更新

Model Pruning

orchestra-research

指导使用 Wanda、SparseGPT、幅度剪枝和 N:M 方法压缩 LLM，减小模型体积并加速推理。

AI & Agents13K3个月前更新

Outlines

orchestra-research

指导使用 Outlines 库，通过本地模型进行受结构约束的文本生成。

AI & Agents13K3个月前更新

Ml Paper Writing

orchestra-research

指导撰写可发表的机器学习/人工智能会议论文，从研究仓库到 LaTeX 投稿定稿。

Writing & Content13K3个月前更新

更多AI & Agents技能

Ppt Template Creator

anthropics

精选

将用户的 PowerPoint 模板转化为可复用技能，用于生成品牌演示文稿。

AI & Agents39K2周前更新

Skill Development

anthropics

精选

指导创建 Claude Code 插件技能，涵盖结构、描述、渐进式披露与验证。

AI & Agents37K今天更新

Plugin Structure

anthropics

精选

指导 Claude Code 插件的结构、清单与组件布局。

AI & Agents37K今天更新

Command Development

anthropics

精选

指导创建 Claude Code 斜杠命令，涵盖结构、YAML frontmatter、参数与插件功能。

AI & Agents37K今天更新

Claude Md Improver

anthropics

精选

审查仓库中的 CLAUDE.md 文件，评估其质量，并在获得批准后应用有针对性的改进。

AI & Agents37K今天更新

Find Skills

vercel-labs

精选

指导通过 npx skills 命令行工具发现、评估并安装智能体技能。

AI & Agents33K今天更新