# 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 (`<|channel>thought`), CoT-trained |
| EOT token | `<turn|>` (id **106**) — must be set as EOS at inference |

## Quickstart (transformers)

```python
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:

```bash
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.

## Training your own

`axolotl/` contains a ready-to-adapt QLoRA config with a reasoning
(thought-channel) chat template, plus documentation for opening/closing the
router and experts independently. `scripts/` contains the Gemma4 → Solon-MoE
structure converter and the LoRA merge tool. See **USAGE.md** for the full
pipeline: convert → train → merge → serve.

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