[fix:]: initialize ModernBertUnpaddedRotaryEmbedding when loading from a checkpoint

src/transformers/models/modernbert/modeling_modernbert.py

@@ -685,6 +685,12 @@ def init_weight(module: nn.Module, std: float):
                 curr_inv_freq, _ = rope_init_fn(module.config, layer_type=layer_type)
                 init.copy_(getattr(module, f"{layer_type}_inv_freq"), curr_inv_freq)
                 init.copy_(getattr(module, f"{layer_type}_original_inv_freq"), curr_inv_freq)
+        elif isinstance(module, ModernBertUnpaddedRotaryEmbedding):
+            inv_freq = module._compute_inv_freq()
+            module.register_buffer("inv_freq", inv_freq, persistent=False)
+            module._seq_len_cached = 0
+            module._cos_cached = None
+            module._sin_cached = None
 
     def _check_and_adjust_attn_implementation(
         self, attn_implementation: Optional[str], is_init_check: bool = False

src/transformers/models/modernbert/modular_modernbert.py

@@ -879,6 +879,12 @@ def init_weight(module: nn.Module, std: float):
                 curr_inv_freq, _ = rope_init_fn(module.config, layer_type=layer_type)
                 init.copy_(getattr(module, f"{layer_type}_inv_freq"), curr_inv_freq)
                 init.copy_(getattr(module, f"{layer_type}_original_inv_freq"), curr_inv_freq)
+        elif isinstance(module, ModernBertUnpaddedRotaryEmbedding):
+            inv_freq = module._compute_inv_freq()
+            module.register_buffer("inv_freq", inv_freq, persistent=False)
+            module._seq_len_cached = 0
+            module._cos_cached = None
+            module._sin_cached = None
 
     def _check_and_adjust_attn_implementation(
         self, attn_implementation: Optional[str], is_init_check: bool = False

tests/models/modernbert/test_modeling_modernbert.py

@@ -27,6 +27,7 @@
     require_flash_attn,
     require_torch,
     require_torch_accelerator,
+    require_torch_gpu,
     slow,
     torch_device,
 )
@@ -599,6 +600,63 @@ def test_inference_sequence_classification(self):
         expected = torch.tensor([[1.6466, 4.5662]])
         torch.testing.assert_close(output, expected, rtol=1e-4, atol=1e-4)
 
+    @slow
+    @require_flash_attn
+    @require_torch_gpu
+    def test_inference_sequence_classification_flash_attention_2(self):
+        if version.parse(torch.__version__) < version.parse("2.4.0"):
+            self.skipTest(reason="This test requires torch >= 2.4 to run.")
+
+        device = torch.device("cuda")
+
+        model = ModernBertForSequenceClassification.from_pretrained(
+            "hf-internal-testing/tiny-random-ModernBertForSequenceClassification",
+            reference_compile=False,
+            attn_implementation="flash_attention_2",
+            device_map=device,
+        )
+        tokenizer = AutoTokenizer.from_pretrained(
+            "hf-internal-testing/tiny-random-ModernBertForSequenceClassification"
+        )
+
+        inputs = tokenizer("Hello World!", return_tensors="pt")
+        inputs = {k: v.to(device=device) for k, v in inputs.items()}
+        with torch.no_grad():
+            output = model(**inputs)[0]
+            self.assertIsInstance(output, torch.Tensor)
+        expected_shape = torch.Size((1, 2))
+        self.assertEqual(output.shape, expected_shape)
+        self.assertFalse(torch.isnan(output).any().item())
+
+        expected = torch.tensor([[1.6466, 4.5662]], device=device)
+        torch.testing.assert_close(output, expected, rtol=1e-2, atol=1e-2)
+
+    @slow
+    @require_flash_attn
+    @require_torch_gpu
+    def test_inference_sequence_classification_flash_attention_2_modernbert_base(self):
+        if version.parse(torch.__version__) < version.parse("2.4.0"):
+            self.skipTest(reason="This test requires torch >= 2.4 to run.")
+
+        device = torch.device("cuda")
+
+        model = ModernBertForSequenceClassification.from_pretrained(
+            "answerdotai/ModernBERT-base",
+            reference_compile=False,
+            attn_implementation="flash_attention_2",
+            device_map=device,
+        )
+        tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")
+
+        inputs = tokenizer("Hello World!", return_tensors="pt")
+        inputs = {k: v.to(device=device) for k, v in inputs.items()}
+        with torch.no_grad():
+            output = model(**inputs)[0]
+            self.assertIsInstance(output, torch.Tensor)
+        expected_shape = torch.Size((1, 2))
+        self.assertEqual(output.shape, expected_shape)
+        self.assertFalse(torch.isnan(output).any().item())
+
     @pytest.mark.torch_export_test
     @slow
     def test_export(self):