Add Llama

lib/bumblebee.ex

@@ -125,6 +125,9 @@ defmodule Bumblebee do
     "LayoutLMForTokenClassification" =>
       {Bumblebee.Multimodal.LayoutLm, :for_token_classification},
     "LayoutLMModel" => {Bumblebee.Multimodal.LayoutLm, :base},
+    "LlamaModel" => {Bumblebee.Text.Llama, :base},
+    "LlamaForCausalLM" => {Bumblebee.Text.Llama, :for_causal_language_modeling},
+    "LlamaForSequenceClassification" => {Bumblebee.Text.Llama, :for_sequence_classification},
     "MBartForCausalLM" => {Bumblebee.Text.Mbart, :for_causal_language_modeling},
     "MBartForConditionalGeneration" => {Bumblebee.Text.Mbart, :for_conditional_generation},
     "MBartForQuestionAnswering" => {Bumblebee.Text.Mbart, :for_question_answering},
@@ -194,6 +197,7 @@ defmodule Bumblebee do
     "clip" => Bumblebee.Text.ClipTokenizer,
     "gpt2" => Bumblebee.Text.Gpt2Tokenizer,
     "layoutlm" => Bumblebee.Text.LayoutLmTokenizer,
+    "llama" => Bumblebee.Text.LlamaTokenizer,
     "mbart" => Bumblebee.Text.MbartTokenizer,
     "roberta" => Bumblebee.Text.RobertaTokenizer,
     "t5" => Bumblebee.Text.T5Tokenizer,

lib/bumblebee/layers.ex

@@ -965,4 +965,51 @@ defmodule Bumblebee.Layers do
 
     x * weight
   end
+
+  @doc """
+  Adds a rotary embedding layer to the network.
+  """
+  def rotary_embedding(query, key, position_ids, size, opts \\ []) do
+    opts = Keyword.validate!(opts, [:name, max_positions: 2048, base: 10_000])
+
+    output =
+      Axon.layer(&apply_rotary_embedding/4, [query, key, position_ids], [size: size] ++ opts)
+
+    unwrap_tuple(output, 2)
+  end
+
+  deftransformp create_sinusoidal_positions(max_positions, size, base) do
+    range = Nx.multiply(Nx.iota({div(size, 2)}), 2)
+    inv_frequency = Nx.divide(1.0, Nx.pow(base, range))
+
+    position = Nx.iota({max_positions})
+    angle = Nx.outer(position, inv_frequency)
+
+    angle = Nx.concatenate([angle, angle], axis: -1)
+
+    {Nx.cos(angle), Nx.sin(angle)}
+  end
+
+  defnp apply_rotary_embedding(query, key, position_ids, opts \\ []) do
+    opts = keyword!(opts, [:size, mode: :inference, max_positions: 2048, base: 10_000])
+
+    {cos, sin} = create_sinusoidal_positions(opts[:max_positions], opts[:size], opts[:base])
+
+    position_ids = Nx.as_type(position_ids, :s64)
+
+    cos = cos |> Nx.take(position_ids) |> Nx.new_axis(2)
+    sin = sin |> Nx.take(position_ids) |> Nx.new_axis(2)
+
+    rotated_query = query * cos + rotate_half(query) * sin
+    rotated_key = key * cos + rotate_half(key) * sin
+
+    {rotated_query, rotated_key}
+  end
+
+  defnp rotate_half(x) do
+    size = div(Nx.axis_size(x, -1), 2)
+    x1 = x[[.., .., .., 0..(size - 1)//1]]
+    x2 = x[[.., .., .., size..-1//1]]
+    Nx.concatenate([-x2, x1], axis: -1)
+  end
 end

lib/bumblebee/layers/transformer.ex

@@ -56,7 +56,8 @@ defmodule Bumblebee.Layers.Transformer do
       :layer_norm,
       :norm_placement,
       :output_shortcut,
-      :scale_query?
+      :scale_query?,
+      :rotary_embedding
     ]
 
     opts =
@@ -260,6 +261,15 @@ defmodule Bumblebee.Layers.Transformer do
     * `:scale_query?` - whether to scale query in the traditional style of
       multi-headed attention. Defaults to `true`
 
+    * `:rotary_embedding` - configuration of rotary embedding. If set,
+      will apply rotary position embedding with the given options. Valid
+      options are:
+
+        * `:position_ids` (required) - input position ids used for the
+          embedding
+
+        * `:max_positions` - the maximum number of distinct positions
+
     * `:name` - the prefix for layer names
 
   ## References
@@ -296,7 +306,8 @@ defmodule Bumblebee.Layers.Transformer do
         norm_placement: :last,
         layer_norm: [],
         output_shortcut: true,
-        scale_query?: true
+        scale_query?: true,
+        rotary_embedding: nil
       ])
 
     name = opts[:name]
@@ -324,6 +335,7 @@ defmodule Bumblebee.Layers.Transformer do
     norm_placement = opts[:norm_placement]
     output_shortcut = opts[:output_shortcut]
     scale_query? = opts[:scale_query?]
+    rotary_embedding = opts[:rotary_embedding]
 
     ffn_fun =
       case ffn do
@@ -384,6 +396,7 @@ defmodule Bumblebee.Layers.Transformer do
         value_use_bias: value_use_bias,
         output_use_bias: output_use_bias,
         scale_query?: scale_query?,
+        rotary_embedding: rotary_embedding,
         name: join(name, "self_attention")
       )
 
@@ -424,6 +437,7 @@ defmodule Bumblebee.Layers.Transformer do
               value_use_bias: value_use_bias,
               output_use_bias: output_use_bias,
               scale_query?: scale_query?,
+              rotary_embedding: rotary_embedding,
               name: join(name, "cross_attention")
             )
 
@@ -549,6 +563,15 @@ defmodule Bumblebee.Layers.Transformer do
     * `:output_use_bias` - whether to use bias in the output projection.
       Defaults to `true`
 
+    * `:rotary_embedding` - configuration of rotary embedding. If set,
+      will apply rotary position embedding with the given options. Valid
+      options are:
+
+        * `:position_ids` (required) - input position ids used for the
+          embedding
+
+        * `:max_positions` - the maximum number of distinct positions
+
     * `:name` - the prefix for layer names
 
   ## References
@@ -577,7 +600,8 @@ defmodule Bumblebee.Layers.Transformer do
         query_use_bias: true,
         key_use_bias: true,
         value_use_bias: true,
-        output_use_bias: true
+        output_use_bias: true,
+        rotary_embedding: nil
       ])
 
     attention_mask = opts[:attention_mask]
@@ -592,6 +616,7 @@ defmodule Bumblebee.Layers.Transformer do
     causal? = opts[:causal?]
     scale_query? = opts[:scale_query?]
     dropout_rate = opts[:dropout_rate]
+    rotary_embedding = opts[:rotary_embedding]
 
     query_use_bias = opts[:query_use_bias]
     key_use_bias = opts[:key_use_bias]
@@ -607,6 +632,8 @@ defmodule Bumblebee.Layers.Transformer do
         hidden_size
       end
 
+    head_size = div(hidden_size, num_heads)
+
     query =
       query
       |> Axon.dense(inner_size,
@@ -634,6 +661,25 @@ defmodule Bumblebee.Layers.Transformer do
       )
       |> Layers.split_heads(num_heads)
 
+    {query, key} =
+      case rotary_embedding do
+        opts when is_list(opts) ->
+          validate_required_keys!(opts, [:position_ids])
+          opts = Keyword.validate!(opts, [:position_ids, :max_positions])
+          {position_ids, opts} = Keyword.pop(opts, :position_ids)
+
+          Layers.rotary_embedding(
+            query,
+            key,
+            position_ids,
+            head_size,
+            [name: join(name, "rotary_embedding")] ++ opts
+          )
+
+        nil ->
+          {query, key}
+      end
+
     {key, value, attention_cache} =
       Layers.Decoder.cached_attention_key_values(key, value, attention_cache, offset)