mimo_v2_flash¶
MiMo-V2-Flash (MiMoV2FlashForCausalLM).
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. Queries and keys
are head_dim wide and values v_head_dim wide, so the family defines
qk_head_dim as the config's head_dim and head_dim as v_head_dim. The
values are scaled by attention_value_scale before the interface.
layer_types alternates full and sliding-window attention. A sliding block has
twice num_key_value_heads key/value heads and an attention sink per head: a
learned logit that joins the softmax as one extra key column and is dropped
afterwards, so that block's pattern rows sum to less than one. num_kv_heads is
the config's, the full blocks' count. The softmax input is shifted by its row max
(and on a sliding block is one key wider), so attention_scores is read one step
earlier, at the masked scores, on both block types. mlp_layer_types says which
blocks are dense and which a mixture of experts returning the routed hidden states.
module-attribute
¶
RENAME = {'model.embed_tokens': 'embed_tokens', 'model.layers': 'layers', 'model.norm': 'norm', 'gate': 'router'}
module-attribute
¶
ENVOYS = {MiMoV2FlashDecoderLayer: Layer, MiMoV2FlashAttention: Attention, MiMoV2FlashMLP: Mlp, MiMoV2FlashMoE: Moe}
head_dim(model: StandardizedTransformer) -> int
Width of one head's values and outputs: v_head_dim.
qk_head_dim(model: StandardizedTransformer) -> int
Width of one head's queries and keys: the config's head_dim.