Upload modeling_gemmoe.py

modeling_gemmoe.py

@@ -13,7 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-""" PyTorch Gemma model."""
+""" PyTorch Gemmoe model."""
 
 import math
 import warnings
@@ -43,7 +43,7 @@ from transformers.utils import (
     replace_return_docstrings,
 )
 from transformers.utils.import_utils import is_torch_fx_available
-from .configuration_gemma import GemmaConfig
+from .configuration_gemmoe import GemmoeConfig
 
 
 if is_flash_attn_2_available():
@@ -62,7 +62,7 @@ if is_torch_fx_available():
 
 logger = logging.get_logger(__name__)
 
-_CONFIG_FOR_DOC = "GemmaConfig"
+_CONFIG_FOR_DOC = "GemmoeConfig"
 
 def approx_gelu(x):
     return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * x**3)))
@@ -79,7 +79,7 @@ def _get_unpad_data(attention_mask):
     )
 
 
-class GemmaRMSNorm(nn.Module):
+class GemmoeRMSNorm(nn.Module):
     def __init__(self, dim: int, eps: float = 1e-6):
         super().__init__()
         self.eps = eps
@@ -97,10 +97,10 @@ class GemmaRMSNorm(nn.Module):
         return normed_x * (self.weight + 1)
 
 
-ALL_LAYERNORM_LAYERS.append(GemmaRMSNorm)
+ALL_LAYERNORM_LAYERS.append(GemmoeRMSNorm)
 
 
-class GemmaRotaryEmbedding(nn.Module):
+class GemmoeRotaryEmbedding(nn.Module):
     def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
         super().__init__()
         self.dim = dim
@@ -164,8 +164,8 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
     return q_embed, k_embed
 
 
-# Copied from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->Gemma
-class GemmaMLP(nn.Module):
+# Copied from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->Gemmoe
+class GemmoeMLP(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.config = config
@@ -193,11 +193,11 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
     return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
 
 
-class GemmaAttention(nn.Module):
+class GemmoeAttention(nn.Module):
     """Multi-headed attention from 'Attention Is All You Need' paper"""
 
     # Ignore copy
-    def __init__(self, config: GemmaConfig, layer_idx: Optional[int] = None):
+    def __init__(self, config: GemmoeConfig, layer_idx: Optional[int] = None):
         super().__init__()
         self.config = config
         self.layer_idx = layer_idx
@@ -228,7 +228,7 @@ class GemmaAttention(nn.Module):
         self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
         self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
         self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)
-        self.rotary_emb = GemmaRotaryEmbedding(
+        self.rotary_emb = GemmoeRotaryEmbedding(
             self.head_dim,
             max_position_embeddings=self.max_position_embeddings,
             base=self.rope_theta,
@@ -298,10 +298,10 @@ class GemmaAttention(nn.Module):
         return attn_output, attn_weights, past_key_value
 
 
-# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2 with Llama->Gemma
-class GemmaFlashAttention2(GemmaAttention):
+# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2 with Llama->Gemmoe
+class GemmoeFlashAttention2(GemmoeAttention):
     """
-    Gemma flash attention module. This module inherits from `GemmaAttention` as the weights of the module stays
+    Gemmoe flash attention module. This module inherits from `GemmoeAttention` as the weights of the module stays
     untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
     flash attention and deal with padding tokens in case the input contains any of them.
     """
@@ -363,7 +363,7 @@ class GemmaFlashAttention2(GemmaAttention):
         # therefore the input hidden states gets silently casted in float32. Hence, we need
         # cast them back in the correct dtype just to be sure everything works as expected.
         # This might slowdown training & inference so it is recommended to not cast the LayerNorms
-        # in fp32. (GemmaRMSNorm handles it correctly)
+        # in fp32. (GemmoeRMSNorm handles it correctly)
 
         input_dtype = query_states.dtype
         if input_dtype == torch.float32:
@@ -422,7 +422,7 @@ class GemmaFlashAttention2(GemmaAttention):
         if not self._flash_attn_uses_top_left_mask:
             causal = self.is_causal
         else:
-            # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in GemmaFlashAttention2 __init__.
+            # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in GemmoeFlashAttention2 __init__.
             causal = self.is_causal and query_length != 1
 
         # Contains at least one padding token in the sequence
@@ -495,11 +495,11 @@ class GemmaFlashAttention2(GemmaAttention):
         )
 
 
-# Copied from transformers.models.llama.modeling_llama.LlamaSdpaAttention with Llama->Gemma
-class GemmaSdpaAttention(GemmaAttention):
+# Copied from transformers.models.llama.modeling_llama.LlamaSdpaAttention with Llama->Gemmoe
+class GemmoeSdpaAttention(GemmoeAttention):
     """
-    Gemma attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
-    `GemmaAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
+    Gemmoe attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
+    `GemmoeAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
     SDPA API.
     """
 
@@ -517,7 +517,7 @@ class GemmaSdpaAttention(GemmaAttention):
         if output_attentions:
             # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
             logger.warning_once(
-                "GemmaModel is using GemmaSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
+                "GemmoeModel is using GemmoeSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
                 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
             )
             return super().forward(
@@ -580,24 +580,24 @@ class GemmaSdpaAttention(GemmaAttention):
         return attn_output, None, past_key_value
 
 
-GEMMA_ATTENTION_CLASSES = {
-    "eager": GemmaAttention,
-    "flash_attention_2": GemmaFlashAttention2,
-    "sdpa": GemmaSdpaAttention,
+GEMMOE_ATTENTION_CLASSES = {
+    "eager": GemmoeAttention,
+    "flash_attention_2": GemmoeFlashAttention2,
+    "sdpa": GemmoeSdpaAttention,
 }
 
 
-# Copied from transformers.models.llama.modeling_llama.LlamaDecoderLayer with LLAMA->GEMMA,Llama->Gemma
-class GemmaDecoderLayer(nn.Module):
-    def __init__(self, config: GemmaConfig, layer_idx: int):
+# Copied from transformers.models.llama.modeling_llama.LlamaDecoderLayer with LLAMA->GEMMOE,Llama->Gemmoe
+class GemmoeDecoderLayer(nn.Module):
+    def __init__(self, config: GemmoeConfig, layer_idx: int):
         super().__init__()
         self.hidden_size = config.hidden_size
 
-        self.self_attn = GEMMA_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)
+        self.self_attn = GEMMOE_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)
 
-        self.mlp = GemmaMLP(config)
-        self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.mlp = GemmoeMLP(config)
+        self.input_layernorm = GemmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = GemmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
     def forward(
         self,
@@ -663,7 +663,7 @@ class GemmaDecoderLayer(nn.Module):
         return outputs
 
 
-GEMMA_START_DOCSTRING = r"""
+GEMMOE_START_DOCSTRING = r"""
     This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
     library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
     etc.)
@@ -673,7 +673,7 @@ GEMMA_START_DOCSTRING = r"""
     and behavior.
 
     Parameters:
-        config ([`GemmaConfig`]):
+        config ([`GemmoeConfig`]):
             Model configuration class with all the parameters of the model. Initializing with a config file does not
             load the weights associated with the model, only the configuration. Check out the
             [`~PreTrainedModel.from_pretrained`] method to load the model weights.
@@ -681,15 +681,15 @@ GEMMA_START_DOCSTRING = r"""
 
 
 @add_start_docstrings(
-    "The bare Gemma Model outputting raw hidden-states without any specific head on top.",
-    GEMMA_START_DOCSTRING,
+    "The bare Gemmoe Model outputting raw hidden-states without any specific head on top.",
+    GEMMOE_START_DOCSTRING,
 )
-class GemmaPreTrainedModel(PreTrainedModel):
-    config_class = GemmaConfig
+class GemmoePreTrainedModel(PreTrainedModel):
+    config_class = GemmoeConfig
     base_model_prefix = "model"
     supports_gradient_checkpointing = True
     _keep_in_fp32_modules = ["inv_freq", "rotary_emb", "cos_cached", "sin_cached"]
-    _no_split_modules = ["GemmaDecoderLayer"]
+    _no_split_modules = ["GemmoeDecoderLayer"]
     _skip_keys_device_placement = ["past_key_values", "causal_mask"]
     _supports_flash_attn_2 = True
     _supports_sdpa = True
@@ -728,7 +728,7 @@ class GemmaPreTrainedModel(PreTrainedModel):
             layer.self_attn.past_key_value = None
 
 
-GEMMA_INPUTS_DOCSTRING = r"""
+GEMMOE_INPUTS_DOCSTRING = r"""
     Args:
         input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
             Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
@@ -803,28 +803,28 @@ GEMMA_INPUTS_DOCSTRING = r"""
 
 
 @add_start_docstrings(
-    "The bare Gemma Model outputting raw hidden-states without any specific head on top.",
-    GEMMA_START_DOCSTRING,
+    "The bare Gemmoe Model outputting raw hidden-states without any specific head on top.",
+    GEMMOE_START_DOCSTRING,
 )
-# Copied from transformers.models.llama.modeling_llama.LlamaModel with LLAMA->GEMMA,Llama->Gemma
-class GemmaModel(GemmaPreTrainedModel):
+# Copied from transformers.models.llama.modeling_llama.LlamaModel with LLAMA->GEMMOE,Llama->Gemmoe
+class GemmoeModel(GemmoePreTrainedModel):
     """
-    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`GemmaDecoderLayer`]
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`GemmoeDecoderLayer`]
 
     Args:
-        config: GemmaConfig
+        config: GemmoeConfig
     """
 
-    def __init__(self, config: GemmaConfig):
+    def __init__(self, config: GemmoeConfig):
         super().__init__(config)
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
 
         self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
         self.layers = nn.ModuleList(
-            [GemmaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+            [GemmoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
         )
-        self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.norm = GemmoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.gradient_checkpointing = False
 
         # Register a causal mask to separate causal and padding mask creation. Merging happens in the attention class.
@@ -842,7 +842,7 @@ class GemmaModel(GemmaPreTrainedModel):
     def set_input_embeddings(self, value):
         self.embed_tokens = value
 
-    @add_start_docstrings_to_model_forward(GEMMA_INPUTS_DOCSTRING)
+    @add_start_docstrings_to_model_forward(GEMMOE_INPUTS_DOCSTRING)
     # Ignore copy
     def forward(
         self,
@@ -1021,13 +1021,13 @@ class GemmaModel(GemmaPreTrainedModel):
         return causal_mask
 
 
-# Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM with LLAMA->GEMMA,Llama->Gemma,llama->gemma
-class GemmaForCausalLM(GemmaPreTrainedModel):
+# Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM with LLAMA->GEMMOE,Llama->Gemmoe,llama->gemmoe
+class GemmoeForCausalLM(GemmoePreTrainedModel):
     _tied_weights_keys = ["lm_head.weight"]
 
     def __init__(self, config):
         super().__init__(config)
-        self.model = GemmaModel(config)
+        self.model = GemmoeModel(config)
         self.vocab_size = config.vocab_size
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
 
@@ -1053,7 +1053,7 @@ class GemmaForCausalLM(GemmaPreTrainedModel):
         return self.model
 
     # Ignore copy
-    @add_start_docstrings_to_model_forward(GEMMA_INPUTS_DOCSTRING)
+    @add_start_docstrings_to_model_forward(GEMMOE_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
     def forward(
         self,
@@ -1081,10 +1081,10 @@ class GemmaForCausalLM(GemmaPreTrainedModel):
         Example:
 
         ```python
-        >>> from transformers import AutoTokenizer, GemmaForCausalLM
+        >>> from transformers import AutoTokenizer, GemmoeForCausalLM
 
-        >>> model = GemmaForCausalLM.from_pretrained("google/gemma-7b")
-        >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b")
+        >>> model = GemmoeForCausalLM.from_pretrained("google/gemmoe-7b")
+        >>> tokenizer = AutoTokenizer.from_pretrained("google/gemmoe-7b")
 
         >>> prompt = "What is your favorite condiment?"
         >>> inputs = tokenizer(prompt, return_tensors="pt")
@@ -1236,9 +1236,9 @@ class GemmaForCausalLM(GemmaPreTrainedModel):
 
 @add_start_docstrings(
     """
-    The Gemma Model transformer with a sequence classification head on top (linear layer).
+    The Gemmoe Model transformer with a sequence classification head on top (linear layer).
 
-    [`GemmaForSequenceClassification`] uses the last token in order to do the classification, as other causal models
+    [`GemmoeForSequenceClassification`] uses the last token in order to do the classification, as other causal models
     (e.g. GPT-2) do.
 
     Since it does classification on the last token, it requires to know the position of the last token. If a
@@ -1247,14 +1247,14 @@ class GemmaForCausalLM(GemmaPreTrainedModel):
     padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
     each row of the batch).
     """,
-    GEMMA_START_DOCSTRING,
+    GEMMOE_START_DOCSTRING,
 )
-# Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with LLAMA->GEMMA,Llama->Gemma
-class GemmaForSequenceClassification(GemmaPreTrainedModel):
+# Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with LLAMA->GEMMOE,Llama->Gemmoe
+class GemmoeForSequenceClassification(GemmoePreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels
-        self.model = GemmaModel(config)
+        self.model = GemmoeModel(config)
         self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
 
         # Initialize weights and apply final processing
@@ -1266,7 +1266,7 @@ class GemmaForSequenceClassification(GemmaPreTrainedModel):
     def set_input_embeddings(self, value):
         self.model.embed_tokens = value
 
-    @add_start_docstrings_to_model_forward(GEMMA_INPUTS_DOCSTRING)
+    @add_start_docstrings_to_model_forward(GEMMOE_INPUTS_DOCSTRING)
     def forward(
         self,
         input_ids: torch.LongTensor = None,