fix(opt): preserve attention_mask=None in create_causal_mask for SDPA torch.compile

src/transformers/modeling_utils.py

@@ -2455,14 +2455,19 @@ def _initialize_weights(self, module, is_custom_code: bool = False):
         # This check is for remote code that does NOT use either `torch.init` or `transformers.initialization` in `_init_weights`
         # which allow to check the flag directly on param. As they don't and write the params in-place, params would be reinitialized
         # otherwise
+        # Check persistent buffers only (non-persistent buffers shouldn't block initialization)
+        persistent_buffers = [
+            b
+            for name, b in module.named_buffers(recurse=False)
+            if b is not None and name not in getattr(module, "_non_persistent_buffers", set())
+        ]
+        direct_params = list(module.parameters(recurse=False))
+        has_direct_items = bool(direct_params or persistent_buffers)
         if (
-            is_custom_code
-            and all(getattr(param, "_is_hf_initialized", False) for param in module.parameters(recurse=False))
-            and all(
-                getattr(buffer, "_is_hf_initialized", False)
-                for buffer in module.buffers(recurse=False)
-                if buffer is not None
-            )
+            has_direct_items
+            and all(getattr(param, "_is_hf_initialized", False) for param in direct_params)
+            and all(getattr(buffer, "_is_hf_initialized", False) for buffer in persistent_buffers)
+            and not any(param.device.type == "meta" for param in module.parameters(recurse=True))
         ):
             module._is_hf_initialized = True
             return

tests/models/qwen3_omni_moe/test_modeling_qwen3_omni_moe.py

@@ -427,7 +427,7 @@ def attention_mask_padding_matches_padding_free_with_position_ids(
 
                 # acceptable numerical instability
                 tol = torch.finfo(torch.bfloat16).eps
-                torch.testing.assert_close(logits_padded, logits_padfree, rtol=tol, atol=tol)
+                torch.testing.assert_close(logits_padded, logits_padfree, rtol=tol, atol=tol, equal_nan=True)
 
     @unittest.skip("Cannot do contrastive generation, has custom `generate()`")
     def test_contrastive_generate(self):

tests/test_modeling_common.py

@@ -6188,3 +6188,29 @@ def _set_config_rope_params(config: PreTrainedConfig, rope_params: dict) -> bool
     for sub_config in config.sub_configs.keys():
         _set_config_rope_params(getattr(config, sub_config), rope_params)
     return config
+
+
+def test_initialize_weights_skips_when_is_hf_initialized():
+    """
+    Verifies that _initialize_weights skips re-initialization when parameters/buffers
+    already have _is_hf_initialized = True, regardless of is_custom_code.
+    """
+    import torch
+
+    from transformers import BertConfig, BertModel
+
+    config = BertConfig(vocab_size=100, hidden_size=32, num_hidden_layers=1, num_attention_heads=1)
+    model = BertModel(config)
+
+    for param in model.parameters():
+        param._is_hf_initialized = True
+    for buffer in model.buffers():
+        if buffer is not None:
+            buffer._is_hf_initialized = True
+
+    first_param = next(model.parameters())
+    original_weight = first_param.clone()
+
+    model._initialize_weights(model, is_custom_code=False)
+
+    assert torch.equal(first_param, original_weight)