Add Gemma (#358)

Co-authored-by: Jonatan Kłosko <[email protected]>
elixir-nx · Mar 4, 2024 · 20855ad · 20855ad
1 parent a09b230
commit 20855ad
Show file tree

Hide file tree

Showing 5 changed files with 547 additions and 3 deletions.
diff --git a/lib/bumblebee.ex b/lib/bumblebee.ex
@@ -133,6 +133,9 @@ defmodule Bumblebee do
     "DistilBertForQuestionAnswering" => {Bumblebee.Text.Distilbert, :for_question_answering},
     "DistilBertForTokenClassification" => {Bumblebee.Text.Distilbert, :for_token_classification},
     "DistilBertForMultipleChoice" => {Bumblebee.Text.Distilbert, :for_multiple_choice},
+    "GemmaModel" => {Bumblebee.Text.Gemma, :base},
+    "GemmaForCausalLM" => {Bumblebee.Text.Gemma, :for_causal_language_modeling},
+    "GemmaForSequenceClassification" => {Bumblebee.Text.Gemma, :for_sequence_classification},
     "GPT2ForSequenceClassification" => {Bumblebee.Text.Gpt2, :for_sequence_classification},
     "GPT2ForTokenClassification" => {Bumblebee.Text.Gpt2, :for_token_classification},
     "GPT2LMHeadModel" => {Bumblebee.Text.Gpt2, :for_causal_language_modeling},
@@ -231,6 +234,7 @@ defmodule Bumblebee do
     "distilbert" => :distilbert,
     "camembert" => :camembert,
     "clip" => :clip,
+    "gemma" => :gemma,
     "gpt_neox" => :gpt_neo_x,
     "gpt2" => :gpt2,
     "gpt_bigcode" => :gpt2,

diff --git a/lib/bumblebee/layers.ex b/lib/bumblebee/layers.ex
@@ -1091,7 +1091,13 @@ defmodule Bumblebee.Layers do
   # TODO: Add to Axon
   def rms_norm(input, opts \\ []) do
     opts =
-      Keyword.validate!(opts, [:name, channel_index: -1, epsilon: 1.0e-6, initializer: :ones])
+      Keyword.validate!(opts, [
+        :name,
+        shift: 0.0,
+        channel_index: -1,
+        epsilon: 1.0e-6,
+        initializer: :ones
+      ])
 
     weight =
       Axon.param("weight", &Axon.Shape.norm_param(&1, opts[:channel_index]),
@@ -1100,13 +1106,14 @@ defmodule Bumblebee.Layers do
 
     Axon.layer(&rms_norm_impl/3, [input, weight],
       name: opts[:name],
+      shift: opts[:shift],
       epsilon: opts[:epsilon],
       op_name: :rms_norm
     )
   end
 
   defnp rms_norm_impl(input, weight, opts \\ []) do
-    opts = keyword!(opts, epsilon: 1.0e-6, channel_index: -1, mode: :train)
+    opts = keyword!(opts, shift: 0.0, epsilon: 1.0e-6, channel_index: -1, mode: :train)
 
     variance =
       input
@@ -1117,7 +1124,7 @@ defmodule Bumblebee.Layers do
       input
       |> Nx.multiply(Nx.rsqrt(variance + opts[:epsilon]))
 
-    x * weight
+    x * (opts[:shift] + weight)
   end
 
   @doc """