core/nn/moe_transformer library
Transformer stack with the dense FFN sublayer replaced by a Mixture-of-Experts block.
Same pre-LN skeleton as AFTBlock / the vanilla TransformerBlock
— the only structural difference is ffn -> MoEFeedForward. The
self-attention path is standard multi-head causal attention so this
module drops in wherever TransformerLM / GPT would.
MoELanguageModel mirrors TransformerLM's public surface
(call(tokens) -> logits) so training loops built for one can
train the other unchanged. Call
model.updateRoutingBias() once per epoch (or similar cadence) to
let the aux-loss-free load balancer nudge routing.
Classes
- MoEBlock
- MoELanguageModel
- Decoder-only language model with MoE FFNs.