TOTORONG/Solon_MOE_22B_A17B

Solon-MoE-22B-A17B

A sparse Mixture-of-Experts upcycled from Gemma4-12B-IT — 22B total / ~17B active parameters, with the original dense model fully preserved as a shared path.

Solon-MoE converts 15 of Gemma4's 48 decoder layers (L18–L34, excluding 23 & 29) into MoE layers using the native transformers Gemma4 MoE block, applied selectively per layer. Each MoE layer adds 4 experts (initialized from the layer's dense FFN) with top-2 routing, running in parallel with the preserved dense FFN. A learnable per-layer scale λ (post_feedforward_layernorm_2) gates the routed contribution — setting λ=0 exactly recovers the original dense model, which makes λ both a safety knob and a diagnostic tool.

Base model	google/gemma4-12b-it
Architecture	SolonMoEForCausalLM (57-line shim over native Gemma4)
MoE layers	15 of 48 (L18–22, 24–28, 30–34)
Experts / top-k	4 / 2 (softmax → top-k → renormalize → per-expert scale)
Expert weights	packed 3D: gate_up_proj [4, 30720, 3840], down_proj [4, 3840, 15360]
λ convention	trained at 0.15, served at 0.10
Reasoning	thought-channel chat template (`<
EOT token	`<turn

Quickstart (transformers)

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

m = "your-org/solon-moe-22b-a17b"
tok = AutoTokenizer.from_pretrained(m, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    m, dtype=torch.bfloat16, trust_remote_code=True, device_map="auto").eval()

text = tok.apply_chat_template(
    [{"role": "user", "content": "If a train covers 240 km, ..."}],
    add_generation_prompt=True, tokenize=False)
ids = tok(text, return_tensors="pt", add_special_tokens=False).input_ids.to(model.device)
out = model.generate(ids, max_new_tokens=2048, do_sample=False,
                     eos_token_id=tok.convert_tokens_to_ids("<turn|>"))
print(tok.decode(out[0][ids.shape[1]:], skip_special_tokens=False))

Always pass the EOT token (<turn|>, id 106) as EOS — via eos_token_id here, via generation_config.json for servers, or via stop_token_ids: [106] in vLLM requests. Without it the model's turn-end signal is ignored and generation degenerates into repetition.

Serving (vLLM)

vLLM ships a native Gemma4 implementation including the MoE block; Solon-MoE only needs a thin plugin that makes the MoE flag per-layer. See vllm/ in this repo:

pip install -e vllm/            # registers SolonMoEForCausalLM via entry point
SOLON_LAMBDA=0.10 vllm serve your-org/solon-moe-22b-a17b \
    --trust-remote-code --max-model-len 8192 \
    --max-num-seqs 32 --attention-backend TRITON_ATTN

SOLON_LAMBDA=0.10 applies the inference-time λ at load (no checkpoint copy needed). On Blackwell (sm120) use --attention-backend TRITON_ATTN; the FlashInfer decode kernel does not support this head configuration.

Results (internal evaluation, n=50 per subject)

After CoT training (260k Korean/English reasoning samples, experts + router, dense frozen), Solon-MoE exceeds the original dense model on three Korean legal-reasoning benchmarks while remaining statistically indistinguishable on GPQA-diamond (paired McNemar):

Benchmark (KMMLU-Hard, CoT)	Solon-MoE	Gemma4-12B-IT
law	0.38	0.32
patent	0.36	0.32
criminal law	0.32	0.30
GPQA-diamond (English)	0.56	0.64 (p=0.45, n.s.)

Notable training-dynamics findings: adaptation concentrates in deep MoE layers (norm, effective rank, and orthogonality all increase with depth); per-expert learned updates are near-orthogonal (mean pairwise cos ≈ 0.01–0.09), evidence of parameter-space expert differentiation; router rewiring localizes to deep layers while mid-band routing is preserved.

Limitations

• Korean taxation prompts can occasionally produce runaway generations (~12% at n=50); a stability-replay fine-tune is the known mitigation.
• λ > 0.10 at inference degrades quality; keep the 0.10 convention.
• Requires trust_remote_code=True (57-line architecture shim).

License

Inherits the Gemma license and usage terms from the base model. The conversion and training code in this repository is released under Apache-2.0.