Skip to content

falcon_h1

falcon_h1

Falcon-H1 (FalconH1ForCausalLM).

Llama's containers over a parallel hybrid block: every block runs a Mamba-2 mixer and attention side by side on the same normed input, then an MLP::

model.layers[i]          FalconH1DecoderLayer
    .input_layernorm
    .mamba               FalconH1Mixer, an SSD mixer
    .self_attn           FalconH1Attention
    .pre_ff_layernorm
    .feed_forward        FalconH1MLP

h = x + mamba(n) * ssm_out_multiplier + self_attn(n * attention_in_multiplier) * attention_out_multiplier
out = h + feed_forward(pre_ff_layernorm(h))          with n = input_layernorm(x)

mamba is linear_attn (nnterp.StateSpace), feed_forward is mlp and pre_ff_layernorm is post_attention_layernorm. Every block has both mixers, so the contribution identity has four terms, input + self_attn.attention_output + linear_attn.attention_output + mlp_output == layer_output.

The config's µP multipliers, and where each is handled:

  • ssm_out_multiplier / attention_out_multiplier: the block scales each mixer's output before adding it, so attention_output on each is that product, the block's own binding of it (mamba_hidden_states_1, attention_hidden_states_0): the tensor the block adds, so reads, writes and in-place edits need no arithmetic.
  • lm_head_multiplier: the model multiplies the head's output by it, so the family defines project_on_vocab with that step.
  • embedding_multiplier: the model multiplies the embedding module's output by it before the first block, so token_embeddings (the module's output) times the multiplier is layers[0].input.
  • attention_in_multiplier: the attention's input is the normed stream times it (1.0 on the released checkpoints), so self_attn.input is the norm's output only then.
  • mlp_multipliers, key_multiplier, ssm_in_multiplier, ssm_multipliers: applied inside the MLP, the attention and the mixer; the values read where the forward already applied them.

With mamba_rms_norm false the decode step passes the gate into the update kernel, so a decode step's attention_head_outputs is gated by silu(z) where a prompt's is not; with it true (Falcon-H1 1.5B and up, H1R-7B) both are ungated. Falcon-H1-0.5B and the Falcon-H1-Tiny checkpoints set it false, and their linear_attn.norm is an identity. The block returns a one-element tuple.

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'}

MAMBA_SCALED module-attribute 

MAMBA_SCALED = 'mamba_hidden_states_1'

ATTENTION_SCALED module-attribute 

ATTENTION_SCALED = 'attention_hidden_states_0'

ENVOYS module-attribute 

ENVOYS = {FalconH1DecoderLayer: Layer, FalconH1Attention: Attention, FalconH1Mixer: StateSpace, FalconH1MLP: Mlp}

project_on_vocab 

project_on_vocab(model: StandardizedTransformer, hidden: Tensor) -> Tensor

The logit lens as the model makes its logits: the final norm, lm_head, then times lm_head_multiplier.