Fix #2661 ... don't skip reset_parameters/init when meta device detected as it breaks use of accelerate and similar dispatch override context managers

Author: rwightman

timm/layers/blur_pool.py

@@ -51,8 +51,8 @@ def __init__(
         filt_shape = (channels or 1, 1, filt_size, filt_size)
         self.register_buffer('filt', torch.empty(filt_shape, device=device, dtype=dtype), persistent=False)
 
-        if not self.filt.is_meta:
-            self.reset_parameters()
+        # TODO: skip init when on meta device when safe to do so
+        self.reset_parameters()
 
     def reset_parameters(self) -> None:
         """Initialize buffers."""

timm/layers/lambda_layer.py

@@ -125,8 +125,8 @@ def __init__(
 
         self.pool = nn.AvgPool2d(2, 2) if stride == 2 else nn.Identity()
 
-        if not self.qkv.weight.is_meta:
-            self.reset_parameters()
+        # TODO: skip init when on meta device when safe to do so
+        self.reset_parameters()
 
     def reset_parameters(self) -> None:
         """Initialize parameters and buffers."""

timm/layers/pos_embed_rel.py

@@ -299,8 +299,8 @@ def __init__(
             persistent=False,
         )
 
-        if not self.relative_position_bias_table.is_meta:
-            self.reset_parameters()
+        # TODO: skip init when on meta device when safe to do so
+        self.reset_parameters()
 
     def reset_parameters(self) -> None:
         """Initialize parameters and buffers."""
@@ -420,8 +420,8 @@ def __init__(
             persistent=False,
         )
 
-        if not self.mlp.fc1.weight.is_meta:
-            self.reset_parameters()
+        # TODO: skip init when on meta device when safe to do so
+        self.reset_parameters()
 
     def get_bias(self) -> torch.Tensor:
         relative_position_bias = self.mlp(self.rel_coords_log)
@@ -554,8 +554,8 @@ def __init__(
         self.register_buffer('height_lookup', torch.empty(height_lookup_shape, **dd), persistent=False)
         self.register_buffer('width_lookup', torch.empty(width_lookup_shape, **dd), persistent=False)
 
-        if not self.relative_position_bias_table.is_meta:
-            self.reset_parameters()
+        # TODO: skip init when on meta device when safe to do so
+        self.reset_parameters()
 
     def reset_parameters(self) -> None:
         """Initialize parameters and buffers."""

timm/layers/pos_embed_sincos.py

@@ -11,8 +11,6 @@
 from ._fx import register_notrace_function
 from .grid import ndgrid
 from .trace_utils import _assert
-from .weight_init import is_meta_device
-
 
 def pixel_freq_bands(
         num_bands: int,
@@ -188,8 +186,8 @@ def __init__(
         self.keep_spatial = keep_spatial
         self.register_buffer('bands', torch.empty(num_bands, device=device, dtype=dtype), persistent=False)
 
-        if not is_meta_device(device):
-            self.reset_parameters()
+        # TODO: skip init when on meta device when safe to do so
+        self.reset_parameters()
 
     def reset_parameters(self) -> None:
         """Initialize parameters and buffers."""
@@ -447,8 +445,8 @@ def __init__(
             self.register_buffer('pos_embed_sin', torch.empty(emb_shape, device=device, dtype=dtype), persistent=False)
             self.register_buffer('pos_embed_cos', torch.empty(emb_shape, device=device, dtype=dtype), persistent=False)
 
-        if not is_meta_device(device):
-            self.reset_parameters()
+        # TODO: skip init when on meta device when safe to do so
+        self.reset_parameters()
 
     def reset_parameters(self) -> None:
         """Initialize parameters and buffers."""
@@ -583,8 +581,8 @@ def __init__(
             emb_shape = (num_pos, dim * 2)  # concatenated sin & cos
             self.register_buffer('pos_embed', torch.empty(emb_shape, device=device, dtype=dtype), persistent=False)
 
-        if not is_meta_device(device):
-            self.reset_parameters()
+        # TODO: skip init when on meta device when safe to do so
+        self.reset_parameters()
 
     def reset_parameters(self) -> None:
         """Initialize parameters and buffers."""
@@ -851,8 +849,8 @@ def __init__(
                 num_pos *= s
             self.register_buffer('t_x', torch.empty(num_pos, device=device, dtype=dtype), persistent=False)
             self.register_buffer('t_y', torch.empty(num_pos, device=device, dtype=dtype), persistent=False)
-            if not is_meta_device(device):
-                self._init_buffers()
+            # TODO: skip init when on meta device when safe to do so
+            self._init_buffers()
         else:
             self.t_x = self.t_y = None
 
@@ -1087,8 +1085,8 @@ def __init__(
         else:
             self.pos_embed_cached = None
 
-        if not is_meta_device(device):
-            self.reset_parameters()
+        # TODO: skip init when on meta device when safe to do so
+        self.reset_parameters()
 
     def reset_parameters(self) -> None:
         """Initialize parameters and buffers."""

timm/models/beit.py

@@ -179,8 +179,8 @@ def __init__(
         self.proj = nn.Linear(all_head_dim, dim, **dd)
         self.proj_drop = nn.Dropout(proj_drop)
 
-        if not self.proj.weight.is_meta:
-            self.reset_parameters()
+        # TODO: skip init when on meta device when safe to do so
+        self.reset_parameters()
 
     def _get_rel_pos_bias(self) -> torch.Tensor:
         """Get relative position bias for the attention window.
@@ -362,8 +362,8 @@ def __init__(
         else:
             self.gamma_1, self.gamma_2 = None, None
 
-        if not self.norm1.weight.is_meta:
-            self.reset_parameters()
+        # TODO: skip init when on meta device when safe to do so
+        self.reset_parameters()
 
     def reset_parameters(self) -> None:
         """Initialize parameters."""
@@ -416,8 +416,8 @@ def __init__(self, window_size: Tuple[int, int], num_heads: int, device=None, dt
             persistent=False,
         )
 
-        if not self.relative_position_bias_table.is_meta:
-            self.reset_parameters()
+        # TODO: skip init when on meta device when safe to do so
+        self.reset_parameters()
 
     def reset_parameters(self) -> None:
         """Initialize parameters and buffers."""
@@ -569,8 +569,8 @@ def __init__(
         self.head = nn.Linear(embed_dim, num_classes, **dd) if num_classes > 0 else nn.Identity()
         self.head_init_scale = head_init_scale
 
-        if not self.patch_embed.proj.weight.is_meta:
-            self.init_weights(needs_reset=False)
+        # TODO: skip init when on meta device when safe to do so
+        self.init_weights(needs_reset=False)
 
     def init_weights(self, needs_reset: bool = True) -> None:
         """Initialize model weights.

timm/models/csatv2.py

@@ -226,8 +226,8 @@ def __init__(
         self.register_buffer('imagenet_mean', torch.empty(3, 1, 1, device=device, dtype=dtype), persistent=False)
         self.register_buffer('imagenet_std', torch.empty(3, 1, 1, device=device, dtype=dtype), persistent=False)
 
-        if not self.mean.is_meta:
-            self.reset_parameters()
+        # TODO: skip init when on meta device when safe to do so
+        self.reset_parameters()
 
     def reset_parameters(self) -> None:
         """Initialize buffers."""
@@ -609,8 +609,8 @@ def __init__(
 
         self.head = NormMlpClassifierHead(dims[-1], num_classes, pool_type=global_pool, **dd)
 
-        if not self.stem_dct.conv_y.weight.is_meta:
-            self.init_weights(needs_reset=False)
+        # TODO: skip init when on meta device when safe to do so
+        self.init_weights(needs_reset=False)
 
     def init_weights(self, needs_reset: bool = True):
         self.apply(partial(self._init_weights, needs_reset=needs_reset))

timm/models/deit.py

@@ -46,7 +46,8 @@ def __init__(self, *args, **kwargs):
         self.distilled_training = False  # must set this True to train w/ distillation token
 
         self.weight_init_mode = 'reset' if weight_init == 'skip' else weight_init
-        if weight_init != 'skip' and not next(self.parameters()).is_meta:
+        # TODO: skip init when on meta device when safe to do so
+        if weight_init != 'skip':
             self.init_weights(needs_reset=False)
 
     def init_weights(self, mode='', needs_reset=True):

timm/models/efficientformer_v2.py

@@ -159,8 +159,8 @@ def __init__(
         )
         self.attention_bias_cache = {}
 
-        if not self.attention_biases.is_meta:
-            self.reset_parameters()
+        # TODO: skip init when on meta device when safe to do so
+        self.reset_parameters()
 
     @torch.no_grad()
     def train(self, mode=True):
@@ -300,8 +300,8 @@ def __init__(
         )
         self.attention_bias_cache = {}
 
-        if not self.attention_biases.is_meta:
-            self.reset_parameters()
+        # TODO: skip init when on meta device when safe to do so
+        self.reset_parameters()
 
     @torch.no_grad()
     def train(self, mode=True):
@@ -719,8 +719,8 @@ def __init__(
         else:
             self.head_dist = None
 
-        if not self.norm.weight.is_meta:
-            self.init_weights(needs_reset=False)
+        # TODO: skip init when on meta device when safe to do so
+        self.init_weights(needs_reset=False)
 
         self.distilled_training = False
 

timm/models/efficientvit_msra.py

@@ -207,8 +207,8 @@ def __init__(
         )
         self.attention_bias_cache = {}
 
-        if not self.attention_bias_idxs.is_meta:
-            self.reset_parameters()
+        # TODO: skip init when on meta device when safe to do so
+        self.reset_parameters()
 
     def reset_parameters(self) -> None:
         """Initialize parameters and buffers."""
@@ -537,8 +537,8 @@ def __init__(
         self.head = NormLinear(
             self.num_features, num_classes, drop=self.drop_rate, **dd) if num_classes > 0 else torch.nn.Identity()
 
-        if not self.patch_embed.conv1.conv.weight.is_meta:
-            self.init_weights(needs_reset=False)
+        # TODO: skip init when on meta device when safe to do so
+        self.init_weights(needs_reset=False)
 
     def init_weights(self, needs_reset: bool = True):
         self.apply(partial(self._init_weights, needs_reset=needs_reset))

timm/models/eva.py

@@ -178,8 +178,8 @@ def __init__(
         self.proj = nn.Linear(attn_dim, dim, **dd)
         self.proj_drop = nn.Dropout(proj_drop)
 
-        if not self.proj.weight.is_meta:
-            self.reset_parameters()
+        # TODO: skip init when on meta device when safe to do so
+        self.reset_parameters()
 
     def reset_parameters(self) -> None:
         """Initialize parameters and buffers."""
@@ -369,8 +369,8 @@ def __init__(
         self.gamma_2 = nn.Parameter(torch.empty(dim, **dd)) if init_values is not None else None
         self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
 
-        if not self.norm1.weight.is_meta:
-            self.reset_parameters()
+        # TODO: skip init when on meta device when safe to do so
+        self.reset_parameters()
 
     def reset_parameters(self) -> None:
         """Initialize parameters."""
@@ -738,8 +738,8 @@ def __init__(
         self.head = nn.Linear(embed_dim, num_classes, **dd) if num_classes > 0 else nn.Identity()
         self.head_init_scale = head_init_scale
 
-        if not self.patch_embed.proj.weight.is_meta:
-            self.init_weights(needs_reset=False)
+        # TODO: skip init when on meta device when safe to do so
+        self.init_weights(needs_reset=False)
 
     def init_weights(self, needs_reset: bool = True):
         self.apply(partial(self._init_weights, needs_reset=needs_reset))