Repository: huggingface/transformers
Pull request: fix: initialize BatchNorm2d buffers only on meta

# What does this PR do?

Commit 8dd9c999a6262d6ceb48f4a2da7acaccfa80e3bc introduced a regression by unconditionally reinitializing BatchNorm2d buffers (running_mean, running_var, num_batches_tracked) in the _init_weights() method.

The problem: When loading pretrained timm models, the flow is:
1. timm.create_model(pretrained=True) correctly loads pretrained BatchNorm statistics
2. post_init() calls _init_weights() on all modules
3. _init_weights() overwrites the pretrained BatchNorm buffers with zeros/ones
4. Model produces incorrect results because normalization uses wrong statistics

The error shows a mismatch at index (0, 178, 0, 0) with a difference of 80224.0 on CircleCI on AVX 512 CPUs

The fix: added a check to only initialize BatchNorm2d buffers when they're on the meta device (uninitialized): This ensures pretrained buffers are preserved while still handling the meta device initialization case.


Changed files:
filename: src/transformers/models/timm_backbone/modeling_timm_backbone.py
patch: @@ -117,10 +117,17 @@ def _init_weights(self, module):
         assume weights and persistent buffers will be part of checkpoint as we have no way to control timm inits)"""
         if hasattr(module, "init_non_persistent_buffers"):
             module.init_non_persistent_buffers()
-        elif isinstance(module, nn.BatchNorm2d) and getattr(module, "running_mean", None) is not None:
-            init.zeros_(module.running_mean)
-            init.ones_(module.running_var)
-            init.zeros_(module.num_batches_tracked)
+        elif isinstance(module, nn.BatchNorm2d):
+            # Skip initialization if using pretrained backbone - buffers are already loaded from checkpoint
+            if self.config.use_pretrained_backbone:
+                return
+
+            # For non-pretrained models, always initialize buffers (handles both meta device and to_empty() cases)
+            running_mean = getattr(module, "running_mean", None)
+            if running_mean is not None:
+                init.zeros_(module.running_mean)
+                init.ones_(module.running_var)
+                init.zeros_(module.num_batches_tracked)
 
     def forward(
         self,

filename: src/transformers/models/timm_wrapper/modeling_timm_wrapper.py
patch: @@ -115,10 +115,14 @@ def _init_weights(self, module):
         # Also, reinit all non-persistemt buffers if any!
         if hasattr(module, "init_non_persistent_buffers"):
             module.init_non_persistent_buffers()
-        elif isinstance(module, nn.BatchNorm2d) and getattr(module, "running_mean", None) is not None:
-            init.zeros_(module.running_mean)
-            init.ones_(module.running_var)
-            init.zeros_(module.num_batches_tracked)
+        elif isinstance(module, nn.BatchNorm2d):
+            # TimmWrapper always creates models with pretrained=False, so buffers are never pre-loaded
+            # Always initialize buffers (handles both meta device and to_empty() cases)
+            running_mean = getattr(module, "running_mean", None)
+            if running_mean is not None:
+                init.zeros_(module.running_mean)
+                init.ones_(module.running_var)
+                init.zeros_(module.num_batches_tracked)
 
     def _timm_model_supports_gradient_checkpointing(self):
         """