64f6526160* 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>
150 lines
6 KiB
Python
150 lines
6 KiB
Python
"""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
|