phimoe

phimoe

Phi-3.5-MoE (PhimoeForCausalLM).

Llama's tree and Llama's block: model.{embed_tokens, layers[i].{input_layernorm, self_attn, post_attention_layernorm, mlp}, norm} and lm_head, the residual added in the block, attention through the shared eager forward. The norms are nn.LayerNorm. Every block's MLP is a sparse mixture of experts that returns the routed hidden states as a bare tensor; the experts are intermediate_size wide.

RENAME module-attribute

RENAME = {'model.embed_tokens': 'embed_tokens', 'model.layers': 'layers', 'model.norm': 'norm'}

ENVOYS module-attribute

ENVOYS = {PhimoeDecoderLayer: Layer, PhimoeAttention: Attention, PhimoeSparseMoeBlock: Mlp}

Mlp

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

Bases: Moe

A mixture of experts: the module returns the routed hidden states as a bare tensor, so the base holds.

The router is an nn.Linear subclass whose forward calls its parent's (super().forward), so the logits are that call's output; sparsemixer turns them into the weights.

SCORING class-attribute instance-attribute

SCORING = 'sparsemixer'

router_logits

router_logits(value) -> RouterLogits