aaie-moe-pretrain-mixtral

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.