OLMoE 完全ガイド

🧪 OLMoE-1B-7B-Instruct — 全体アーキテクチャ

AllenAI (Allen Institute for AI) 製の完全オープン MoE。訓練データ・コード・中間チェックポイント全て公開。 このモデルは あなたが実際に LoRA 学習できる唯一の MoE。

⚖ Qwen3-30B との比較

🔄 推論フロー

入力 "hello"
  │
  ▼  [① Tokenize] BPE で 50,304 語彙に分割
  │    "hello" → [19659]   ← 1 トークンに収まる (英語が得意)
  ▼
[② Embedding]  token_embd [2048 × 50304]
  │  各 token → 2048 次元ベクトル
  ▼
┌─────── 16 層ブロックをループ (Qwen3 は 48) ──────┐
│                                                      │
│ [③ Self-Attention (MHA)]                      │
│    q/k/v 全て 16 head                                │
│    ※ GQA なし → KV cache が Qwen3 より大きい         │
│                                                      │
│ [④ Router]  ffn_gate_inp [2048 × 64]          │
│    2048 次元 → 64 スコア → softmax                   │
│                                                      │
│ [⑤ top-8 選出]                                │
│    64 人から最適な 8 人を選ぶ                         │
│                                                      │
│ [⑥ Expert 8 人実行]                           │
│    選ばれた 8 expert だけ FFN (1024 dim)             │
│    残り 56 expert は計算しない = 高速                │
│                                                      │
│ [⑦ 重み付き合成]                              │
│    8 人の出力 × routing_weight で足し合わせ          │
└───────────────────────────────────────────────────┘
  │
  ▼  [⑧ Final Norm + LM Head]
  │    2048 次元 → 50,304 次元ロジット
  ▼  [⑨ Sampling] top-p=0.9 で次 token 決定
  │
"お"  次のトークン!

📖 OLMoE ソースコードの歩き方 (1500 行)

HuggingFace transformers の modeling_olmoe.py は 1,500 行しかない。 MoE の実装を読む上で最短の教材。読む順番をクラス単位で示します。

🔬 MoE の心臓 70 行を擬似コードで

class OlmoeSparseMoeBlock:
    # 入力: [B, T, 2048]   ← T トークン分の隠れ状態
    # 出力: [B, T, 2048]   ← MoE 処理後の隠れ状態

    def forward(self, x):
        B, T, H = x.shape
        x_flat = x.view(B*T, H)                     # [B*T, 2048]

        ### ① Router: 64 expert のスコアを計算 ###
        router_logits = self.gate(x_flat)           # [B*T, 64]
        weights       = softmax(router_logits, -1)  # 確率化

        ### ② top-K: 上位 8 人を選ぶ ###
        w, idx = torch.topk(weights, k=8, dim=-1)   # idx: [B*T, 8]

        ### ③ 選ばれた expert だけ実行 ###
        out = torch.zeros_like(x_flat)
        for e_id in range(64):                      # 64 expert 全部ループ
            mask = (idx == e_id)                    # この expert が選ばれた位置
            if not mask.any():
                continue                            # 呼ばれなければ skip (高速化)

            tok_idx, topk_idx = mask.nonzero(as_tuple=True)
            expert_out = self.experts[e_id](x_flat[tok_idx])  # FFN 実行

            ### ④ 重み付きで元の位置に足し戻す ###
            out.index_add_(0, tok_idx,
                           expert_out * w[tok_idx, topk_idx, None])

        return out.view(B, T, H), router_logits

# ★ この 30 行が MoE の全て。Mixtral も Qwen3 も同じパターン。expert 数 (64 or 128) と top-K (8 or 2) が違うだけ。

🔍 部品別の役割 — 何がどこで何をしているか

OLMoE-1B-7B を動かす 8 つの部品と、その中で「どこに何が入っているか」を表で整理。「無いとどうなる?」列が初心者の理解に効きます。

🔤 語彙 50,304 個はどこ?

「50,304」は OLMoE が知っている語彙の種類数。3 つの場所に散らばっています。

📊 4.2 GB の内訳 (Q4_K_M)

OLMoE-1B-7B-0125-Instruct-Q4_K_M.gguf  (4.2 GB)
├── ヘッダ                 ~1 KB    // "GGUF" magic + version
├── メタデータ             ~2 MB    // tokenizer + architecture info
│   ├── tokenizer.ggml.tokens          // 50304 文字列
│   ├── tokenizer.ggml.merges          // BPE 規則
│   └── architecture="olmoe", layers=16, experts=64, …
│
├── テンソル本体          ~4.2 GB
│   │
│   ├── token_embd.weight          52 MB   // 埋め込み
│   ├── output.weight              52 MB   // LM head
│   │
│   ├── blk.0 〜 blk.15 (16 層)
│   │   ├── attn_q/k/v/output         17 MB × 16  // Attention
│   │   ├── attn_norm / ffn_norm          8 KB × 16
│   │   │
│   │   ├── ffn_gate_inp              0.5 MB × 16  // Router ★小さい
│   │   ├── ffn_gate_exps (×64)        50 MB × 16   // Expert gate
│   │   ├── ffn_up_exps   (×64)        50 MB × 16   // Expert up
│   │   └── ffn_down_exps (×64)        50 MB × 16   // Expert down
│   │
│   └── output_norm.weight             8 KB

💪 強くする方法の選び方 — 7 つの手法と意思決定

「モデルを賢くする」には目的別に 7 つの選択肢。何をしたいかで選び方が変わる。まず自分のゴールを決めてから手法を選ぶこと。

🌲 意思決定フローチャート

Q1: 新しい知識を追加したい? それとも style/振る舞いだけ?
  │
  ├─ style だけ (敬語変更 / キャラ付け / 分野特化返答)
  │       │
  │       └─→ LoRA + SFT
  │           (100〜1000 サンプル、3〜6 時間、VRAM 6GB)
  │           ★ 最小コスト、最初に試すべき ★
  │
  ├─ 知識追加 (ドメイン特化・専門用語)
  │       │
  │       ├─ データが 1 万件未満  →  LoRA + SFT (Q&A 形式)
  │       ├─ データが 10 万件以上 →  Continued Pre-training (クラウド推奨)
  │       └─ 最新情報で常に更新したい → RAG (検索連携) が正解
  │
  ├─ 応答の品質/安全性を高めたい
  │       │
  │       └─→ DPO or ORPO
  │           (選好ペア 1000 組、VRAM 8GB)
  │
  └─ ChatGPT レベルを自作したい
          │
          └─→ RLHF (PPO)
              → ローカル GPU 不可能、クラウドで 100 万円規模

🏗 プロ級の組み合わせレシピ

🇯🇵 日本語データセット集

OLMoE を日本語で強化するための実際に使える日本語データセット 10 種類。 全て HuggingFace datasets ライブラリで 1 行ダウンロード可能。

📗 整形レシピ ① Dolly-15k-ja

指示→回答 15,015 件。この 1 つだけで初回 LoRA は完走できる。

# prepare_dolly_ja.py
from datasets import load_dataset
import json

ds = load_dataset("kunishou/databricks-dolly-15k-ja", split="train")

with open("my_data.jsonl", "w", encoding="utf-8") as f:
    for row in ds:
        user_content = row["instruction"]
        if row.get("input"):
            user_content += "\n" + row["input"]
        messages = [
            {"role": "user",      "content": user_content},
            {"role": "assistant", "content": row["output"]},
        ]
        f.write(json.dumps({"messages": messages}, ensure_ascii=False) + "\n")

print("✅ my_data.jsonl done:", len(ds), "件")
# → ✅ my_data.jsonl done: 15015 件

🎓 LoRA 学習チュートリアル — RTX 4070 Super 完全手順

あなたの RTX 4070 Super 12GB で実際に動く完全手順。 所要時間 3〜6 時間 (データセット次第)。 OLMoE は個人が本格 fine-tune できる唯一の MoE。

① 環境構築 (5 分)

Windows の場合は WSL2 Ubuntu 推奨。ネイティブ Windows も可だが bitsandbytes で詰まりやすい。

# venv を作る (Python 3.11 推奨)
python3.11 -m venv ~/olmoe-env
source ~/olmoe-env/bin/activate     # Windows: .\olmoe-env\Scripts\activate

# PyTorch (CUDA 12.1 版)
pip install torch --index-url https://download.pytorch.org/whl/cu121

# 方式 A: Unsloth (速い。推奨)
pip install "unsloth[cu121-torch230] @ git+https://github.com/unslothai/unsloth.git"

# 方式 B: 素の HF transformers + PEFT + bitsandbytes
pip install transformers peft datasets accelerate bitsandbytes trl

# 動作確認
python -c "import torch; print('CUDA:', torch.cuda.is_available(), torch.cuda.get_device_name(0))"
# → CUDA: True NVIDIA GeForce RTX 4070 SUPER

③ 学習スクリプト (train_lora.py)

Unsloth 版。コピペでそのまま動く。target_modules は Expert FFN + Attention を指定。

# train_lora.py
from unsloth import FastLanguageModel
from datasets import load_dataset
from trl import SFTTrainer, SFTConfig

# ① モデル読み込み (4-bit 量子化で VRAM 4GB)
model, tok = FastLanguageModel.from_pretrained(
    model_name = "allenai/OLMoE-1B-7B-0125-Instruct",
    max_seq_length = 2048,
    load_in_4bit = True,
    dtype = None,  # 自動 (bf16 on Ampere+)
)

# ② LoRA 設定 (target_modules に expert を含める)
model = FastLanguageModel.get_peft_model(
    model,
    r = 16,               # LoRA rank (8〜32 が定石)
    lora_alpha = 32,
    lora_dropout = 0.05,
    target_modules = [
        "q_proj", "k_proj", "v_proj", "o_proj",  # Attention
        "gate_proj", "up_proj", "down_proj",       # ★ Expert FFN ★
    ],
    use_gradient_checkpointing = "unsloth",
    random_state = 3407,
)

# ③ データロード (JSONL → chat template)
def fmt(sample):
    return {"text": tok.apply_chat_template(sample["messages"], tokenize=False)}

ds = load_dataset("json", data_files="my_data.jsonl", split="train").map(fmt)

# ④ 学習設定
trainer = SFTTrainer(
    model = model,
    tokenizer = tok,
    train_dataset = ds,
    dataset_text_field = "text",
    max_seq_length = 2048,
    args = SFTConfig(
        output_dir = "./olmoe-lora-out",
        per_device_train_batch_size = 2,
        gradient_accumulation_steps = 4,
        num_train_epochs = 3,
        learning_rate = 2e-4,
        warmup_steps = 10,
        logging_steps = 5,
        save_steps = 100,
        bf16 = True,
        optim = "adamw_8bit",
    ),
)

# ⑤ 実行!
trainer.train()
model.save_pretrained("./olmoe-lora-out")

⑤ LoRA マージ & GGUF 変換

# ① LoRA をベースにマージ
python -c "
from unsloth import FastLanguageModel
model, tok = FastLanguageModel.from_pretrained('./olmoe-lora-out', load_in_4bit=False)
model.save_pretrained_merged('./olmoe-merged', tok, save_method='merged_16bit')
"

# ② GGUF 変換 (llama.cpp の convert スクリプト)
git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp && make LLAMA_CUDA=1

python convert_hf_to_gguf.py ../olmoe-merged \
    --outfile olmoe-mytuned.gguf \
    --outtype f16

# ③ Q4_K_M 量子化 (4.2 GB に圧縮)
./llama-quantize olmoe-mytuned.gguf olmoe-mytuned-Q4_K_M.gguf Q4_K_M

# ④ Ollama 登録
ollama create my-olmoe -f Modelfile
ollama run my-olmoe "自己紹介して"

💎 MoE 特有のチューニング Tips

🚨 よくあるエラーと対処

🔬 Expert 作成・強化 — MoE 特有の深い改造 (上級)

MoE 特有の改造術。単なる LoRA を超えて、Expert そのものを操作する方法。 Expert を新規追加・特定 expert だけ強化・不要 expert を剪定。フルオープンの OLMoE だからこそ可能。

🧬 Expert 1 個の中身 (SwiGLU FFN)

Expert i の forward (入力 x [2048] → 出力 [2048])

        x  [2048]
        │
        ├──→  gate_proj[i]  ──→  gate  [1024]
        │     (重み [2048→1024])
        │                          │
        ├──→  up_proj[i]    ──→  up    [1024]
        │     (重み [2048→1024])    │
        │                           ▼
        │           silu(gate) * up  ── hidden [1024]
        │                           │
        │     down_proj[i]    │
        │     (重み [1024→2048])    ▼
        └───────────────────── output [2048]

# OLMoE の場合: Expert 1 個 = 3 × 2048 × 1024 = 6.3M パラ
# × 64 expert × 16 層 = 6.5 B パラ (全体の 90%)

🔧 改造① 特定 Expert だけ強化する

「数学が苦手なモデル」を直したい → 数学で発火する expert を特定して LoRA をかぶせる。

# Step 1: どの expert が何で発火してるか調べる
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model = AutoModelForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-0125-Instruct",
                                             output_router_logits=True)
tok = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0125-Instruct")

prompts = ["2+2=?", "Solve x^2=4", "derivative of sin(x)"]
expert_count = torch.zeros(16, 64)  # 16 層 × 64 expert

for p in prompts:
    inputs = tok(p, return_tensors="pt")
    with torch.no_grad():
        out = model(**inputs)
    for layer, logits in enumerate(out.router_logits):
        topk_idx = logits.topk(8, dim=-1).indices
        for e in topk_idx.flatten():
            expert_count[layer, e] += 1

top_math_experts = expert_count.view(-1).topk(20).indices
print("数学エキスパート候補:", [(i//64, i%64) for i in top_math_experts])

# Step 2: その expert だけに LoRA
from peft import LoraConfig, get_peft_model

cfg = LoraConfig(
    r=32,
    target_modules=["gate_proj","up_proj","down_proj"],
    layers_to_transform=[8,9,10],          # 数学で発火した層だけ
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, cfg)
# → 数学データで SFT 実行

🧪 改造② 新しい Expert を追加する (上級)

OLMoE は 64 expert。これを 80 にしたい → Router と Expert テンソルをコード側で拡張。

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-0125-Instruct")
config = model.config
NEW_N = 80                                 # 64 → 80 に増やす

for i, layer in enumerate(model.model.layers):
    mlp = layer.mlp  # OlmoeSparseMoeBlock

    # ① Router (ffn_gate_inp) 拡張 [2048 × 64] → [2048 × 80]
    old_w = mlp.gate.weight.data     # [64, 2048]
    new_w = torch.zeros(NEW_N, old_w.size(1))
    new_w[:64] = old_w
    new_w[64:] = old_w[:16].clone()        # 新 16 個は既存の複製
    mlp.gate = torch.nn.Linear(2048, NEW_N, bias=False)
    mlp.gate.weight.data = new_w

    # ② Expert (ModuleList) を 16 個追加
    for j in range(16):
        src = mlp.experts[j % 64]
        import copy
        mlp.experts.append(copy.deepcopy(src))

    mlp.num_experts = NEW_N

config.num_experts = NEW_N
model.save_pretrained("./olmoe-80experts")
# ★ 続いて SFT で新 expert を分化させる学習が必要

✂ 改造③ 使われていない Expert を削除 (モデル圧縮)

# ② 各層で下位 16 個を削除 (64 → 48 expert)
for layer_idx, layer in enumerate(model.model.layers):
    mlp = layer.mlp
    usage = expert_count[layer_idx]              # [64]
    keep_idx = usage.topk(48).indices             # 上位 48 残す

    mlp.experts = torch.nn.ModuleList([mlp.experts[i] for i in keep_idx])
    mlp.gate.weight.data = mlp.gate.weight.data[keep_idx]
    mlp.num_experts = 48

# ③ 再微調整 (剪定の穴埋め)
# → SFT を軽く回す (2000 サンプル程度)

# 結果: 7B → 5.4B (25% 削減)、品質はほぼ維持

💎 Expert 改造のベストプラクティス