diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index 86f675c281..7998d62191 100644 --- a/unsloth/_gpu_init.py +++ b/unsloth/_gpu_init.py @@ -179,6 +179,7 @@ from .import_fixes import ( patch_trackio, patch_datasets, patch_enable_input_require_grads, + patch_unsafe_trainer_rng_load, fix_openenv_no_vllm, patch_openspiel_env_async, fix_executorch, @@ -206,6 +207,7 @@ patch_ipykernel_hf_xet() patch_trackio() patch_datasets() patch_enable_input_require_grads() +patch_unsafe_trainer_rng_load() fix_openenv_no_vllm() patch_openspiel_env_async() fix_executorch() diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index eb43666d56..87003b5985 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -628,6 +628,75 @@ def patch_enable_input_require_grads(): logger.info("Unsloth: Patched enable_input_require_grads for vision model compatibility") +def patch_unsafe_trainer_rng_load(): + """Harden Trainer._load_rng_state against CVE-2026-1839 (RCE from a malicious + rng_state.pth on resume). Hardens only the rng torch.load, via a thread-local + flag, so it forces weights_only=True (defeats TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD) + and refuses torch < 2.6 (CVE-2025-32434), while rng-less resumes and unrelated + torch.load calls are untouched. No-op if transformers is absent or already + guards the load (>= 5.0.0rc3).""" + if importlib.util.find_spec("transformers") is None: + return + try: + from transformers.trainer import Trainer + except Exception: + return + load_rng_state = getattr(Trainer, "_load_rng_state", None) + if load_rng_state is None or getattr(load_rng_state, "_unsloth_safe_rng_load", False): + return + try: + source = inspect.getsource(load_rng_state) + except Exception: + return + if "torch.load" not in source or "check_torch_load_is_safe" in source: + return + + import threading, torch + + try: + # Older supported transformers (>= 4.51.3) may not export the helper. + from transformers.utils.import_utils import check_torch_load_is_safe + except Exception: + + def check_torch_load_is_safe(): + if TrueVersion(torch.__version__.split("+")[0]) < TrueVersion("2.6"): + raise RuntimeError( + "Unsloth: refusing to load checkpoint RNG state on torch < 2.6 " + "(CVE-2026-1839 / CVE-2025-32434); upgrade to torch >= 2.6." + ) + + # Install one process-wide torch.load shim that stays inert unless the calling + # thread is inside _load_rng_state, so we gate only at the real rng load with + # no global-swap race and no effect on other torch.load callers. + if not getattr(torch.load, "_unsloth_rng_guard", False): + _orig_load = torch.load + _rng_active = threading.local() + + @functools.wraps(_orig_load) + def _guarded_torch_load(*args, **kwargs): + if getattr(_rng_active, "on", False): + check_torch_load_is_safe() # raises on torch < 2.6 (CVE-2025-32434) + kwargs.setdefault("weights_only", True) + return _orig_load(*args, **kwargs) + + _guarded_torch_load._unsloth_rng_guard = True + _guarded_torch_load._unsloth_rng_flag = _rng_active + torch.load = _guarded_torch_load + _rng_active = torch.load._unsloth_rng_flag + + @functools.wraps(load_rng_state) + def _unsloth_safe_load_rng_state(self, checkpoint): + _rng_active.on = True + try: + return load_rng_state(self, checkpoint) + finally: + _rng_active.on = False + + _unsloth_safe_load_rng_state._unsloth_safe_rng_load = True + Trainer._load_rng_state = _unsloth_safe_load_rng_state + logger.info("Unsloth: Hardened Trainer._load_rng_state rng loading (CVE-2026-1839).") + + def _is_custom_torch_build(raw_version_str): """Check if a raw version string indicates a custom or source build.