Mixture of Wide Experts for OpenLanguageModel (OLM)

This repository is the minimal public implementation of Mixture of Wide Experts (MoWE) as an extension to OpenLanguageModel (source, paper). It is not a standalone transformer stack.

MoWE follows one architectural rule: read all, write disjoint. A router selects a small set of experts for each token. Every selected expert reads the complete residual representation and writes only to its fixed output slice.

src/mowe/model.py       OLM-native language-model assembly
src/mowe/ffn.py         Dense, standard-MoE, and MoWE FFNs
src/mowe/routing.py     Top-k softmax routing and auxiliary losses
src/mowe/kernels/       Optional forward-only CUDA and Triton kernels
examples/train.py       Generic training through olm.train.Trainer
configs/                Released architectures and a small training example
tests/                  CPU architecture and optional CUDA correctness tests

Pretrained checkpoints are released on Hugging Face: MoWE, dense, and standard MoE. This repository contains no datasets, corpus loaders, private paths, experiment logs, optimizer states, evaluation pipelines, or paper assets.

python -m venv .venv
source .venv/bin/activate
pip install -e .

For the optional optimized MoWE forward:

pip install -e '.[kernels]'

The kernel path requires CUDA, Triton, a local CUDA toolkit, Ninja, and a C++ compiler. The CUDA/CUB router extension compiles on first use.

from mowe import ModelConfig, RoutingLanguageModel

config = ModelConfig.from_json("configs/mowe-768-19l.json")
model = RoutingLanguageModel(config)

RoutingLanguageModel subclasses OLM's Block. It uses OLM token embeddings, layer normalization, RoPE attention, output head, and training interface. The routed FFNs are the research contribution supplied by this package.

For a small end-to-end example, place ordinary UTF-8 .txt files in a directory, then run:

python examples/train.py \
  --config configs/mowe-tiny.json \
  --data-dir /path/to/text \
  --output runs/mowe-tiny

Training uses OLM's Trainer, RoutingCrossEntropyLoss, and the readable, differentiable selected-expert reference path. The public release does not include custom backward kernels. The 768-by-19 configs describe the released checkpoints and require hardware and batch sizes appropriate for their scale.

import torch
from mowe import load_checkpoint

model = load_checkpoint(
    "/path/to/downloaded/checkpoint",
    device="cuda",
    dtype=torch.bfloat16,
).eval()

For the optimized inference path, pass inference_backend="optimized". This backend specializes to the released MoWE shape: d_model = G = 768, H = 8, and k = 38.

The optional forward backend contains the exact optimization family used for the paper's 97.7% throughput result:

• exact BF16 radix top-38 selection over 768 experts;
• selected-softmax, dispatch, and grouped-schedule construction in one CUDA router pass;
• a grouped Triton tensor-core kernel that evaluates selected experts only;
• direct writes into each expert's residual coordinate, with optional fused residual addition.

The reported measurement used BF16, batch 16, sequence length 1,024, 19 layers, CUDA Graph replay, PyTorch 2.4.1 with CUDA 12.4, Triton 3.0.0, and an NVIDIA RTX A6000. The transformer backbone excluding the vocabulary head took 73.269 ms for dense and 74.989 ms for MoWE, so MoWE retained 97.7% of dense forward throughput. This is a fixed-shape inference benchmark, not a training-throughput claim.

pytest -q

The CUDA correctness test is skipped when CUDA is unavailable. Run it explicitly on a configured CUDA machine with pytest -q -m cuda.

The code is provided under the Apache License 2.0. Citation metadata for MoWE and its OpenLanguageModel foundation is in CITATION.cff.