AAIE MoE Pretrain (Mixtral-compatible export)
A MixtralForCausalLM-format export of AAIE MoE Pretrain — same weights, same architecture, just registered under transformers'/vLLM's native Mixtral implementation instead of custom trust_remote_code.
Why this exists
Unlike the dense AAIE models (which map directly onto Llama), this MoE model has a routed-expert FFN, so it isn't Llama-compatible. It turns out to be structurally identical to Mixtral's mixture-of-experts design instead: the same routing math (softmax over router logits, top-k selection, renormalize the selected weights to sum to 1) and the same per-expert SwiGLU FFN shape as Mixtral-8x7B (also 8 experts, top-2 routing) — just with different tensor names. This export is a pure state_dict + config remap, not a retrain.
Verified byte-for-byte identical to the original: same prompt, same greedy decode, torch.equal()
on the output token ids between the original custom-class model and this Mixtral-remap returned True.
Payoff: loads natively in vLLM (no trust_remote_code, no custom backend) using vLLM's own
MixtralForCausalLM implementation, including its expert-parallel/fused-MoE kernels.
Usage
from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("namquangstudy/aaie-moe-pretrain-mixtral")
model = AutoModelForCausalLM.from_pretrained("namquangstudy/aaie-moe-pretrain-mixtral")
# no trust_remote_code needed - this is a plain MixtralForCausalLM
from vllm import LLM, SamplingParams
llm = LLM(model="namquangstudy/aaie-moe-pretrain-mixtral")
out = llm.generate(["The purpose of a database index is"], SamplingParams(max_tokens=100))
print(out[0].outputs[0].text)
Important caveat: this checkpoint is not converged
Training was stopped at 41.9% of its planned budget (step 41,500 / 100,000) — see the original
model's README and the project's TECHNICAL_REPORT_MOE.md for why and what's still open. Expect
repetitive, undertrained-sounding output; this is a faithful export of an in-progress checkpoint,
not a finished model.
Full details
8-expert, top-2-routed MoE variant (617.91M total / 222.73M active parameters per token),
distilled from Qwen/Qwen2.5-3B on FineWeb-Edu. No instruction-tuning. Full architecture, training
data, and hyperparameters are documented on the original
AAIE MoE Pretrain model card — this
export carries the exact same weights and behavior, just packaged for vLLM/ecosystem compatibility.