moe-training

MoE トレーニング：Mixture of Experts

このスキルを使用する場面

次のような場合に MoE トレーニングを使用します。

• 限られた計算リソースでより大きなモデルをトレーニングする（密なモデルと比較してコストを 5 分の 1 に削減）
• 計算量を比例して増やさずにモデルの容量を拡張する
• 密なモデルよりも計算予算あたりの性能を向上させる
• 異なるドメイン／タスク／言語向けにエキスパートを特化させる
• スパースな活性化により推論レイテンシを短縮する（Mixtral では 130 億／470 億パラメーターのみがアクティブ）
• Mixtral 8x7B、DeepSeek-V3、Switch Transformers などの SOTA モデルを実装する

主な MoE モデル：Mixtral 8x7B（Mistral AI）、DeepSeek-V3、Switch Transformers（Google）、GLaM（Google）、NLLB-MoE（Meta）

インストール

# DeepSpeed with MoE support
pip install deepspeed>=0.6.0

# Megatron-DeepSpeed for large-scale training
git clone https://github.com/microsoft/Megatron-DeepSpeed
cd Megatron-DeepSpeed
pip install -r requirements.txt

# Alternative: HuggingFace Transformers
pip install transformers accelerate

クイックスタート

基本的な MoE アーキテクチャ

import torch
import 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 script with MoE
deepspeed 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

主な概念

1. MoE アーキテクチャ

主な構成要素：

• エキスパート：複数の特化型 FFN ネットワーク（通常 8～128 個）
• ルーター／ゲート：使用するエキスパートを選択する学習済みネットワーク
• Top-k ルーティング：トークンごとに k 個のエキスパートのみをアクティブ化（k=1 または k=2）
• 負荷分散：エキスパートが均等に利用されるようにすること

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. ルーティングの仕組み

Top-1 ルーティング（Switch Transformer）：

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

Top-2 ルーティング（Mixtral）：

# Top-2: two experts per token
gate_scores = torch.softmax(router(x), dim=-1)
top2_scores, top2_indices = torch.topk(gate_scores, k=2, dim=-1)

# Normalize scores
top2_scores = top2_scores / top2_scores.sum(dim=-1, keepdim=True)

# Combine expert outputs
output = (top2_scores[:, :, 0:1] * expert_outputs[top2_indices[:, :, 0]] +
          top2_scores[:, :, 1:2] * expert_outputs[top2_indices[:, :, 1]])

エキスパート選択ルーティング：

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

3. 負荷分散

補助損失：

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 loss
total_loss = language_model_loss + 0.01 * load_balancing_loss(...)

ルーター Z 損失（安定性）：

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. エキスパート並列処理

# 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
  }
}

トレーニング設定

DeepSpeed MoE 設定

{
  "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
  }
}

トレーニングスクリプト

#!/bin/bash

# Mixtral-style MoE training
deepspeed --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

高度なパターン

詳細、コード例、パラメーター表については、references/advanced-patterns.md を参照してください。この手順が該当する場合に読み進めてください。

ベストプラクティス

1. エキスパート数の選択

# 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 = (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. 学習率のガイドライン

# MoE models need lower LR than dense models
# - Dense model: lr = 6e-4
# - MoE model: lr = 1e-4 (3-6× lower)

# Also extend decay schedule
dense_lr_decay_iters = 300000
moe_lr_decay_iters = 500000  # 1.5-2× longer

4. 損失係数の調整

# Start with standard values
moe_loss_coeff = 0.01    # Auxiliary loss (load balancing)
router_z_loss_coeff = 0.001  # Router entropy (stability)

# If load imbalance persists, increase aux loss
if max_expert_usage / min_expert_usage > 2.0:
    moe_loss_coeff = 0.1  # Stronger load balancing

# If training unstable, increase z-loss
if grad_norm > 10.0:
    router_z_loss_coeff = 0.01

5. よくある落とし穴を避ける

# ❌ Bad: Using same LR as dense model
optimizer = Adam(model.parameters(), lr=6e-4)

# ✅ Good: Lower LR for MoE
optimizer = Adam([
    {'params': model.non_moe_params, 'lr': 6e-4},
    {'params': model.moe_params, 'lr': 1e-4}
])

# ❌ Bad: No load balancing
loss = lm_loss

# ✅ Good: Add auxiliary loss
loss = lm_loss + 0.01 * aux_loss + 0.001 * z_loss

# ❌ Bad: Too many experts for small dataset
num_experts = 128  # Overfitting risk

# ✅ Good: Match experts to data diversity
num_experts = 8  # Better for small datasets

推論の最適化

スパース推論

# 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

リソース

• DeepSpeed MoE チュートリアル：https://www.deepspeed.ai/tutorials/mixture-of-experts-nlg/
• Mixtral 論文：https://arxiv.org/abs/2401.04088
• Switch Transformers：https://arxiv.org/abs/2101.03961
• HuggingFace MoE ガイド：https://huggingface.co/blog/moe
• NVIDIA MoE ブログ：https://developer.nvidia.com/blog/applying-mixture-of-experts-in-llm-architectures/

関連項目

• references/architectures.md - MoE モデルのアーキテクチャ（Mixtral、Switch、DeepSeek-V3）
• references/training.md - 高度なトレーニング手法と最適化
• references/inference.md - 本番環境へのデプロイとサービングのパターン

エージェントの作業手順

1. 環境を確認します。 GPU メモリ、モデルのチェックポイントとライセンス、フレームワークのバージョン、入力形式を確認します。
2. 入力を特定します。 データまたはユーザーから形式、識別子、パラメーターを確認します。結果に影響する値は推測せず、質問してください。
3. まず小規模に実行します。 最初に、単一のサンプルを低解像度または短い長さで推論します。
4. 上記の手順と参考情報に従って、タスク全体を実行します。
5. 結果を検証します。 出力を定性的に確認し、タスクの指標で評価します。生成コンテンツのライセンス制限も確認してください。
6. 報告します。 実行内容（バージョン、コマンド、パラメーター）、確認した内容、未確定事項を明記します。

このような場合	対応方法
大きな入力でメモリ不足が発生する	解像度またはシーケンス長を下げる、半精度を使用する、または入力を分割します。
この手順に記載されている関数、フラグ、またはエンドポイントがインストール済みのバージョンにない	インストール済みバージョンのドキュメント（help()、--help、公式ドキュメント）を確認し、適宜対応して、ユーザーに伝えてください。API を決して捏造しないでください。
必須の入力、識別子、またはパラメーターが曖昧である	ユーザーに質問するか、実行前に前提を明示します。

誠実性に関するルール

• 結果、パラメーター、識別子、引用、統計を決して捏造しないでください。実行または検証できない場合は、その旨を明確に伝えてください。
• 商用利用または公開利用の前に、モデルとデータセットのライセンスを確認してください。
• ここに記載されているバージョン固有の詳細は古くなっている可能性があるため、インストール済みバージョンの公式ドキュメントと照合してください。
• 費用が発生する、共有 GPU やクラウドの利用枠を消費する、個人データや患者データに触れる、または元に戻せない操作を行う前に、許可を求めてください。