Repository: huggingface/transformers
Pull request: Fix output capturing for Backbones

# What does this PR do?

As per the title. Supersedes https://github.com/huggingface/transformers/pull/44614.
This one is the result of a long debugging session and discussion with @vasqu.
The issue is as follow:
- Backbone ALWAYS need to capture their outputs
- In `test_sdpa_can_compile_dynamic`, we EXTERNALLY compile the model, i.e. we do `model = torch.compile(model, dynamic=True)` before calling any `forward`
- Since Backbone ALWAYS need to capture outputs, lazy hook installation will happen while the model is already compiling
- Due to the `lock`, hook installation creates a graph break, but that graph break extends a little bit to the later call to `set`
- `set` is called as-if we were not compiling, so uses the ContextVar
- Later `get` calls are still under compile, so they don't get the ContextVar, but the Global var.
- Containers are mixed up, so outputs are not gathered correctly
- Model later crashes as all outputs are required, but due to the mixup it only gets a part, or nothing at all

## Remedial

We can easily solve by forcing hook installation during `post_init` for Backbones. Since they always need to capture anyway, any `forward` with them would install them, so we do not incur any additional bloat on the model state that would not be there otherwise by installing early.
Then, we won't have the problematic earlier graph break when calling `set`, as the hooks are already installed. It also allows external compilation with `model = torch.compile(model, fullgrah=True)` for Backbone, which would otherwise not be possible as lazy installation requires a break in the graph as mentionned multiple times!


Changed files:
filename: src/transformers/backbone_utils.py
patch: @@ -20,6 +20,7 @@
 from huggingface_hub import repo_exists
 
 from .utils import logging
+from .utils.output_capturing import maybe_install_capturing_hooks
 
 
 logger = logging.get_logger(__name__)
@@ -181,6 +182,18 @@ def __init__(self, *args, **kwargs) -> None:
         else:
             raise ValueError(f"backbone_type {self.backbone_type} not supported.")
 
+    def post_init(self):
+        """
+        Override `post_init` to always install capturing hooks, as backbone will ALWAYS capture outputs. We need to do
+        it in `post_init`, as modules need to be already instantiated.
+        It avoids some mixups with `torch.compile`, as the first hook installation will need/create a graph break,
+        which can clash with external user call such as `model = torch.compile(model...)`.
+        """
+        # NOTE: Since this class is ALWAYS used as a Mixin with another PreTrainedModel class, this `super` call
+        # will call the PreTrained's `post_init`
+        super().post_init()
+        maybe_install_capturing_hooks(self)
+
     def _init_timm_backbone(self, backbone) -> None:
         """
         Initialize the backbone model from timm. The backbone must already be loaded to backbone

filename: src/transformers/utils/output_capturing.py
patch: @@ -63,22 +63,17 @@ class CompileableContextVar:
     that the access to the underlying variable is not thread-safe when compilation is triggered.
     """
 
-    def __init__(self, name, default):
-        self.context_var = ContextVar(name, default=default)
-        self.global_var = default
+    def __init__(self, name):
+        self.context_var = ContextVar(name, default=None)
+        self.global_var = None
         self.compiling = False
 
     def get(self):
         # Set was called before and compilation was already detected
         if self.compiling:
             return self.global_var
         else:
-            # Set was maybe never called, so still check it here
-            if is_torchdynamo_compiling():
-                self.is_compiling = True
-                return self.global_var
-            else:
-                return self.context_var.get()
+            return self.context_var.get()
 
     def set(self, value):
         if is_torchdynamo_compiling():
@@ -97,7 +92,7 @@ def reset(self, token):
 
 
 # Thread/context-safe global variable
-_active_collector = CompileableContextVar("output_collector", default=None)
+_active_collector = CompileableContextVar("output_collector")
 
 
 def install_output_capuring_hook(module: nn.Module, key: str, index: int) -> None:

filename: utils/tests_fetcher.py
patch: @@ -80,6 +80,7 @@
     "src/transformers/core_model_loading.py",
     "src/transformers/cache_utils.py",
     "src/transformers/generation/utils.py",
+    "src/transformers/utils/output_capturing.py",
 )