jamba¶
Jamba (JambaForCausalLM): Mamba-1 blocks, attention blocks and a mixture of experts.
Its tree is model.{embed_tokens, layers[i], final_layernorm} plus
lm_head, with two block classes (config.layers_block_type): most
blocks are JambaMambaDecoderLayer, {input_layernorm, mamba,
pre_ff_layernorm, feed_forward}, and one in attn_layer_period is
JambaAttentionDecoderLayer, the same with self_attn in place of
mamba. Both are Llama's pre-norm sequential block, so the identity is
Llama's. Each block has one mixer, never both: on a Mamba block every
self_attn value is reported missing and the selective-scan values live
at layers[i].linear_attn (see nnterp.SelectiveScan); on an attention
block the reverse. feed_forward is a dense MLP or, one block in
expert_layer_period, a sparse mixture of experts; both are mlp, and
pre_ff_layernorm is post_attention_layernorm.
The attention has no positional encoding (the Mamba blocks carry position)
and runs the shared eager interface, so the base holds. The Mamba mixer puts
RMS norms on the step size, B and C before the scan (with weights,
where Falcon-Mamba's are weightless); the values are read at the kernel
call, after them.
With mamba_ssm installed, call nnterp.route_kernels(model.family,
"torch") before the first trace.
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'}
module-attribute
¶
ENVOYS = {JambaAttentionDecoderLayer: Layer, JambaMambaDecoderLayer: Layer, JambaAttention: Attention, JambaMambaMixer: SelectiveScan, JambaMLP: Mlp, JambaSparseMoeBlock: Moe}