frankenmoe / AGENT_GUIDE.md

Add AGENT_GUIDE — reproducible MoE pipeline for agents

862ec11 verified 5 months ago

7.4 kB

AGENT_GUIDE.md — FrankenMoE Reproducible Pipeline

Purpose: Step-by-step guide for building a Mixture-of-Experts (MoE) from LoRA fine-tuned dense models. Target reader: AI agents, MLOps engineers, future you. Last verified: May 2026, mergekit 0.1.4, Qwen2.5-1.5B-Instruct

Quick Reference Card

ARCH: QwenMoE (mergekit) → GGUF (llama.cpp)
BASE: Qwen2.5-1.5B-Instruct
EXPERTS: 2 (coding, math) + 1 shared
SIZE: 3.86B params, 8.2 GB safetensors, 8.2 GB GGUF F16
INFRA: RTX 4060 Ti 16GB (train) + RTX 8000 48GB (merge, convert)

Phase 1: LoRA Fine-Tuning

# Install
pip install torch transformers peft datasets accelerate

# Train (example — use your own training script)
python train_lora.py \
  --base_model unsloth/Qwen2.5-1.5B-Instruct \
  --domain coding \
  --lora_r 16 --lora_alpha 32 \
  --target_modules q_proj,k_proj,v_proj,o_proj \
  --learning_rate 2e-5 --batch_size 4 \
  --epochs 3 --precision bf16

Key config

Param	Value	Why
lora_r	16	Good balance size/quality
lora_alpha	32	Standard alpha=2*r
target_modules	q,k,v,o_proj	Attention only (FFN stays frozen)
precision	bf16	Required for RTX 4060 Ti

Output

outputs/coding_lora/
├── adapter_model.safetensors   (~71 MB)
├── adapter_config.json
└── tokenizer files...

Phase 2: LoRA → Dense Merge

from peft import PeftModel
from transformers import AutoModelForCausalLM

base = AutoModelForCausalLM.from_pretrained(
    "unsloth/Qwen2.5-1.5B-Instruct",
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
model = PeftModel.from_pretrained(base, "outputs/coding_lora")
model = model.merge_and_unload()
model.save_pretrained("outputs/dense_coding")

⚠️ CRITICAL: merge_and_unload() is REQUIRED. mergekit cannot use LoRA adapters directly.

Phase 3: MoE Assembly (mergekit)

3a. Install & Patch mergekit

pip install mergekit==0.1.4

3b. PATCH REQUIRED — `router.py` line 122

mergekit 0.1.4 passes load_in_4bit / load_in_8bit directly to from_pretrained(). This is BROKEN in transformers >= 4.40.

# Find the file
ROUTER=$(python3 -c "import mergekit.moe; print(mergekit.moe.router.__file__)")

# Patch: remove load_in_4bit and load_in_8bit params
sed -i 's/load_in_4bit=load_in_4bit,//' $ROUTER
sed -i 's/load_in_8bit=load_in_8bit,//' $ROUTER

3c. MoE Config

# moe_config.yaml
base_model: unsloth/Qwen2.5-1.5B-Instruct
gate_mode: random          # "hidden" also works but more complex
dtype: bfloat16
experts_per_token: 1       # top-1 routing

experts:
  - source_model: outputs/dense_coding
    positive_prompts:
      - "Write a Python function to sort a list"
      - "Debug this code snippet"
  - source_model: outputs/dense_math
    positive_prompts:
      - "Solve x^2 + 5x + 6 = 0"
      - "Find the derivative of f(x) = x^3"

shared_experts:
  - source_model: unsloth/Qwen2.5-1.5B-Instruct
    positive_prompts:
      - "Hello, how are you?"
      - "What is the capital of France?"

3d. Run mergekit-moe

mergekit-moe moe_config.yaml moe_output/ --trust-remote-code

3e. CRITICAL RULES

Rule	Wrong ❌	Right ✅
Shared experts	0	Exactly 1
Routed experts count	3	2, 4, or 8 (power of 2)
Expert source	LoRA adapter	Merged dense model
llama.cpp compatibility	Non-power-of-2	2^n only

Phase 4: GGUF Conversion (THE HARD PART)

Problem: Qwen2.5 uses Tied Embeddings

Qwen2.5 has tie_word_embeddings: true → no separate lm_head.weight tensor. llama.cpp requires explicit output.weight in the GGUF.

Fix: Clone embed_tokens → lm_head

import torch, os, shutil, json
from safetensors.torch import save_file, load_file

SRC = "moe_output"
DST = "moe_output_fixed"
os.makedirs(DST, exist_ok=True)

# Copy config files
for f in ["config.json", "tokenizer.json", "tokenizer_config.json"]:
    src_f = os.path.join(SRC, f)
    if os.path.exists(src_f):
        shutil.copy2(src_f, os.path.join(DST, f))

# Fix each shard
for sf in sorted(f for f in os.listdir(SRC) if f.endswith(".safetensors")):
    tensors = load_file(os.path.join(SRC, sf))
    if "model.embed_tokens.weight" in tensors:
        tensors["lm_head.weight"] = tensors["model.embed_tokens.weight"].clone()
    save_file(tensors, os.path.join(DST, sf))

# Update config
with open(os.path.join(DST, "config.json")) as f:
    config = json.load(f)
config["tie_word_embeddings"] = False
with open(os.path.join(DST, "config.json"), "w") as f:
    json.dump(config, f)

Convert to GGUF

cd llama.cpp
python3 convert_hf_to_gguf.py moe_output_fixed --outtype f16 --outfile model-F16.gguf

Verify output.weight exists

# GGUF binary check: output.weight should appear multiple times
strings model-F16.gguf | grep -c "output.weight"
# Expected: > 0 (found 29 in our build)

Phase 5: Test

transformers

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "moe_output_fixed", trust_remote_code=True,
    torch_dtype=torch.bfloat16, device_map="auto"
)

llama.cpp

llama-cli -m model-F16.gguf -p "Write a Python function to sort a list"

⚠️ Expected: With gate_mode: random, output is coherent but not domain-optimal. Router training is needed for production quality.

Phase 6: Quantize (Optional)

# F16 → Q4_K_M (~4x smaller)
llama-quantize model-F16.gguf model-Q4_K_M.gguf Q4_K_M

Complete File Checklist

✅ moe_output_fixed/
   ├── config.json              (tie_word_embeddings: false)
   ├── tokenizer.json
   ├── model-00001-of-00002.safetensors   (has lm_head.weight)
   ├── model-00002-of-00002.safetensors
   └── model.safetensors.index.json
✅ model-F16.gguf               (has output.weight)
✅ model-Q4_K_M.gguf            (optional, 4x smaller)

Common Pitfalls

Error	Cause	Fix
`missing tensor 'output.weight'`	tied embeddings, no lm_head	Clone embed→lm_head, tie=False
`MistralForCausalLM got load_in_4bit`	mergekit 0.1.4 bug	Patch router.py line 122
`3 experts not power of two`	Wrong expert count	Use 2, 4, or 8 experts
`QwenMoE requires 1 shared expert`	No shared_expert in config	Add shared_experts section
Garbage output	LoRA not merged, or 3 experts	merge_and_unload(), use 2^n
OOM during GGUF convert	16GB GPU not enough	Use CPU: `--outtype f16` (CPU-only)

Environment Used

OS: Ubuntu 22.04
GPU: Quadro RTX 8000 48GB
Python: 3.12
mergekit: 0.1.4
transformers: 4.49+
torch: 2.5+
peft: latest
safetensors: latest
llama.cpp: latest (git clone)

Scaling Up

For larger base models (Qwen2.5-7B, 14B, 32B):

Same pipeline works — just more VRAM needed
Training: multi-GPU or cloud GPU with >24GB
MoE assembly: CPU-only works (no GPU needed for merge)
GGUF F16 size ≈ params × 2 bytes (3.86B → 7.7 GB, 7B → 14 GB)
Router training: use cloud GPU (48GB+) with classification loss
4 experts ideal: coding, math, chat, medical — fills all 2^n slots

Built by UKA — May 2026 — Bangkok, Thailand 🇹🇭