# 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

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

```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):
1. **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).
2. **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`:

```json
{"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.Parameter`s, not Linear
modules); the router and dense weights via `lora_target_modules`:

```yaml
# ── 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:

1. **`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.
2. **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

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

```bash
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):
1. no "old-style model class" warning; weights load (786 tensors);
2. `[solon_moe] SOLON_LAMBDA=0.1 applied to 15 layer(s)` — inference λ active;
3. 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.