diff --git a/src/transformers/utils/output_capturing.py b/src/transformers/utils/output_capturing.py index 0aa4c0290e05..f7b23e657850 100644 --- a/src/transformers/utils/output_capturing.py +++ b/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: 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", )