bamba

bamba

Bamba (BambaForCausalLM).

Llama's containers over one block class that holds either a Mamba-2 mixer or attention, by config.layers_block_type (attn_layer_indices name the attention blocks)::

model.layers[i]          BambaDecoderLayer
    .input_layernorm
    .mamba               BambaMixer, an SSD mixer   (or .self_attn, BambaAttention)
    .pre_ff_layernorm
    .feed_forward        BambaMLP

A pre-norm block, as Llama: the mixer's output and then the MLP's are added to the stream as the modules return them. The names are Llama's but for three: mamba is linear_attn (nnterp.StateSpace), feed_forward is mlp and pre_ff_layernorm is post_attention_layernorm. Each block has self_attn or linear_attn, never both. The block returns (hidden_states, attention_weights).

RENAME module-attribute

RENAME = {'model.embed_tokens': 'embed_tokens', 'model.layers': 'layers', 'model.final_layernorm': 'norm', 'mamba': 'linear_attn', 'feed_forward': 'mlp', 'pre_ff_layernorm': 'post_attention_layernorm'}

ENVOYS module-attribute

ENVOYS = {BambaDecoderLayer: Layer, BambaAttention: Attention, BambaMixer: StateSpace, BambaMLP: Mlp}