[Compile] Fix capture outputs during compile

src/transformers/utils/output_capturing.py

@@ -63,41 +63,37 @@ 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
-        self.compiling = False
+    def __init__(self, name):
+        self.context_var = ContextVar(name, default=None)  # real thread-safe, local var
+        self.global_var = None  # mirror for compile, global var
 
     def get(self):
-        # Set was called before and compilation was already detected
-        if self.compiling:
+        # Always check the current dynamo state dynamically! See the note below
+        #
+        # NOTE: We can move between compiled and eager regions (e.g. after
+        #       graph breaks). A cached flag would force the choice on a past
+        #       state, so later eager accesses could incorrectly use the shared global
+        #       value (`global_var`) instead of the real local value (`context_var`)
+        if is_torchdynamo_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):
+        # Always set `global_var` to ensure consistent states between eager and compiled regions
+        # See the note above --> sometimes switches between compile/eager can happen
+        self.global_var = value
         if is_torchdynamo_compiling():
-            self.global_var = value
-            self.compiling = True
             return None
-        else:
-            return self.context_var.set(value)
+        return self.context_var.set(value)
 
     def reset(self, token):
-        if self.compiling or token is None:
-            self.global_var = None
-            self.compiling = False
-        else:
+        self.global_var = None
+        if token is not None:
             self.context_var.reset(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:

utils/tests_fetcher.py

@@ -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",
 )