granitemoe

Model type

granitemoe

Class

GraniteMoECausalLMModel

Task

text-generation

Source

models/granite.py

Description

GraniteMoE model: MoE with Granite scaling multipliers.

Combines MoE routing (TopK-then-softmax gating, per HuggingFace GraniteMoeTopKGating) with Granite-specific scaling factors:

  • embedding_multiplier: scales embeddings after lookup

  • attention_multiplier: replaces 1/sqrt(head_dim) attention scale

  • logits_scaling: divides final logits

  • residual_multiplier: scales attention/MLP outputs before residual add

HuggingFace class: GraniteMoeForCausalLM.

Usage

mobius build --model <MODEL_ID> output_dir/

from mobius import build

model = build("<MODEL_ID>")