moe-40b-a18b-lora

moe-40b-a18b-lora

LoRA adapters + trained router weights for an experimental 39.98B-total / 17.96B-active Mixture-of-Experts student, frankenstein-upcycled from Qwen/Qwen3.6-35B-A3B.

This repo contains the trained parts (adapters + router). The base student weights are not included — they are built deterministically from the public base model with the included code (see Reproduce below). A ready-to-run NF4 checkpoint is in Leonther/moe-40b-a18b-q4.

Architecture

Built by arch_builder.py: streamed split of the fused expert tensors into per-expert modules + expert growth (256 → 288) + router growth. This makes the fused base trainable (bitsandbytes cannot 4-bit fused 3D expert params).

Training

QLoRA: NF4 base, LoRA rank 2 / alpha 4 on expert + attention projections, plus the router weights (raw fp32, never quantized). SFT data: GLM-5.3-Flash gold + qwen38 synthetic imitation (see Leonther/moe-40b-a18b-dataset).

Training beyond top-16 OOMs on a 32GB card at this model size. Future work: slim rebuild (256 experts, no growth) to unlock top-32; substantially more SFT data to avoid overfitting (the current set is only a few hundred examples).

Files

iter_1/  adapter_model.safetensors, adapter_config.json, router.pt   (top-8)
iter_2/  ...                                                          (top-16, best)
iter_5/  ...                                                          (top-16, overfit)
student_config.json      # built student architecture config (288 experts)
build_report.json        # param accounting of the build
quantize_report.json     # NF4 pre-quantization stats
code/                    # arch_builder.py, quantize_student.py, student_moe.py, trainer.py, ...

Reproduce

# 1. get the public base
hf download Qwen/Qwen3.6-35B-A3B --local-dir base

# 2. build the 40B-A18B student (288 experts, per-expert layout)
python code/arch_builder.py --base base --out student

# 3. pre-quantize to NF4 (bnb-native layout, fast loads)
python code/quantize_student.py --src student --out student_q4

# 4. load base + adapter (custom loader patches the fused-expert block)
python - <<'PY'
import student_moe as sm
model = sm.load_student_model("student_q4", top_k=16, for_training=False)
model = sm.attach_lora(model, {"lora_rank": 2, "lora_alpha": 4})
sm.load_adapter_bundle(model, "iter_2")
PY

Status

One-day experimental run, shared for reference. Expect rough edges; the training loop and builder are in code/.