granitemoeshared¶
GraniteMoE with a shared expert (GraniteMoeSharedForCausalLM).
GraniteMoE's block (granitemoe.py) with a dense shared expert beside the
mixture: the block computes block_sparse_moe(h) + shared_mlp(h) from the
pre-MLP norm's output h and adds that sum times residual_multiplier. No
module returns the sum, so block_sparse_moe is mlp (it is on every block;
shared_mlp is None when shared_intermediate_size is 0) and its
mlp_output is the block's binding of the sum (hidden_states_4, or
hidden_states_3, the mixture's output alone, without a shared expert) times
the multiplier: what the block adds. As on Granite it is a computed copy, divided
on assignment and carried back by a transform after an in-place edit; the write
lands on the sum, so it replaces both experts' contribution. mlp.output is the
mixture's output alone and layers[i].shared_mlp.output the shared expert's.
The binding is in the block's forward, read after the block has started, so the
family's Layer sets sourced; the mixture keeps no config, so its Mlp reads
the multiplier and whether the block has a shared expert off its parent block
(residual_multiplier, shared). The attention, embedding_multiplier and
logits_scaling are as on Granite. The mixture is a Moe (GraniteMoE's routing);
its shared_expert_output is the block's shared_mlp's output, unscaled, so
routed_output + shared_expert_output == mlp_output / residual_multiplier.
module-attribute
¶
RENAME = {'model.embed_tokens': 'embed_tokens', 'model.layers': 'layers', 'model.norm': 'norm', 'block_sparse_moe': 'mlp'}
module-attribute
¶
ENVOYS = {GraniteMoeSharedDecoderLayer: Layer, GraniteMoeSharedAttention: Attention, GraniteMoeSharedMoE: Mlp}
Bases: Attention
GraniteMoE-Shared's attention: Granite's, the block adds its output times residual_multiplier.
Bases: Moe
GraniteMoE-Shared's mixture of experts; the block adds it plus the shared expert, times residual_multiplier.
class-attribute
instance-attribute
¶
residual_multiplier = property(residual_multiplier)
property
¶
Whether a shared expert runs beside the mixture: the block's shared_mlp.