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:
parent
6d674e5cc9
commit
0d4bd50768
2 changed files with 136 additions and 2 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue