Fix export-time trust_remote_code bypass in FP8/INT8/GGUF-LoRA export (#6869)

* Fix export-time trust_remote_code bypass in FP8/INT8/GGUF-LoRA export

The torchao, compressed-tensors, and LoRA GGUF export paths re-read the merged
checkpoint and used to set trust_remote_code from the checkpoint config's static
auto_map (the torchao path also scanned the staged tokenizer/processor configs).
A model that loads with built-in Transformers classes can carry an auto_map entry,
which skips the load-time remote-code consent scan (that only runs when the load
already requested trust_remote_code) yet flips trust_remote_code on at export,
running unvetted custom code.

Derive the reload trust_remote_code from the approved load decision instead: a new
_loaded_via_remote_code() checks whether the in-memory model / tokenizer was itself
loaded from custom code (its class lives in the transformers_modules package),
walking PEFT / wrapper layers. Built-in-loaded models no longer gain trust from
config metadata; genuine custom-code models (loaded with consent) still reload
correctly. Add CPU-only regression tests.

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

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

* Harden _loaded_via_remote_code against a None/missing __module__

Read type(node).__module__ via getattr and require a string before startswith,
so a dynamically created or C-extension class with a None module does not raise
during export. Add a regression test.

* Split model and tokenizer trust for the compressed subprocess, walk processor components

The compressed-tensors export collapsed model and tokenizer trust into
one --trust-remote-code flag, so an approved custom tokenizer would have
let an unapproved model's custom code run inside the quantization
subprocess. The subprocess now takes --trust-remote-code-tokenizer for
the processor load and keeps --trust-remote-code for the model loads,
matching the torchao path's separate model_trust / tok_trust.

_loaded_via_remote_code now also walks processor components (tokenizer,
image_processor, feature_extractor, video_processor), so an approved
custom tokenizer held inside a built-in ProcessorMixin keeps its trust
on the export reload instead of failing with trust_remote_code=False.
The walk is a bounded BFS with a seen set so wrapper cycles terminate.

* [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>

tests/saving/test_torchao_remote_code_consent.py

@@ -0,0 +1,150 @@
+"""Regression tests for the export-time remote-code trust decision.
+
+FP8/FP4/INT quantization export re-reads the just-merged checkpoint. It used to enable
+trust_remote_code whenever the checkpoint's config carried an ``auto_map`` entry, so a model
+that loads fine with built-in classes (and therefore skips the load-time consent scan) could
+smuggle unvetted remote code that then runs at export. The export paths now derive
+trust_remote_code from ``_loaded_via_remote_code`` - the already approved load decision - instead.
+
+These run on CPU with no torch / unsloth import: they AST-extract the real helper from
+unsloth/save.py and exec it in isolation, plus assert the call sites dropped the auto_map trust.
+"""
+
+import ast
+from pathlib import Path
+
+_SAVE_PY = Path(__file__).resolve().parents[2] / "unsloth" / "save.py"
+_SRC = _SAVE_PY.read_text(encoding = "utf-8")
+
+
+def _load_helper():
+    """Exec just `_loaded_via_remote_code` from save.py (no torch import) and return it."""
+    tree = ast.parse(_SRC)
+    fn = next(
+        n
+        for n in tree.body
+        if isinstance(n, ast.FunctionDef) and n.name == "_loaded_via_remote_code"
+    )
+    ns = {}
+    exec(compile(ast.Module(body = [fn], type_ignores = []), str(_SAVE_PY), "exec"), ns)
+    return ns["_loaded_via_remote_code"]
+
+
+_loaded_via_remote_code = _load_helper()
+
+
+def _obj(module_name, **attrs):
+    """A throwaway instance whose class __module__ is `module_name`, plus given attributes."""
+    cls = type("Fake", (), {})
+    cls.__module__ = module_name
+    inst = cls()
+    for k, v in attrs.items():
+        setattr(inst, k, v)
+    return inst
+
+
+def test_builtin_class_is_not_remote_code():
+    assert _loaded_via_remote_code(_obj("transformers.models.llama.modeling_llama")) is False
+
+
+def test_transformers_modules_class_is_remote_code():
+    assert _loaded_via_remote_code(_obj("transformers_modules.acme.modeling_x")) is True
+
+
+def test_none_is_not_remote_code():
+    assert _loaded_via_remote_code(None) is False
+
+
+def test_none_module_is_not_remote_code():
+    # A class whose __module__ is None must not raise AttributeError.
+    assert _loaded_via_remote_code(_obj(None)) is False
+
+
+def test_auto_map_in_config_alone_does_not_grant_trust():
+    # The core bypass: a built-in-loadable model whose config merely declares auto_map must NOT
+    # be treated as remote-code-loaded (that is exactly what enabled the consent-gate bypass).
+    cfg = type("Cfg", (), {"auto_map": {"AutoModelForCausalLM": "modeling_x.Model"}})()
+    assert (
+        _loaded_via_remote_code(_obj("transformers.models.llama.modeling_llama", config = cfg))
+        is False
+    )
+
+
+def test_peft_base_model_is_unwrapped():
+    base = _obj("transformers_modules.acme.modeling_x")
+    peft = _obj("peft.peft_model", get_base_model = lambda: base)
+    assert _loaded_via_remote_code(peft) is True
+
+
+def test_wrapper_model_attr_is_walked():
+    inner = _obj("transformers_modules.acme.modeling_x")
+    wrapper = _obj("peft.peft_model", model = inner)
+    assert _loaded_via_remote_code(wrapper) is True
+
+
+def test_wrapper_over_builtin_stays_false():
+    inner = _obj("transformers.models.llama.modeling_llama")
+    wrapper = _obj("peft.peft_model", model = inner)
+    assert _loaded_via_remote_code(wrapper) is False
+
+
+def test_processor_held_custom_tokenizer_is_detected():
+    # A built-in ProcessorMixin can hold an approved custom-code tokenizer; the walk must
+    # descend into processor components or the export reload loses that approved trust.
+    tok = _obj("transformers_modules.acme.tokenization_x")
+    proc = _obj("transformers.processing_utils", tokenizer = tok)
+    assert _loaded_via_remote_code(proc) is True
+
+
+def test_processor_held_custom_image_processor_is_detected():
+    ip = _obj("transformers_modules.acme.image_processing_x")
+    proc = _obj("transformers.processing_utils", image_processor = ip)
+    assert _loaded_via_remote_code(proc) is True
+
+
+def test_builtin_processor_with_builtin_components_stays_false():
+    proc = _obj(
+        "transformers.processing_utils",
+        tokenizer = _obj("transformers.tokenization_utils_fast"),
+        image_processor = _obj("transformers.image_processing_utils"),
+    )
+    assert _loaded_via_remote_code(proc) is False
+
+
+def test_cyclic_wrappers_terminate():
+    a = _obj("peft.peft_model")
+    b = _obj("peft.peft_model", model = a)
+    a.model = b
+    assert _loaded_via_remote_code(a) is False
+
+
+# -- call-site assertions: the auto_map-derived trust is gone from every export path -----------
+
+
+def test_torchao_export_derives_trust_from_load_decision():
+    assert "model_trust = _loaded_via_remote_code(model)" in _SRC
+    assert "tok_trust = _loaded_via_remote_code(tokenizer)" in _SRC
+    assert "trust_remote_code = model_trust" in _SRC
+    assert "trust_remote_code = tok_trust" in _SRC
+    # The staged-config auto_map scan that granted trust is removed.
+    assert 'if "auto_map" in json.load' not in _SRC
+
+
+def test_compressed_and_gguf_lora_paths_drop_auto_map_trust():
+    # No path derives a trust decision straight from config auto_map anymore, and no path
+    # collapses model and tokenizer trust into one flag.
+    assert 'bool(getattr(model.config, "auto_map", None))' not in _SRC
+    assert "_loaded_via_remote_code(model) or _loaded_via_remote_code(tokenizer)" not in _SRC
+    assert "if _loaded_via_remote_code(model):" in _SRC  # GGUF-LoRA converter flag
+
+
+def test_compressed_export_keeps_model_and_tokenizer_trust_separate():
+    # The subprocess gets one flag per component, so an approved custom tokenizer cannot
+    # enable an unapproved model's code during compressed quantization (or vice versa).
+    assert 'cmd.append("--trust-remote-code")' in _SRC
+    assert 'cmd.append("--trust-remote-code-tokenizer")' in _SRC
+    qsrc = (_SAVE_PY.parent / "_compressed_quantize.py").read_text(encoding = "utf-8")
+    assert 'ap.add_argument("--trust-remote-code-tokenizer", action = "store_true")' in qsrc
+    assert "trust_remote_code = args.trust_remote_code_tokenizer" in qsrc
+    # The model loads keep the model flag only.
+    assert "args.model, args.trust_remote_code)" in qsrc

unsloth/_compressed_quantize.py

@@ -203,6 +203,7 @@ def main():
     ap.add_argument("--max-seq-length", type = int, default = 2048)
     ap.add_argument("--is-vlm", action = "store_true")
     ap.add_argument("--trust-remote-code", action = "store_true")
+    ap.add_argument("--trust-remote-code-tokenizer", action = "store_true")
     ap.add_argument("--variant", default = "", help = "weight-filename variant for the output shards")
     args = ap.parse_args()
 
@@ -232,7 +233,11 @@ def main():
     model.eval()
     # A tokenizer may be absent if the caller saved it separately; only calibration needs one.
     try:
-        tokenizer = auto_proc.from_pretrained(args.model, trust_remote_code = args.trust_remote_code)
+        # The tokenizer/processor has its own trust flag: consent for one component must not
+        # let the other's custom code run.
+        tokenizer = auto_proc.from_pretrained(
+            args.model, trust_remote_code = args.trust_remote_code_tokenizer
+        )
     except Exception:
         if args.needs_calibration:
             raise RuntimeError(

unsloth/save.py

@@ -223,6 +223,49 @@ def _normalize_torchao_method(save_method):
     return TORCHAO_EXPORT_SCHEMES.get(key)
 
 
+def _loaded_via_remote_code(obj):
+    """True if `obj`'s class comes from downloaded custom code (an auto_map module).
+
+    Transformers loads auto_map code into the ``transformers_modules`` package, so a
+    ``transformers_modules`` class proves the original load actually ran that remote code
+    (which the caller's / Studio's consent gate scans at load time). Export paths derive their
+    reload trust_remote_code from this - the already approved load decision - instead of from a
+    checkpoint's static ``auto_map``: a model that loads with built-in classes must not have its
+    unvetted remote code run when it is re-read during quantization export. Walks PEFT / wrapper
+    layers so a LoRA over a custom-code base is still detected, and processor components so a
+    custom tokenizer held inside a built-in processor keeps its approved trust.
+    """
+    seen = set()
+    queue = [obj]
+    while queue and len(seen) < 16:
+        node = queue.pop(0)
+        if node is None or id(node) in seen:
+            continue
+        seen.add(id(node))
+        # __module__ can be None/absent on some dynamically created or C-extension classes;
+        # treat anything non-string as "not remote code" rather than crashing the export.
+        module = getattr(type(node), "__module__", None)
+        if isinstance(module, str) and module.startswith("transformers_modules"):
+            return True
+        if hasattr(node, "get_base_model"):
+            try:
+                queue.append(node.get_base_model())
+            except Exception:
+                pass
+        # PEFT / trainer wrappers hold the real model in base_model / model; a built-in
+        # ProcessorMixin holds its (possibly custom-code) components as attributes.
+        for attr in (
+            "base_model",
+            "model",
+            "tokenizer",
+            "image_processor",
+            "feature_extractor",
+            "video_processor",
+        ):
+            queue.append(getattr(node, attr, None))
+    return False
+
+
 def _normalize_compressed_method(save_method):
     """Return (scheme, needs_calibration, suffix) if `save_method` is an FP8/FP4 compressed
     export, else None (so normal lora / merged_16bit / merged_4bit handling proceeds).
@@ -3532,7 +3575,9 @@ def _unsloth_save_lora_gguf(
             cmd += ["--base", base_model_id]
         else:
             cmd += ["--base-model-id", base_model_id]
-        if bool(getattr(model.config, "auto_map", None)):
+        # Only pass --trust-remote-code when the loaded model actually came from custom code (the
+        # approved load decision), not merely because its config carries an auto_map entry.
+        if _loaded_via_remote_code(model):
             cmd.append("--trust-remote-code")
 
         # Expose the token to the converter so it can fetch a gated/private base config from the Hub.
@@ -4387,8 +4432,8 @@ def _unsloth_save_compressed_tensors(
         )
         unsloth_generic_save(**merge_args)
 
-        # 4) Detect VLM + trust_remote_code from the in-memory model config. A vision/multimodal
-        #    model exposes a vision_config or an explicitly vision-named architecture; a bare
+        # 4) Detect VLM from the in-memory model config. A vision/multimodal model exposes a
+        #    vision_config or an explicitly vision-named architecture; a bare
         #    *ForConditionalGeneration also matches text seq2seq models (T5/BART/Whisper), so it
         #    is not treated as a VLM on its own.
         is_vlm = False
@@ -4402,9 +4447,13 @@ def _unsloth_save_compressed_tensors(
                 "Unsloth: FP8/FP4 compressed export for vision / multimodal models is "
                 "experimental; vision-tower layers may be affected."
             )
-        trust_remote_code = (
-            bool(getattr(model.config, "auto_map", None)) if hasattr(model, "config") else False
-        )
+        # trust_remote_code must reflect the approved load decision (whether the model / tokenizer
+        # was actually loaded from custom code), not the config's static auto_map, so a
+        # built-in-loadable model carrying auto_map cannot run unvetted code in the subprocess.
+        # Model and tokenizer trust stay separate, like the torchao path: an approved custom
+        # tokenizer must not enable an unapproved model's code in the subprocess (or vice versa).
+        model_trust = _loaded_via_remote_code(model)
+        tok_trust = _loaded_via_remote_code(tokenizer)
 
         # 5) Marshal the calibration dataset for the subprocess: None -> ultrachat default; a
         #    str/PathLike is a local save_to_disk dir if it exists else a Hub id; Dataset -> temp.
@@ -4479,8 +4528,10 @@ def _unsloth_save_compressed_tensors(
             cmd += ["--calibration-dataset", calib_value]
         if is_vlm:
             cmd.append("--is-vlm")
-        if trust_remote_code:
+        if model_trust:
             cmd.append("--trust-remote-code")
+        if tok_trust:
+            cmd.append("--trust-remote-code-tokenizer")
         if variant:
             cmd += ["--variant", variant]
 
@@ -4679,35 +4730,19 @@ def _unsloth_save_torchao(
         )
         unsloth_generic_save(**merge_args)
 
-        # 2) Detect VLM + trust_remote_code so the right auto class reloads the staged checkpoint.
-        #    A bare *ForConditionalGeneration also matches text seq2seq (T5/BART/Whisper), so key off
-        #    vision_config / a vision-named architecture only, like the compressed path.
+        # 2) Detect VLM + reload class. A bare *ForConditionalGeneration also matches text seq2seq
+        #    (T5/BART/Whisper), so key off vision_config / a vision-named architecture only.
         is_vlm = False
-        trust_remote_code = False
         if hasattr(model, "config"):
             archs = getattr(model.config, "architectures", None) or []
             is_vlm = hasattr(model.config, "vision_config") or any(
                 x.endswith("ForVisionText2Text") for x in archs
             )
-            trust_remote_code = bool(getattr(model.config, "auto_map", None))
-        # Custom code can be declared only in the tokenizer/processor config, so also honor an
-        # auto_map in any staged config (the original load already had the user's consent).
-        if not trust_remote_code:
-            for _cfg in (
-                "config.json",
-                "tokenizer_config.json",
-                "processor_config.json",
-                "preprocessor_config.json",
-            ):
-                try:
-                    _p = os.path.join(staging, _cfg)
-                    if os.path.exists(_p):
-                        with open(_p, "r", encoding = "utf-8") as _f:
-                            if "auto_map" in json.load(_f):
-                                trust_remote_code = True
-                                break
-                except Exception:
-                    pass
+        # trust_remote_code must reflect the approved load decision - whether the in-memory model /
+        # tokenizer was itself loaded from custom code - not the staged config's auto_map, which an
+        # attacker can set on a built-in-loadable model to run unvetted code past the consent gate.
+        model_trust = _loaded_via_remote_code(model)
+        tok_trust = _loaded_via_remote_code(tokenizer)
         # Reload with the class that matches the checkpoint: an image-text VLM class (with a
         # fallback for older Transformers that lack AutoModelForImageTextToText); the model's own
         # architecture class for encoder-decoder seq2seq (T5/BART/Whisper are not causal LMs, and
@@ -4766,12 +4801,10 @@ def _unsloth_save_torchao(
             staging,
             device_map = "auto",
             quantization_config = TorchAoConfig(quant_type = quant_type),
-            trust_remote_code = trust_remote_code,
+            trust_remote_code = model_trust,
             **dtype_kw,
         )
-        staged_tokenizer = auto_processor.from_pretrained(
-            staging, trust_remote_code = trust_remote_code
-        )
+        staged_tokenizer = auto_processor.from_pretrained(staging, trust_remote_code = tok_trust)
 
         quantized_model.save_pretrained(out_dir, safe_serialization = safe_serialization)
         staged_tokenizer.save_pretrained(out_dir)