moe-lab

Mixture-of-experts architecture lab

Implement model.py and run the supplied anonymous checkpoint using PyTorch, without loading a pretrained-model implementation from another library. You may use safetensors and tokenizers for file loading and tokenization.

Required model

The checkpoint describes a decoder-only sparse mixture-of-experts Transformer. It accepts token IDs, returns next-token logits, and supports autoregressive generation. Read exact values from model.json; this artifact uses:

flowchart TD
    A[Token IDs<br/>batch × sequence] --> B[Token table<br/>65,536 × 1,024]
    B --> C[Multiply by embedding scale]
    C --> D[25 causal MoE decoder units]
    D --> E[Final RMSNorm]
    E --> F[Tied token-table projection]
    F --> G[Divide by logits scale]
    G --> H[Keep first 49,152 logits]
    H --> I[Select or sample next token]
    I -. append and repeat .-> A

Every unit contains two pre-normalized residual branches:

flowchart LR
    X[Input x] --> N1[RMSNorm]
    X --> R1((+))
    N1 --> Q[Q: 16 heads<br/>K,V: 8 heads]
    Q --> A[RoPE + causal GQA]
    A --> O[Output projection × residual scale]
    O --> R1
    R1 --> Y[Intermediate y]
    Y --> N2[RMSNorm]
    Y --> R2((+))
    N2 --> ROUTE[Router 1,024 → 32]
    ROUTE --> TOP[Top 8 + softmax]
    TOP --> EXP[Selected SwiGLU experts]
    EXP --> SUM[Weighted sum × residual scale]
    SUM --> R2
    R2 --> Z[Unit output]

Grouped-query causal attention

Queries have 16 heads, but keys and values have only 8. Repeat each key/value head twice so it serves one group of query heads. Apply rotary position information to queries and keys before repeating them. Use the explicit attention scale from model.json, not PyTorch's default scale. Tokens may attend only to valid keys at their own or earlier positions.

Sparse expert routing

For every token independently:

1. Project the normalized token to 32 router logits in float32.
2. Select the eight largest logits.
3. Apply softmax only across those eight values.
4. Send the token through each selected expert.
5. Multiply each expert output by its routing probability and sum the results.

Each expert owns two bias-free matrices. The first projects 1,024 values to 1,024 values, which split into two 512-wide halves. Apply SiLU to the first half and multiply it elementwise by the second half. The second expert matrix projects the result back from 512 to 1,024.

The expert weights are stored together:

• expand.weight: [32, 1,024, 1,024];
• contract.weight: [32, 1,024, 512];
• route.weight: [32, 1,024].

Scaling and normalization

RMSNorm computes its variance in float32, then restores the input dtype before multiplying by its learned weight. Multiply token embeddings by embedding_scale. Multiply both attention and expert outputs by residual_scale before their residual additions. Divide tied-projection logits by logits_scale.

Checkpoint contract

Load safetensor_moe.safetensors. Module names must match the neutral checkpoint schema: tokens, units, pre_mix, mix, pre_experts, experts, route, expand, contract, and output_norm. Instantiate and execute every unit declared in model.json.

The stored token table has more rows than the tokenizer vocabulary. It is also the output projection weight. Return only the first vocabulary_size logits.

Input and output examples

forward()

Tokenizing Once upon a time produces:

input_ids = torch.tensor([[
    9297, 12683, 312, 1133
]])
# tokens: Once Ġupon Ġa Ġtime

Calling model(input_ids) returns:

ModelOutput(
    logits=...,         # shape [1, 4, 49_152]
    hidden_states=...,  # shape [1, 4, 1_024]
    router_logits=(...),
)

router_logits is a tuple with one tensor per executed unit:

len(output.router_logits) == 25
output.router_logits[0].shape == (1, 4, 32)

Each [1, 4, 32] tensor contains one score for every expert at every token. The eight largest scores select the active experts. Router logits are useful for inspection; generation uses the vocabulary logits.

Although the tied token table has 65,536 rows, the returned vocabulary axis is 49,152 because storage-only rows are sliced away.

generate()

With greedy decoding and three new tokens:

generated = model.generate(input_ids, max_new_tokens=3, temperature=0)

generated = torch.tensor([[
    9297, 12683, 312, 1133, 30, 2017, 1597
]])

Input:  Once upon a time
Output: Once upon a time, there was

The output contains the four-token input prefix followed by three generated tokens, giving shape [1, 7]. generate() returns token IDs directly, whereas forward() returns ModelOutput.

Completion criteria

• All NotImplementedError sites in model.py are implemented.
• Causal and padding masks are correct.
• Grouped-query attention uses the configured scale and rotary base.
• Routing uses top-8 softmax probabilities and combines all selected experts.
• Earlier logits do not change when later tokens are appended.
• Every checkpoint tensor loads with strict=True.
• python inference.py . "Once upon a time" --max-new-tokens 20 runs.
• The implementation uses the supplied weights rather than Transformers.