[Jetmoe] Fix RoPE (#40819)

* fix

* remove prints

* why was this there...

Author: vasqu

src/transformers/models/jetmoe/configuration_jetmoe.py

@@ -94,6 +94,7 @@ class JetMoeConfig(PretrainedConfig):
 
     model_type = "jetmoe"
     keys_to_ignore_at_inference = ["past_key_values"]
+    attribute_map = {"head_dim": "kv_channels"}
 
     def __init__(
         self,

src/transformers/models/jetmoe/modeling_jetmoe.py

@@ -936,15 +936,6 @@ def forward(
         if position_ids is None:
             position_ids = cache_position.unsqueeze(0)
 
-        if attention_mask is not None and self._attn_implementation == "flash_attention_2" and use_cache:
-            batch_size = inputs_embeds.shape[0]
-            is_padding_right = attention_mask[:, -1].sum().item() != batch_size
-            if is_padding_right:
-                raise ValueError(
-                    "You are attempting to perform batched generation with padding_side='right'"
-                    " this may lead to unexpected behaviour for Flash Attention version of JetMoe. Make sure to "
-                    " call `tokenizer.padding_side  = 'left'` before tokenizing the input. "
-                )
         causal_mask = self._update_causal_mask(
             attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
         )

tests/models/jetmoe/test_modeling_jetmoe.py

@@ -184,12 +184,10 @@ def test_model_8b_batched_generation(self):
         tokenizer = AutoTokenizer.from_pretrained("jetmoe/jetmoe-8b", use_fast=False)
         model = JetMoeForCausalLM.from_pretrained("jetmoe/jetmoe-8b")
         input_ids = tokenizer(prompt, return_tensors="pt", padding=True).to(model.model.embed_tokens.weight.device)
-        print(input_ids)
 
         # greedy generation outputs
         generated_ids = model.generate(**input_ids, max_new_tokens=10, temperature=0)
         text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
-        print(text)
         self.assertEqual(EXPECTED_TEXT_COMPLETION, text)
 
         del model