Gemma4 model implementation

lib/bumblebee.ex

@@ -142,6 +142,8 @@ defmodule Bumblebee do
     "Gemma3TextForCausalLM" => {Bumblebee.Text.Gemma3Text, :for_causal_language_modeling},
     "Gemma3TextForSequenceClassification" =>
       {Bumblebee.Text.Gemma3Text, :for_sequence_classification},
+    "Gemma4ForConditionalGeneration" =>
+      {Bumblebee.Text.Gemma4Text, :for_causal_language_modeling},
     "GPT2ForSequenceClassification" => {Bumblebee.Text.Gpt2, :for_sequence_classification},
     "GPT2ForTokenClassification" => {Bumblebee.Text.Gpt2, :for_token_classification},
     "GPT2LMHeadModel" => {Bumblebee.Text.Gpt2, :for_causal_language_modeling},
@@ -273,6 +275,7 @@ defmodule Bumblebee do
     "clip" => :clip,
     "gemma" => :gemma,
     "gemma3_text" => :gemma,
+    "gemma4" => :gemma,
     "gpt_neox" => :gpt_neo_x,
     "gpt2" => :gpt2,
     "gpt_bigcode" => :gpt2,
@@ -777,11 +780,15 @@ defmodule Bumblebee do
   end
 
   defp params_file_loader_fun(".safetensors", opts) do
-    opts[:safetensors_reader] || (&Safetensors.read!(&1, lazy: true))
+    opts[:safetensors_reader] || (&read_safetensors_chunked/1)
   end
 
   defp params_file_loader_fun(_, _opts), do: &Bumblebee.Conversion.PyTorchLoader.load!/1
 
+  defp read_safetensors_chunked(path) do
+    Safetensors.read!(path, lazy: true)
+  end
+
   @doc """
   Featurizes `input` with the given featurizer.
 

lib/bumblebee/conversion/pytorch_params.ex

@@ -150,7 +150,7 @@ defmodule Bumblebee.Conversion.PyTorchParams do
 
             {value, diff} =
               if all_sources_found? do
-                source_values = Enum.map(source_values, &Nx.to_tensor/1)
+                source_values = Enum.map(source_values, &lazy_to_tensor/1)
                 value = builder_fun.(Enum.reverse(source_values))
 
                 case verify_param_shape(param_expr, value) do
@@ -188,6 +188,51 @@ defmodule Bumblebee.Conversion.PyTorchParams do
 
   defp prepend(diff, key, values), do: Map.update!(diff, key, &(values ++ &1))
 
+  # macOS pread(2) returns EINVAL when byte count > INT_MAX (~2 GB).
+  # For large safetensors tensors, read in 1 GB chunks instead.
+  @pread_chunk 1_073_741_824
+
+  defp lazy_to_tensor(%Safetensors.FileTensor{byte_size: size} = ft)
+       when size > @pread_chunk do
+    # Force BinaryBackend: the GPU backend (EMLX) cannot allocate tensors
+    # this large in a single call, and we must also avoid the macOS pread
+    # INT_MAX limit by reading in chunks.
+    Nx.with_default_backend(Nx.BinaryBackend, fn ->
+      File.open!(ft.path, [:read, :raw], fn file ->
+        binary = pread_chunked(file, ft.byte_offset, ft.byte_size)
+        Safetensors.Shared.build_tensor(binary, ft.shape, ft.type)
+      end)
+    end)
+  end
+
+  defp lazy_to_tensor(value), do: Nx.to_tensor(value)
+
+  defp pread_chunked(file, offset, size) when size <= @pread_chunk do
+    {:ok, binary} = :file.pread(file, offset, size)
+    binary
+  end
+
+  defp pread_chunked(file, offset, size) do
+    full = div(size, @pread_chunk)
+    rest = rem(size, @pread_chunk)
+
+    chunks =
+      for i <- 0..(full - 1) do
+        {:ok, chunk} = :file.pread(file, offset + i * @pread_chunk, @pread_chunk)
+        chunk
+      end
+
+    chunks =
+      if rest > 0 do
+        {:ok, tail} = :file.pread(file, offset + full * @pread_chunk, rest)
+        chunks ++ [tail]
+      else
+        chunks
+      end
+
+    IO.iodata_to_binary(chunks)
+  end
+
   defp infer_prefixes(layers, pytorch_state, params_mapping) do
     # Note: target refers to the parameters we are initializing, while
     # source refers to the state we are loading from

lib/bumblebee/layers.ex

@@ -1237,7 +1237,14 @@ defmodule Bumblebee.Layers do
   Adds a rotary embedding layer to the network.
   """
   def rotary_embedding(query, key, position_ids, attention_mask, size, opts \\ []) do
-    opts = Keyword.validate!(opts, [:name, :scaling_strategy, max_positions: 2048, base: 10_000])
+    opts =
+      Keyword.validate!(opts, [
+        :name,
+        :scaling_strategy,
+        :rotary_dim,
+        max_positions: 2048,
+        base: 10_000
+      ])
 
     output =
       Axon.layer(
@@ -1254,15 +1261,24 @@ defmodule Bumblebee.Layers do
                   max_positions,
                   size,
                   base,
-                  scaling_strategy
+                  scaling_strategy,
+                  rotary_dim \\ nil
                 ) do
     position = Nx.iota({sequence_length})
 
-    range = Nx.iota({div(size, 2)}) |> Nx.multiply(2) |> Nx.divide(size)
+    {num_freqs, denominator} =
+      if rotary_dim do
+        {div(rotary_dim, 2), size}
+      else
+        {div(size, 2), size}
+      end
+
+    range = Nx.iota({num_freqs}) |> Nx.multiply(2) |> Nx.divide(denominator)
 
     case scaling_strategy do
       %{type: :linear, factor: factor} ->
         inv_frequency = inv_frequency(base, range)
+        inv_frequency = maybe_pad_inv_frequency(inv_frequency, div(size, 2), rotary_dim)
         position = Nx.divide(position, factor)
         positions_cos_sin(position, inv_frequency)
 
@@ -1273,6 +1289,7 @@ defmodule Bumblebee.Layers do
           |> Nx.pow(size / (size - 2))
 
         inv_frequency = inv_frequency(base, range)
+        inv_frequency = maybe_pad_inv_frequency(inv_frequency, div(size, 2), rotary_dim)
         positions_cos_sin(position, inv_frequency)
 
       %{
@@ -1300,6 +1317,7 @@ defmodule Bumblebee.Layers do
           end
 
         inv_frequency = inv_frequency(base, range) |> Nx.divide(factor)
+        inv_frequency = maybe_pad_inv_frequency(inv_frequency, div(size, 2), rotary_dim)
         {cos, sin} = positions_cos_sin(position, inv_frequency)
         {Nx.multiply(cos, cos_sin_factor), Nx.multiply(sin, cos_sin_factor)}
 
@@ -1321,14 +1339,29 @@ defmodule Bumblebee.Layers do
             original_max_positions
           )
 
+        inv_frequency = maybe_pad_inv_frequency(inv_frequency, div(size, 2), rotary_dim)
         positions_cos_sin(position, inv_frequency)
 
       _other ->
         inv_frequency = inv_frequency(base, range)
+        inv_frequency = maybe_pad_inv_frequency(inv_frequency, div(size, 2), rotary_dim)
         positions_cos_sin(position, inv_frequency)
     end
   end
 
+  defp maybe_pad_inv_frequency(inv_frequency, _target_size, nil), do: inv_frequency
+
+  defp maybe_pad_inv_frequency(inv_frequency, target_size, _rotary_dim) do
+    pad_size = target_size - Nx.axis_size(inv_frequency, 0)
+
+    if pad_size > 0 do
+      padding = Nx.broadcast(Nx.tensor(0.0, type: Nx.type(inv_frequency)), {pad_size})
+      Nx.concatenate([inv_frequency, padding])
+    else
+      inv_frequency
+    end
+  end
+
   defnp llama3_inv_frequency(
           inv_frequency,
           factor,
@@ -1381,6 +1414,7 @@ defmodule Bumblebee.Layers do
       keyword!(opts, [
         :size,
         :scaling_strategy,
+        :rotary_dim,
         mode: :inference,
         max_positions: 2048,
         base: 10_000
@@ -1400,7 +1434,8 @@ defmodule Bumblebee.Layers do
         opts[:max_positions],
         opts[:size],
         opts[:base],
-        opts[:scaling_strategy]
+        opts[:scaling_strategy],
+        opts[:rotary_dim]
       )
 
     position_ids = Nx.as_type(position_ids, :s64)

lib/bumblebee/layers/transformer.ex

@@ -52,7 +52,6 @@ defmodule Bumblebee.Layers.Transformer do
       :hidden_size,
       :ffn,
       :kernel_initializer,
-      :attention_head_size,
       :dropout_rate,
       :attention_dropout_rate,
       :query_use_bias,
@@ -75,6 +74,7 @@ defmodule Bumblebee.Layers.Transformer do
             :num_blocks,
             :rotary_embedding,
             :attention_window_size,
+            :attention_head_size,
             attention_mask: Layers.none(),
             attention_head_mask: Layers.none(),
             attention_relative_bias: nil,
@@ -97,6 +97,7 @@ defmodule Bumblebee.Layers.Transformer do
     cache = opts[:cache]
     rotary_embedding = opts[:rotary_embedding]
     attention_window_size = opts[:attention_window_size]
+    attention_head_size = opts[:attention_head_size]
 
     block_opts = Keyword.take(opts, block_opts_keys)
 
@@ -142,12 +143,20 @@ defmodule Bumblebee.Layers.Transformer do
               size -> size
             end
 
+          block_attention_head_size =
+            case attention_head_size do
+              nil -> nil
+              fun when is_function(fun, 1) -> fun.(idx)
+              size -> size
+            end
+
           {hidden_state, attention, cross_attention, block_cache, attention_relative_bias} =
             block(
               state.hidden_state,
               [
                 attention_mask: attention_mask,
                 attention_head_mask: block_attention_head_mask,
+                attention_head_size: block_attention_head_size,
                 attention_relative_bias: attention_relative_bias,
                 cross_hidden_state: cross_hidden_state,
                 cross_attention_mask: cross_attention_mask,
@@ -856,6 +865,7 @@ defmodule Bumblebee.Layers.Transformer do
               :position_ids,
               :max_positions,
               :scaling_strategy,
+              :rotary_dim,
               base: 10_000,
               percentage: 1.0
             ])

lib/bumblebee/text/gemma4_chat.ex

@@ -0,0 +1,86 @@
+defmodule Bumblebee.Text.Gemma4Chat do
+  @moduledoc """
+  Builds Gemma 4 chat prompts.
+
+  Gemma 4 checkpoints ship a `tokenizer.json` whose post-processor adds nothing, so no
+  `<bos>` is prepended even with `add_special_tokens: true`. The model degenerates badly
+  without it, so the prompt string must carry both `<bos>` and the turn markers. This
+  module mirrors the `chat_template.jinja` shipped with the official checkpoints for the
+  text-only cases.
+  """
+
+  @typedoc """
+  A single chat message.
+
+  The role is one of `:system`, `:user` or `:assistant` (`:developer` is treated as
+  `:system`, and `:model` as `:assistant`); string roles are accepted too.
+  """
+  @type message :: %{required(:role) => atom() | String.t(), required(:content) => String.t()}
+
+  @doc """
+  Formats `messages` into a Gemma 4 prompt string.
+
+  ## Options
+
+    * `:add_generation_prompt` - append the `<|turn>model\\n` cue that tells the model to
+      start an assistant reply. Defaults to `true`
+
+    * `:enable_thinking` - emit the `<|think|>` control token, which turns on the model's
+      reasoning mode. Defaults to `false`
+
+  ## Examples
+
+      iex> Bumblebee.Text.Gemma4Chat.format([%{role: :user, content: "Hi there"}])
+      "<bos><|turn>user\\nHi there<turn|>\\n<|turn>model\\n"
+
+      iex> Bumblebee.Text.Gemma4Chat.format(
+      ...>   [%{role: :system, content: "You are terse."}, %{role: :user, content: "Hi"}]
+      ...> )
+      "<bos><|turn>system\\nYou are terse.<turn|>\\n<|turn>user\\nHi<turn|>\\n<|turn>model\\n"
+
+  A system turn is emitted whenever there is system content or thinking is enabled, since
+  the `<|think|>` token must appear at the top of the first system turn:
+
+      iex> Bumblebee.Text.Gemma4Chat.format(
+      ...>   [%{role: :user, content: "Hi"}], enable_thinking: true
+      ...> )
+      "<bos><|turn>system\\n<|think|>\\n<turn|>\\n<|turn>user\\nHi<turn|>\\n<|turn>model\\n"
+
+  """
+  @spec format([message()], keyword()) :: String.t()
+  def format(messages, opts \\ []) do
+    opts = Keyword.validate!(opts, add_generation_prompt: true, enable_thinking: false)
+    thinking? = opts[:enable_thinking]
+
+    {system, rest} = split_system(messages)
+
+    header =
+      if thinking? or system do
+        "<|turn>system\n" <>
+          if(thinking?, do: "<|think|>\n", else: "") <>
+          (system || "") <> "<turn|>\n"
+      else
+        ""
+      end
+
+    turns =
+      Enum.map_join(rest, fn %{role: role, content: content} ->
+        "<|turn>#{role_name(role)}\n#{String.trim(content)}<turn|>\n"
+      end)
+
+    generation_prompt = if opts[:add_generation_prompt], do: "<|turn>model\n", else: ""
+
+    "<bos>" <> header <> turns <> generation_prompt
+  end
+
+  defp split_system([%{role: role, content: content} | rest])
+       when role in [:system, "system", :developer, "developer"] do
+    {String.trim(content), rest}
+  end
+
+  defp split_system(messages), do: {nil, messages}
+
+  defp role_name(role) when role in [:assistant, "assistant", :model, "model"], do: "model"
+  defp role_name(role) when is_atom(role), do: Atom.to_string(role)
+  defp role_name(role) when is_binary(role), do: role
+end