From f5e420a9e3ae08852aa55ec14d1e2ea7250361c4 Mon Sep 17 00:00:00 2001 From: vasqu Date: Wed, 11 Mar 2026 21:46:06 +0100 Subject: [PATCH 1/4] fix --- src/transformers/utils/output_capturing.py | 37 ++++++++++------------ 1 file changed, 17 insertions(+), 20 deletions(-) diff --git a/src/transformers/utils/output_capturing.py b/src/transformers/utils/output_capturing.py index 0aa4c0290e05..809140d711ea 100644 --- a/src/transformers/utils/output_capturing.py +++ b/src/transformers/utils/output_capturing.py @@ -64,35 +64,32 @@ class CompileableContextVar: """ def __init__(self, name, default): - self.context_var = ContextVar(name, default=default) - self.global_var = default - self.compiling = False + self.context_var = ContextVar(name, default=default) # real thread-safe, local var + self.global_var = default # mirror for compile, global var + self.default = default 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: Local attributes/flags may be interpreted once during tracing + # and effectively treated as constant across compiled frames. If we + # cached such a flag, eager frames could incorrectly fall back to the + # shared global value, breaking visibility between eager wrappers and + # compiled hooks + 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): + self.global_var = value + # Same, check the current state dynamically rather than relying on cached attributes 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 = self.default + if token is not None: self.context_var.reset(token) From e4b038555a627a17c322ee5478fa7c56d07b2408 Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Thu, 12 Mar 2026 10:09:57 +0100 Subject: [PATCH 2/4] add to all trigger --- utils/tests_fetcher.py | 1 + 1 file changed, 1 insertion(+) diff --git a/utils/tests_fetcher.py b/utils/tests_fetcher.py index 3ce20e968146..4644bfb1d97a 100644 --- a/utils/tests_fetcher.py +++ b/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", ) From fb23532575ce7e40b43139937757f5f3d74d8d12 Mon Sep 17 00:00:00 2001 From: vasqu Date: Thu, 12 Mar 2026 13:18:52 +0100 Subject: [PATCH 3/4] fixup comments --- src/transformers/utils/output_capturing.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/transformers/utils/output_capturing.py b/src/transformers/utils/output_capturing.py index 809140d711ea..0237fc61e77e 100644 --- a/src/transformers/utils/output_capturing.py +++ b/src/transformers/utils/output_capturing.py @@ -71,18 +71,18 @@ def __init__(self, name, default): def get(self): # Always check the current dynamo state dynamically! See the note below # - # NOTE: Local attributes/flags may be interpreted once during tracing - # and effectively treated as constant across compiled frames. If we - # cached such a flag, eager frames could incorrectly fall back to the - # shared global value, breaking visibility between eager wrappers and - # compiled hooks + # 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 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 - # Same, check the current state dynamically rather than relying on cached attributes if is_torchdynamo_compiling(): return None return self.context_var.set(value) From 04569d0cbcbb14960dd78a3dcae80830856d44d8 Mon Sep 17 00:00:00 2001 From: vasqu Date: Thu, 12 Mar 2026 13:25:59 +0100 Subject: [PATCH 4/4] remove default exposure and always use None default --- src/transformers/utils/output_capturing.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/src/transformers/utils/output_capturing.py b/src/transformers/utils/output_capturing.py index 0237fc61e77e..f7b23e657850 100644 --- a/src/transformers/utils/output_capturing.py +++ b/src/transformers/utils/output_capturing.py @@ -63,10 +63,9 @@ 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) # real thread-safe, local var - self.global_var = default # mirror for compile, global var - self.default = default + 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): # Always check the current dynamo state dynamically! See the note below @@ -88,13 +87,13 @@ def set(self, value): return self.context_var.set(value) def reset(self, token): - self.global_var = self.default + 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: