Harden Trainer._load_rng_state against malicious checkpoints (CVE-2026-1839) (#6351)

* Harden Trainer._load_rng_state against malicious checkpoints (CVE-2026-1839)

Resuming from an untrusted checkpoint loads its rng_state.pth via transformers' Trainer._load_rng_state, which calls torch.load without weights_only=True. On torch < 2.6 (weights_only not yet the default, safe_globals() does not restrict the unpickler) a crafted rng_state.pth runs arbitrary code. transformers fixed this in 5.0.0rc3, but we pin 4.57.x for compatibility, so add an import-time monkey patch instead.

patch_unsafe_trainer_rng_load() wraps the method and forces torch.load with weights_only=True for the call, mirroring the upstream fix while keeping the surrounding distributed/device logic intact and tracking transformers' own safe_globals() allowlist. It is idempotent and a no-op on torch >= 2.6 (already the default) and transformers >= 5.0.0rc3 (already fixed), so it self-disables once the pin is bumped. Wired into _gpu_init.py so it applies on import unsloth.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Gate Trainer._load_rng_state via check_torch_load_is_safe instead of forcing weights_only

Review feedback was correct: forcing weights_only=True is not the right fix. It is not a safe boundary below torch 2.6 (CVE-2025-32434), where transformers' safe_globals() is also a nullcontext, so a legitimate rng_state.pth (which holds numpy state) would fail to load there. Call transformers' own check_torch_load_is_safe() before the load instead: it raises on torch < 2.6 and is a no-op on torch >= 2.6, where torch.load already defaults to weights_only=True. This also drops the temporary global torch.load swap, removing the thread-safety race.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Trim comments in the RNG-load guard

* Harden only the rng torch.load via a thread-local gate

Address review feedback on the rng-load guard:
- Import Trainer on its own and fall back to a local torch-version check when
  transformers does not export check_torch_load_is_safe, so older supported
  4.51.x installs are still protected.
- Force weights_only=True at the rng load so TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD
  cannot re-enable pickle execution.
- Gate at the actual rng torch.load through a thread-local flag rather than up
  front, so model-only or rng-less resumes still proceed on torch < 2.6 and
  other torch.load callers are unaffected.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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()

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.