USAGE — Solon-MoE end-to-end pipeline

Four stages: convert → train → merge → serve. Each stage is a single command plus a couple of verifications. All tools live in this repository.

scripts/convert_gemma4_to_solon_moe.py   Gemma4-12B-IT → Solon-MoE skeleton
scripts/modeling_solon_moe.py            architecture shim (ships inside the model dir)
axolotl/solon_moe_plugin.py              axolotl plugin (REQUIRED for training)
axolotl/train_solon_moe.yaml             training config example (reasoning template)
scripts/merge_lora.py                    merge + verification + EOS patch
vllm/                                    vLLM serving plugin (pip installable)

1. Convert: build the MoE skeleton from Gemma4-12B-IT

python scripts/convert_gemma4_to_solon_moe.py \
    --source /path/to/gemma4-12b-it \
    --out    /path/to/solon-moe-p0 \
    --lambda-init 0.15

What it does (deterministic):

For each MoE layer (L18–22, 24–28, 30–34): replicates the layer's dense FFN into 4 packed experts (gate_up_proj [4, 2·ffn, hidden] with gate first, down_proj [4, hidden, ffn]), keeping the dense FFN as the shared path.
Inserts post_feedforward_layernorm_2.weight (the λ scale) at --lambda-init.
Inserts router.proj [4, hidden] (zeros by default → uniform routing), router.per_expert_scale (ones), router.scale (one).
Patches config.json: model_type: solon_moe, moe_layers, expert counts, and auto_map pointing at modeling_solon_moe.py (copied into the output).

Optional --expert-noise-std 0.02 adds zero-mean symmetry-breaking noise to the expert copies. Pure copies also differentiate under training (routing provides the asymmetry), so the default is 0.

Verify: load with trust_remote_code=True and generate once — output must be identical in quality to the base model (λ-scaled experts start as copies of the shared FFN, and with --router-init zeros routing is uniform).

2. Train with axolotl

Plugin (required)

Copy axolotl/solon_moe_plugin.py next to your config (or onto PYTHONPATH) and keep this in the YAML:

plugins:
  - solon_moe_plugin.SolonMoEPlugin

The plugin does two things, both at import time (so they land before model load regardless of axolotl version):

Per-layer MoE patch — gates the native global enable_moe_block flag by config.moe_layers. Without it, Gemma4 builds MoE on all 48 layers and the 33 dense layers get randomly-initialized experts (a silent disaster).
PEFT resume guard — PEFT's TP-sharding helper crashes on packed 3-D expert parameters (no .weight); the guard skips it safely in non-TP runs.

Data format (reasoning / CoT)

One JSON object per line; the assistant turn carries the chain-of-thought in reasoning_content and the final answer in content:

{"conversations": [
  {"role": "user", "content": "Question ..."},
  {"role": "assistant",
   "reasoning_content": "Let me work through this step by step ...",
   "content": "Structured final solution ... **Answer: 575**"}
]}

The tokenizer's chat template renders reasoning_content into the thought channel (<|channel>thought) followed by content, ending with <turn|>. Plain (non-reasoning) data — assistant turns with only content — works with the same template.

Verify the tokenization once before any long run (this catches template and label-masking issues that are otherwise silent). With axolotl's debug tokenization output, check four things: user turns are masked (label −100), the assistant turn including reasoning_content is unmasked, nothing is truncated, and the final <turn|> has an active label (106, 106) — not (-100, 106).

Opening and closing the trainable parts

The three parameter groups are controlled independently. Experts are addressed via lora_target_parameters (they are packed 3-D nn.Parameters, not Linear modules); the router and dense weights via lora_target_modules:

# ── Experts: the core of any Solon-MoE fine-tune. Keep these on. ──
lora_target_parameters:
  - experts.gate_up_proj
  - experts.down_proj

# ── Router: OPEN it when teaching new capabilities/domains (lets routing
#    re-organize around new specializations). CLOSE it (delete the line) for
#    stability healing or small fixes, to preserve the existing routing. ──
lora_target_modules:
  - router.proj

# ── Dense path (attention + shared FFN): keep FROZEN to guarantee the base
#    model's general ability is untouched — this is the main safety property
#    of the architecture. Open only if you accept full-model drift: ──
#  - q_proj
#  - k_proj
#  - o_proj
#  - gate_proj
#  - up_proj
#  - down_proj

Two hard rules learned the expensive way:

quantize_moe_experts: false always. Quantizing the packed 3-D expert parameters under QLoRA silently truncates their gradients — loss decreases normally while experts learn nothing.
Never use lora_modules_to_save for quantized modules (silent freeze).

Rank guidance: measured effective rank of learned expert updates is ~1–10 per projection even with large nominal ranks; lora_r: 32, lora_alpha: 64 is ample for most fine-tunes. See axolotl/train_solon_moe.yaml for the full, annotated config.

3. Merge

python scripts/merge_lora.py \
    --base    /path/to/solon-moe-p0 \
    --adapter /path/to/output_dir/checkpoint-final \
    --out     /path/to/solon-moe-merged

The script merges on CPU in bf16 and then verifies the merge in both directions: every trained tensor must differ from the base, every frozen tensor must be bit-identical, and the key sets must match. It also ensures generation_config.json contains eos_token_id: 106 (<turn|>) — required for correct stopping in every inference stack.

4. Serve with vLLM

vLLM's native Gemma4 implementation already contains the full MoE path (FusedMoE experts, router, λ norm). The plugin in vllm/ only makes the MoE flag per-layer and registers the architecture:

pip install -e vllm/

SOLON_LAMBDA=0.10 vllm serve /path/to/solon-moe-merged \
    --trust-remote-code --max-model-len 8192 \
    --max-num-seqs 32 --attention-backend TRITON_ATTN

Startup checks (all three must appear / hold):

no "old-style model class" warning; weights load (786 tensors);
[solon_moe] SOLON_LAMBDA=0.1 applied to 15 layer(s) — inference λ active;
a test request finishes with finish_reason: "stop", stop_reason: 106.

Notes:

SOLON_LAMBDA overrides the stored (training) λ at load time; the train-at-0.15 / serve-at-0.10 convention is validated — higher inference λ degrades output.
On Blackwell (sm120), --attention-backend TRITON_ATTN is required (FlashInfer's decode kernel rejects this head configuration: Unsupported max_mma_kv: 0). The env var VLLM_ATTENTION_BACKEND is removed in recent vLLM — only the CLI flag works.
Clients should still send "stop_token_ids": [106] as belt-and-braces; string-based stops are unreliable because special tokens are stripped from decoded text.