diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 3e2b5e545687..0805d5b98449 100644 --- a/src/transformers/modeling_utils.py +++ b/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 diff --git a/tests/models/qwen3_omni_moe/test_modeling_qwen3_omni_moe.py b/tests/models/qwen3_omni_moe/test_modeling_qwen3_omni_moe.py index 5bd3d7cf5d44..1196f0d40ee3 100644 --- a/tests/models/qwen3_omni_moe/test_modeling_qwen3_omni_moe.py +++ b/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): diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index 2d6f9b41133d..17dc63c09592 100644 --- a/tests/test_modeling_common.py +++ b/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)