Replace transformer blocks config with building blocks

AGENTS.md

@@ -7,7 +7,7 @@ The main steps of adding a new model are the following:
 2. Look at some existing model implementations in Bumblebee. In case of LLMs, copying an existing LLM implementation is typically a good starting point.
 
 3. Implement the model code.
-   - Whenever possible, reuse existing primitives, most notably `Layers.Transformer.blocks/2`, which is shared for most LLM implementations. Sometimes models introduce novelties to the transformer design, in which case it may be necessary to add a new option to `Layers.Transformer.blocks/2`.
+   - Use `Layers.Transformer.blocks/3` for the block stack. It handles the cache and outputs, and the model writes the block body as plain code, in the same order as the Python model. Use `self_attention/3`, `cross_attention/4`, `basic_ffn/4` and `gated_ffn/4` for standard parts. If the model has a new attention type, write it in the model with the attention steps (`project_heads/4`, `rotary_embedding/6`, `Layers.attention/8`, `output_projection/3`). Do not add model-specific options to the shared code.
    - Include relevant options from Python model configuration as Bumblebee model options (with matching defaults).
    - Make sure the `params_mapping/1` maps to correct Python layer names. You can use `Bumblebee.load_model(..., log_params_diff: true)` to get all logs related to params loading.
 

README.md

@@ -160,7 +160,7 @@ The main steps of adding a new model are the following:
 2. Look at some existing model implementations in Bumblebee. In case of LLMs, copying an existing LLM implementation is typically a good starting point.
 
 3. Implement the model code.
-   - Whenever possible, reuse existing primitives, most notably `Layers.Transformer.blocks/2`, which is shared for most LLM implementations. Sometimes models introduce novelties to the transformer design, in which case it may be necessary to add a new option to `Layers.Transformer.blocks/2`.
+   - Use `Layers.Transformer.blocks/3` for the block stack. It handles the cache and outputs, and the model writes the block body as plain code, in the same order as the Python model. Use `self_attention/3`, `cross_attention/4`, `basic_ffn/4` and `gated_ffn/4` for standard parts. If the model has a new attention type, write it in the model with the attention steps (`project_heads/4`, `rotary_embedding/6`, `Layers.attention/8`, `output_projection/3`). Do not add model-specific options to the shared code.
    - Include relevant options from Python model configuration as Bumblebee model options (with matching defaults).
    - Make sure the `params_mapping/1` maps to correct Python layer names. You can use `Bumblebee.load_model(..., log_params_diff: true)` to get all logs related to params loading.
 

lib/bumblebee/audio/whisper.ex

@@ -419,24 +419,47 @@ defmodule Bumblebee.Audio.Whisper do
     name = opts[:name]
 
     outputs =
-      Layers.Transformer.blocks(hidden_state,
-        attention_head_mask: attention_head_mask,
-        num_blocks: spec.encoder_num_blocks,
-        num_attention_heads: spec.encoder_num_attention_heads,
-        hidden_size: spec.hidden_size,
-        kernel_initializer: kernel_initializer(spec),
-        dropout_rate: spec.dropout_rate,
-        attention_dropout_rate: spec.attention_dropout_rate,
-        key_use_bias: false,
-        layer_norm: [
-          epsilon: 1.0e-5
-        ],
-        ffn: [
-          intermediate_size: spec.encoder_intermediate_size,
-          activation: spec.activation
+      Layers.Transformer.blocks(
+        hidden_state,
+        [
+          num_blocks: spec.encoder_num_blocks,
+          attention_head_mask: attention_head_mask,
+          name: join(name, "blocks")
         ],
-        block_type: :norm_first,
-        name: join(name, "blocks")
+        fn hidden_state, block ->
+          name = block.name
+
+          shortcut = hidden_state
+
+          {hidden_state, attention, _self_attention_cache} =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "self_attention_norm"))
+            |> Layers.Transformer.self_attention(
+              block,
+              [name: join(name, "self_attention")] ++
+                attention_opts(spec, spec.encoder_num_attention_heads)
+            )
+
+          hidden_state =
+            hidden_state
+            |> Axon.dropout(rate: spec.dropout_rate)
+            |> Axon.add(shortcut)
+
+          shortcut = hidden_state
+
+          hidden_state =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "output_norm"))
+            |> Layers.Transformer.basic_ffn(spec.encoder_intermediate_size, spec.hidden_size,
+              activation: spec.activation,
+              dropout_rate: spec.dropout_rate,
+              kernel_initializer: kernel_initializer(spec),
+              name: join(name, "ffn")
+            )
+            |> Axon.add(shortcut)
+
+          %{hidden_state: hidden_state, attention: attention}
+        end
       )
 
     hidden_state = Axon.layer_norm(outputs.hidden_state, name: join(name, "norm"))
@@ -461,29 +484,73 @@ defmodule Bumblebee.Audio.Whisper do
     name = opts[:name]
 
     outputs =
-      Layers.Transformer.blocks(hidden_state,
-        attention_mask: attention_mask,
-        attention_head_mask: attention_head_mask,
-        cross_hidden_state: encoder_hidden_state,
-        cross_attention_head_mask: cross_attention_head_mask,
-        cache: cache,
-        causal: true,
-        num_blocks: spec.decoder_num_blocks,
-        num_attention_heads: spec.decoder_num_attention_heads,
-        hidden_size: spec.hidden_size,
-        kernel_initializer: kernel_initializer(spec),
-        dropout_rate: spec.dropout_rate,
-        attention_dropout_rate: spec.attention_dropout_rate,
-        key_use_bias: false,
-        layer_norm: [
-          epsilon: 1.0e-5
+      Layers.Transformer.blocks(
+        hidden_state,
+        [
+          num_blocks: spec.decoder_num_blocks,
+          attention_mask: attention_mask,
+          attention_head_mask: attention_head_mask,
+          cross_attention_head_mask: cross_attention_head_mask,
+          cache: cache,
+          name: join(name, "blocks")
         ],
-        ffn: [
-          intermediate_size: spec.decoder_intermediate_size,
-          activation: spec.activation
-        ],
-        block_type: :norm_first,
-        name: join(name, "blocks")
+        fn hidden_state, block ->
+          name = block.name
+
+          attention_opts = attention_opts(spec, spec.decoder_num_attention_heads)
+
+          shortcut = hidden_state
+
+          {hidden_state, attention, self_attention_cache} =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "self_attention_norm"))
+            |> Layers.Transformer.self_attention(
+              block,
+              [causal: true, name: join(name, "self_attention")] ++ attention_opts
+            )
+
+          hidden_state =
+            hidden_state
+            |> Axon.dropout(rate: spec.dropout_rate)
+            |> Axon.add(shortcut)
+
+          shortcut = hidden_state
+
+          {hidden_state, cross_attention, cross_attention_cache} =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "cross_attention_norm"))
+            |> Layers.Transformer.cross_attention(
+              encoder_hidden_state,
+              block,
+              [name: join(name, "cross_attention")] ++ attention_opts
+            )
+
+          hidden_state =
+            hidden_state
+            |> Axon.dropout(rate: spec.dropout_rate)
+            |> Axon.add(shortcut)
+
+          shortcut = hidden_state
+
+          hidden_state =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "output_norm"))
+            |> Layers.Transformer.basic_ffn(spec.decoder_intermediate_size, spec.hidden_size,
+              activation: spec.activation,
+              dropout_rate: spec.dropout_rate,
+              kernel_initializer: kernel_initializer(spec),
+              name: join(name, "ffn")
+            )
+            |> Axon.add(shortcut)
+
+          %{
+            hidden_state: hidden_state,
+            attention: attention,
+            cross_attention: cross_attention,
+            self_attention_cache: self_attention_cache,
+            cross_attention_cache: cross_attention_cache
+          }
+        end
       )
 
     hidden_state = Axon.layer_norm(outputs.hidden_state, name: join(name, "norm"))
@@ -495,6 +562,16 @@ defmodule Bumblebee.Audio.Whisper do
     }
   end
 
+  defp attention_opts(spec, num_heads) do
+    [
+      num_heads: num_heads,
+      hidden_size: spec.hidden_size,
+      key_use_bias: false,
+      dropout_rate: spec.attention_dropout_rate,
+      kernel_initializer: kernel_initializer(spec)
+    ]
+  end
+
   defp kernel_initializer(spec) do
     Axon.Initializers.normal(scale: spec.initializer_scale)
   end

lib/bumblebee/diffusion/layers/unet.ex

@@ -314,20 +314,59 @@ defmodule Bumblebee.Diffusion.Layers.UNet do
       end
     end)
     |> Layers.Transformer.blocks(
-      cross_hidden_state: cross_hidden_state,
-      num_blocks: depth,
-      num_attention_heads: num_heads,
-      hidden_size: hidden_size,
-      query_use_bias: false,
-      key_use_bias: false,
-      value_use_bias: false,
-      layer_norm: [
-        epsilon: 1.0e-5
-      ],
-      dropout_rate: dropout,
-      ffn: &ffn_geglu(&1, 4 * hidden_size, hidden_size, dropout: dropout, name: &2),
-      block_type: :norm_first,
-      name: join(name, "blocks")
+      [num_blocks: depth, name: join(name, "blocks")],
+      fn hidden_state, block ->
+        name = block.name
+
+        attention_opts = [
+          num_heads: num_heads,
+          hidden_size: hidden_size,
+          query_use_bias: false,
+          key_use_bias: false,
+          value_use_bias: false
+        ]
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, _cache} =
+          hidden_state
+          |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "self_attention_norm"))
+          |> Layers.Transformer.self_attention(
+            block,
+            [name: join(name, "self_attention")] ++ attention_opts
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: dropout)
+          |> Axon.add(shortcut)
+
+        shortcut = hidden_state
+
+        {hidden_state, cross_attention, _cache} =
+          hidden_state
+          |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "cross_attention_norm"))
+          |> Layers.Transformer.cross_attention(
+            cross_hidden_state,
+            block,
+            [name: join(name, "cross_attention")] ++ attention_opts
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: dropout)
+          |> Axon.add(shortcut)
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "output_norm"))
+          |> ffn_geglu(4 * hidden_size, hidden_size, dropout: dropout, name: join(name, "ffn"))
+          |> Axon.add(shortcut)
+
+        %{hidden_state: hidden_state, attention: attention, cross_attention: cross_attention}
+      end
     )
     |> then(fn %{hidden_state: hidden_state} ->
       if use_linear_projection do