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>
This commit is contained in:
Daniel Han 2026-07-09 02:26:24 -07:00 committed by GitHub
commit 0d4bd50768
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 136 additions and 2 deletions

View file

@ -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

View file

@ -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