granitemoehybrid¶

granitemoehybrid ¶

GraniteMoE-Hybrid (GraniteMoeHybridForCausalLM, IBM Granite 4.0 and 4.0-H).

Llama's containers over one block class that holds either a Mamba-2 mixer or attention, by config.layers_block_type, and then a feed-forward::

model.layers[i]          GraniteMoeHybridDecoderLayer
    .input_layernorm
    .mamba               GraniteMoeHybridMambaLayer, an SSD mixer   (or .self_attn)
    .post_attention_layernorm
    .block_sparse_moe    GraniteMoeHybridMoE (only when num_local_experts > 0)
    .shared_mlp          GraniteMoeHybridMLP (on every block)

h = x + mixer(input_layernorm(x)) * residual_multiplier
out = h + (block_sparse_moe(n) + shared_mlp(n)) * residual_multiplier     with n = post_attention_layernorm(h)

(shared_mlp(n) alone on a dense checkpoint, Granite-4.0-H-Micro's.) mamba is linear_attn (nnterp.StateSpace) and shared_mlp, the one feed-forward every block has, is mlp; each block has self_attn or linear_attn, never both. As on Granite (granite.py) the block adds each sublayer's output times residual_multiplier: both mixers' attention_output are the module's output times it, computed copies divided on assignment and carried back by a transform. No module returns the feed-forward's sum, so mlp_output is the block's binding of it (hidden_states_4 with experts, hidden_states_5 without) times the multiplier, and the family's Layer sets sourced for it; the mixer and the shared expert, whose modules keep no config, read the multiplier (and the Mlp whether routed experts run) off their parent block. A write lands on the sum, replacing both experts' contribution. mlp.output is the shared expert's output alone and layers[i].block_sparse_moe.output the routed experts'. The mixture's values live on mlp (a Moe), which reads block_sparse_moe's router and experts through its parent: shared_expert_output is mlp.output and routed_output + shared_expert_output == mlp_output / residual_multiplier. mlp.intermediate_size and the root's intermediate_size are the shared expert's width, shared_intermediate_size; the routed experts are intermediate_size wide. The attention runs the shared eager forward with attention_multiplier as its scale and no rotary under position_embedding_type nope (Granite-4.0-H); embedding_multiplier and logits_scaling are as on Granite.

RENAME `module-attribute` ¶

RENAME = {'model.embed_tokens': 'embed_tokens', 'model.layers': 'layers', 'model.norm': 'norm', 'mamba': 'linear_attn', 'shared_mlp': 'mlp', 'activation': 'act_fn', 'output_linear': 'down_proj'}

EXPERTS_SUM `module-attribute` ¶

EXPERTS_SUM = 'hidden_states_4'

SHARED_ONLY `module-attribute` ¶

SHARED_ONLY = 'hidden_states_5'

ENVOYS `module-attribute` ¶

ENVOYS = {GraniteMoeHybridDecoderLayer: Layer, GraniteMoeHybridAttention: Attention, GraniteMoeHybridMambaLayer: StateSpace, GraniteMoeHybridMLP: Mlp}

Attention ¶

Bases: Attention

GraniteMoE-Hybrid's attention: Granite's, the block adds its output times residual_multiplier.

Mlp ¶

Mlp(*args: Any, **kwargs: Any)

Bases: Moe

GraniteMoE-Hybrid's shared expert, standing for the block's feed-forward: the block adds it plus the routed experts, times residual_multiplier.

It hosts the block's mixture of experts (block_sparse_moe, whose router and experts it reads off its parent block), and is itself the shared expert.

SCORING `class-attribute` `instance-attribute` ¶

SCORING = 'topk_softmax'

residual_multiplier `class-attribute` `instance-attribute` ¶

residual_multiplier = property(residual_multiplier)

has_experts `property` ¶

has_experts: bool

Whether routed experts run beside this shared expert: the block's has_experts.

router `property` ¶

router: Envoy

The block's mixture's router.

experts `property` ¶

experts: Envoy

The block's mixture's experts.

intermediate_size `property` ¶

intermediate_size: int

Width of the shared expert's hidden layer: the module's hidden_size (shared_intermediate_size).

no_mixture ¶

no_mixture() -> str | None

shared_expert_output ¶

shared_expert_output(value) -> Residual

This MLP's own output: the shared expert beside the routed ones, before the block's multiplier.

mlp_output ¶

mlp_output(value, raw)

intermediate_size ¶

intermediate_size(model: StandardizedTransformer) -> int

The MLP width is the shared expert's, shared_intermediate_size (mlp is the shared expert); intermediate_size is the routed experts'.

granitemoehybrid¶