laguna

laguna

Laguna (LagunaForCausalLM, poolside's Laguna).

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, which returns a bare tensor. What differs is inside the sublayers:

• Per-block head counts. num_attention_heads_per_layer gives each block its own number of query heads (48 on the full-attention blocks and 64 or 72 on the sliding-window ones in the released checkpoints); the key/value heads and head_dim are the same everywhere. layers[i].self_attn.num_heads is that block's count, read off the module; the root's num_heads is the config's top-level num_attention_heads.
• Gated attention. q/k norms before the rotary (partial on the full blocks), the shared eager forward, then a softplus gate (g_proj, per head or per channel) on the interface's output before o_proj: attention_head_outputs are ungated and attention_output is the gated projection the block adds.
• Mixture of experts. mlp_layer_types names the dense blocks (LagunaMLP, intermediate_size wide) and the sparse ones (LagunaSparseMoeBlock: routed experts moe_intermediate_size wide, times moe_routed_scaling_factor, plus a shared expert, as one tensor). The shared expert (mlp.shared_experts) is a LagunaMLP too, so it is an Mlp; its mlp_output is unavailable, since the block adds the mixture's sum.

RENAME module-attribute

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

ENVOYS module-attribute

ENVOYS = {LagunaDecoderLayer: Layer, LagunaAttention: Attention, LagunaMLP: Mlp, LagunaSparseMoeBlock: Moe}