From 0d4bd50768ca1d55009b51dfa097d7c71e819f4b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 02:26:24 -0700 Subject: [PATCH] Restore process-global torch.compile config on torch 2.12 so gradient checkpointing backward honors it (#7019) * Mirror dynamo/inductor config sets into defaults so torch 2.12 worker threads honor them torch 2.12 stores config user overrides in ContextVars, so direct assignments like torch._dynamo.config.recompile_limit = 1024 no longer reach the autograd engine worker threads. Gradient checkpointing recomputes fullgraph-compiled gpt-oss kernels inside backward on those threads, which then read the default recompile limit of 8 and raise FailOnRecompileLimitHit at step 0 of GRPO/SFT. Mirror direct config assignments into the process-global entry defaults on torch >= 2.12, restoring the torch <= 2.11 cross-thread semantics while leaving the context-scoped config.patch API untouched. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep config.patch thread-local when mirroring dynamo/inductor sets config.patch(...) also assigns through ConfigModule.__setattr__, so the default-mirror was leaking its scoped, thread-local writes into the process-global entry default. Track patch enter/exit with a per-thread depth counter (wrapping ConfigModule.patch) and skip mirroring while inside a patch, so only genuine direct assignments restore the torch 2.11 cross-thread semantics and config.patch stays context-local. * Also keep config.load_config thread-local when mirroring config sets load_config restores a saved dynamo/inductor config by calling setattr per key, which the default-mirror would otherwise leak process-wide just like config.patch did. Wrap load_config with the same per-thread depth counter (renamed to _scoped_depth) so both scoped writers skip the mirror and stay context-local, while genuine direct assignments still restore the torch 2.11 cross-thread default. * Drop the pre-existing override replay from the config thread fix The replay was redundant: this runs from _gpu_init before unsloth sets any dynamo/inductor config, so the __setattr__ wrapper already mirrors every later assignment (recompile_limit included). It could also read a value that belonged to a config.patch context still active at import time and write that thread-local override into the global default. Removing it keeps the cross-thread fix and drops the now-unused _inductor.config import. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/_gpu_init.py | 6 ++ unsloth/import_fixes.py | 132 +++++++++++++++++++++++++++++++++++++++- 2 files changed, 136 insertions(+), 2 deletions(-) diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index e39b44488a..e6178e60f3 100644 --- a/unsloth/_gpu_init.py +++ b/unsloth/_gpu_init.py @@ -173,6 +173,7 @@ from .import_fixes import ( fix_vllm_guided_decoding_params, fix_vllm_pdl_blackwell, fix_triton_compiled_kernel_missing_attrs, + fix_dynamo_config_thread_visibility, patch_trunc_normal_precision_issue, ignore_logger_messages, patch_ipykernel_hf_xet, @@ -203,6 +204,10 @@ fix_vllm_guided_decoding_params() fix_trl_vllm_ascend() fix_vllm_pdl_blackwell() fix_triton_compiled_kernel_missing_attrs() +# Must run before unsloth_zoo's patch_torch_compile and the gpt-oss temporary +# patches raise the dynamo recompile limits, so those settings reach the +# autograd worker threads on torch >= 2.12. +fix_dynamo_config_thread_visibility() patch_trunc_normal_precision_issue() ignore_logger_messages() patch_ipykernel_hf_xet() @@ -233,6 +238,7 @@ del fix_vllm_guided_decoding_params del fix_trl_vllm_ascend del fix_vllm_pdl_blackwell del fix_triton_compiled_kernel_missing_attrs +del fix_dynamo_config_thread_visibility del patch_trunc_normal_precision_issue del ignore_logger_messages del patch_ipykernel_hf_xet diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index e5a5d01c2f..c5300ed4d0 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -1064,6 +1064,135 @@ def fix_triton_compiled_kernel_missing_attrs(): ) +def fix_dynamo_config_thread_visibility(): + """torch 2.12 made torch._dynamo/_inductor config overrides thread-local + (ContextVars), so `config.recompile_limit = 1024` set on the main thread is + invisible to the autograd worker threads that run backward. Gradient + checkpointing recompiles fullgraph gpt-oss kernels there against the default + limit of 8, raising FailOnRecompileLimitHit at step 0. Mirror direct config + assignments into the process-global entry default (torch <= 2.11 semantics). + config.patch(...) and config.load_config(...) also assign via __setattr__ but + are thread-local by design, so skip mirroring while inside one (tracked per + thread). No-op below torch 2.12 and on any torch without this internal layout. + """ + try: + import torch + + if Version(torch.__version__) < Version("2.12.0"): + return + import torch._dynamo.config as _dynamo_config + from torch.utils._config_module import ConfigModule + from contextvars import ContextVar + except Exception: + return + + try: + probe = getattr(_dynamo_config, "_config", {}).get("recompile_limit", None) + if probe is None or not isinstance(getattr(probe, "user_override", None), ContextVar): + # Overrides are not context-local on this torch; nothing to fix. + return + original_setattr = ConfigModule.__setattr__ + if getattr(original_setattr, "__unsloth_patched__", False): + return + except Exception: + return + + mirrored_modules = ("torch._dynamo.config", "torch._inductor.config") + + # config.patch(...) and config.load_config(...) also assign via __setattr__, but + # their writes are thread-local by design; a per-thread depth counter marks them + # so they are not mirrored into the process-global default. + import threading + + _scoped_depth = threading.local() + + def _in_scoped_write(): + return getattr(_scoped_depth, "n", 0) > 0 + + def _bump(delta): + _scoped_depth.n = getattr(_scoped_depth, "n", 0) + delta + + original_patch = ConfigModule.patch + if not getattr(original_patch, "__unsloth_patched__", False): + + @functools.wraps(original_patch) + def _patched_patch(self, *args, **kwargs): + ctx = original_patch(self, *args, **kwargs) + try: + cls = type(ctx) # patch() builds a fresh ConfigPatch class each call + if not getattr(cls, "__unsloth_patch_wrapped__", False): + _enter0, _exit0 = cls.__enter__, cls.__exit__ + + def _enter(s, _e = _enter0): + _bump(1) + try: + return _e(s) + finally: + _bump(-1) + + def _exit( + s, + *a, + _x = _exit0, + ): + _bump(1) + try: + return _x(s, *a) + finally: + _bump(-1) + + cls.__enter__, cls.__exit__ = _enter, _exit + cls.__unsloth_patch_wrapped__ = True + except Exception: + pass + return ctx + + _patched_patch.__unsloth_patched__ = True + ConfigModule.patch = _patched_patch + + # load_config restores a saved config by calling setattr per key (thread-local). + original_load_config = getattr(ConfigModule, "load_config", None) + if callable(original_load_config) and not getattr( + original_load_config, "__unsloth_patched__", False + ): + + @functools.wraps(original_load_config) + def _patched_load_config(self, *args, **kwargs): + _bump(1) + try: + return original_load_config(self, *args, **kwargs) + finally: + _bump(-1) + + _patched_load_config.__unsloth_patched__ = True + ConfigModule.load_config = _patched_load_config + + @functools.wraps(original_setattr) + def _patched_setattr(self, name, value): + original_setattr(self, name, value) + if _in_scoped_write(): + return # transient patch / load_config write: keep it thread-local + # Aliases (cache_size_limit -> recompile_limit) re-enter with the real name. + if self.__dict__.get("__name__", None) in mirrored_modules: + try: + entry = self.__dict__["_config"].get(name, None) + if entry is not None and entry.alias is None: + entry.default = value + except Exception: + pass + + _patched_setattr.__unsloth_patched__ = True + ConfigModule.__setattr__ = _patched_setattr + + # No replay of existing overrides: unsloth installs this before it sets any + # dynamo/inductor config, so the wrapper mirrors every later assignment. Replaying + # would also bake a still-active config.patch override into the global default. + logger.info( + "Unsloth: Patched torch config modules so dynamo/inductor settings " + "(e.g. recompile_limit) apply across threads on torch >= 2.12." + ) + + def patch_trunc_normal_precision_issue(): """ Patch torch.nn.init.trunc_normal_ for low precision tensors to run init in fp32. @@ -1323,8 +1452,7 @@ def fix_vllm_pdl_blackwell(): if patched: logger.info( - f"Unsloth: Applied PDL fix for SM100 ({sm100_gpu_name}) - " - f"patched: {', '.join(patched)}" + f"Unsloth: Applied PDL fix for SM100 ({sm100_gpu_name}) - patched: {', '.join(patched)}" ) else: # Just set the env var - vLLM might be an older version without supports_pdl