Repository: huggingface/transformers
Pull request: Do not use global variable, and improve context manager coverage

# What does this PR do?

As per the title

Changed files:
filename: src/transformers/initialization.py
patch: @@ -206,3 +206,40 @@ def guard_torch_init_functions():
         for module, functions in originals.items():
             for func_name, func in functions.items():
                 setattr(module, func_name, func)
+
+
+@contextmanager
+def no_init_weights():
+    """
+    Disable weight initialization both at the torch-level, and at the transformers-level (`init_weights`).
+    This is used to speed-up initializing an empty model with deepspeed, as we do not initialize the model on meta device
+    with deepspeed, but we still don't need to run expensive weight initializations as we are loading params afterwards.
+    """
+    from .modeling_utils import PreTrainedModel
+
+    def empty_func(*args, **kwargs):
+        pass
+
+    originals = defaultdict(dict)
+    try:
+        # Replace all torch funcs by empty ones
+        for module_name in TORCH_MODULES_TO_PATCH:
+            if module_name in sys.modules:
+                module = sys.modules[module_name]
+                for func_name in TORCH_INIT_FUNCTIONS.keys():
+                    if hasattr(module, func_name):
+                        originals[module][func_name] = getattr(module, func_name)
+                        setattr(module, func_name, empty_func)
+
+        # Also patch our own `init_weights`
+        original_init_weights = PreTrainedModel.init_weights
+        PreTrainedModel.init_weights = empty_func
+
+        yield
+    finally:
+        # Set back the original torch functions on all modules
+        for module, functions in originals.items():
+            for func_name, func in functions.items():
+                setattr(module, func_name, func)
+        # Set back `init_weights`
+        PreTrainedModel.init_weights = original_init_weights

filename: src/transformers/modeling_utils.py
patch: @@ -151,7 +151,6 @@
 XLA_USE_BF16 = os.environ.get("XLA_USE_BF16", "0").upper()
 XLA_DOWNCAST_BF16 = os.environ.get("XLA_DOWNCAST_BF16", "0").upper()
 SpecificPreTrainedModelType = TypeVar("SpecificPreTrainedModelType", bound="PreTrainedModel")
-_init_weights = True
 _is_quantized = False
 _is_ds_init_called = False
 
@@ -170,51 +169,6 @@ def is_local_dist_rank_0():
     )
 
 
-TORCH_INIT_FUNCTIONS = {
-    "uniform_": nn.init.uniform_,
-    "normal_": nn.init.normal_,
-    "trunc_normal_": nn.init.trunc_normal_,
-    "constant_": nn.init.constant_,
-    "xavier_uniform_": nn.init.xavier_uniform_,
-    "xavier_normal_": nn.init.xavier_normal_,
-    "kaiming_uniform_": nn.init.kaiming_uniform_,
-    "kaiming_normal_": nn.init.kaiming_normal_,
-    "uniform": nn.init.uniform,
-    "normal": nn.init.normal,
-    "xavier_uniform": nn.init.xavier_uniform,
-    "xavier_normal": nn.init.xavier_normal,
-    "kaiming_uniform": nn.init.kaiming_uniform,
-    "kaiming_normal": nn.init.kaiming_normal,
-    "orthogonal_": nn.init.orthogonal_,
-}
-
-
-@contextmanager
-def no_init_weights():
-    """
-    Context manager to globally disable weight initialization to speed up loading large models.
-    """
-    global _init_weights
-    old_init_weights = _init_weights
-
-    _init_weights = False
-
-    def _skip_init(*args, **kwargs):
-        pass
-
-    # Save the original initialization functions
-    for name, init_func in TORCH_INIT_FUNCTIONS.items():
-        setattr(torch.nn.init, name, _skip_init)
-
-    try:
-        yield
-    finally:
-        _init_weights = old_init_weights
-        # Restore the original initialization functions
-        for name, init_func in TORCH_INIT_FUNCTIONS.items():
-            setattr(torch.nn.init, name, init_func)
-
-
 @contextmanager
 def set_quantized_state():
     global _is_quantized
@@ -2987,12 +2941,12 @@ def init_weights(self):
         Maybe initializes weights. If using a custom `PreTrainedModel`, you need to implement any
         initialization logic in `_init_weights`.
         """
-        if not _init_weights or get_torch_context_manager_or_global_device() == torch.device("meta"):
-            return
-        # Initialize weights
-        self.initialize_weights()
-        # Tie weights needs to be called here, but it can use the pre-computed `all_tied_weights_keys`
-        self.tie_weights(recompute_mapping=False)
+        # If we are initializing on meta device, there is no point in trying to run inits
+        if get_torch_context_manager_or_global_device() != torch.device("meta"):
+            # Initialize weights
+            self.initialize_weights()
+            # Tie weights needs to be called here, but it can use the pre-computed `all_tied_weights_keys`
+            self.tie_weights(recompute_mapping=False)
 
     def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
         """
@@ -3537,7 +3491,11 @@ def get_init_context(cls, dtype: torch.dtype, is_quantized: bool, _is_ds_init_ca
             if not is_quantized and not _is_ds_init_called:
                 logger.info("Detected DeepSpeed ZeRO-3: activating zero.init() for this model")
                 init_contexts.extend(
-                    [no_init_weights(), deepspeed.zero.Init(config_dict_or_path=deepspeed_config()), set_zero3_state()]
+                    [
+                        init.no_init_weights(),
+                        deepspeed.zero.Init(config_dict_or_path=deepspeed_config()),
+                        set_zero3_state(),
+                    ]
                 )
             elif is_quantized:
                 init_contexts.extend([torch.device("meta"), set_quantized_state()])