From 697aba6d9968b3a766a8d238ba262110a6025ac1 Mon Sep 17 00:00:00 2001 From: hotragn Date: Mon, 20 Jul 2026 05:50:37 -0500 Subject: [PATCH] Default to trust_remote_code=False, fall back only when required AutoModel.get_module_class (config load) and the base/MLX tokenizer loads hard-coded trust_remote_code=True, so a normal AirLLM load of an arbitrary Hugging Face repo id would execute Python shipped in that repo -- even for standard architectures that don't need it. AirLLMBaseModel already loads the model config/weights with trust_remote_code=False and only falls back to True when transformers can't parse the model natively; this makes every other load boundary consistent with that. - Add load_prefer_no_remote_code() helper: try trust_remote_code=False, fall back to True only if the native load raises. - Use it for the config load in AutoModel.get_module_class, the tokenizer load in AirLLMBaseModel, and the config + tokenizer loads in the MLX path. - Standard models never trigger remote code now; models that genuinely need it (ChatGLM, Baichuan, some Qwen) still work via the fallback. No API change. - Add offline unit tests for the helper and the get_module_class behavior. Addresses #293. --- air_llm/airllm/airllm_base.py | 13 ++-- air_llm/airllm/airllm_llama_mlx.py | 16 ++--- air_llm/airllm/auto_model.py | 12 ++-- air_llm/airllm/utils.py | 18 +++++ air_llm/tests/test_trust_remote_code.py | 95 +++++++++++++++++++++++++ 5 files changed, 134 insertions(+), 20 deletions(-) create mode 100644 air_llm/tests/test_trust_remote_code.py diff --git a/air_llm/airllm/airllm_base.py b/air_llm/airllm/airllm_base.py index 8da7ab91..fec81526 100644 --- a/air_llm/airllm/airllm_base.py +++ b/air_llm/airllm/airllm_base.py @@ -16,8 +16,8 @@ from .profiler import LayeredProfiler from .utils import clean_memory, load_layer, layer_tensor_names, load_layer_subset, \ - find_or_create_local_splitted_path, load_merged_ngram_embedding, \ - open_ngram_mmap_table, MmapEmbedding, _force_meta_embeddings + find_or_create_local_splitted_path, load_prefer_no_remote_code, \ + load_merged_ngram_embedding, open_ngram_mmap_table, MmapEmbedding, _force_meta_embeddings from .persist import ModelPersister try: @@ -211,10 +211,11 @@ def get_generation_config(self): return GenerationConfig() def get_tokenizer(self, hf_token=None): - if hf_token is not None: - return AutoTokenizer.from_pretrained(self.model_local_path, token=hf_token, trust_remote_code=True) - else: - return AutoTokenizer.from_pretrained(self.model_local_path, trust_remote_code=True) + token_kwargs = {'token': hf_token} if hf_token is not None else {} + # Prefer transformers' native tokenizer; only trust the repo's remote code if it's required + # (custom tokenizers). Matches how the config/model are loaded above. + return load_prefer_no_remote_code( + AutoTokenizer.from_pretrained, self.model_local_path, **token_kwargs) # ---- model construction ----------------------------------------------------------------- diff --git a/air_llm/airllm/airllm_llama_mlx.py b/air_llm/airllm/airllm_llama_mlx.py index e47a0bd4..26c45522 100644 --- a/air_llm/airllm/airllm_llama_mlx.py +++ b/air_llm/airllm/airllm_llama_mlx.py @@ -15,7 +15,7 @@ import psutil from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, AutoModel, GenerationMixin, LlamaForCausalLM, GenerationConfig from .utils import clean_memory, load_layer, \ - find_or_create_local_splitted_path + find_or_create_local_splitted_path, load_prefer_no_remote_code @@ -227,10 +227,9 @@ def __init__(self, model_local_path_or_repo_id, device="cuda:0", dtype=None, max layer_names=self.layer_names_dict, hf_token=hf_token, delete_original=delete_original) - if hf_token is not None: - self.config = AutoConfig.from_pretrained(self.model_local_path, token=hf_token, trust_remote_code=True) - else: - self.config = AutoConfig.from_pretrained(self.model_local_path, trust_remote_code=True) + token_kwargs = {'token': hf_token} if hf_token is not None else {} + self.config = load_prefer_no_remote_code( + AutoConfig.from_pretrained, self.model_local_path, **token_kwargs) self.model_args = get_model_args_from_config(self.config) @@ -243,10 +242,9 @@ def __init__(self, model_local_path_or_repo_id, device="cuda:0", dtype=None, max def get_tokenizer(self, hf_token=None): - if hf_token is not None: - return AutoTokenizer.from_pretrained(self.model_local_path, token=hf_token, trust_remote_code=True) - else: - return AutoTokenizer.from_pretrained(self.model_local_path, trust_remote_code=True) + token_kwargs = {'token': hf_token} if hf_token is not None else {} + return load_prefer_no_remote_code( + AutoTokenizer.from_pretrained, self.model_local_path, **token_kwargs) def generate(self, x, temperature=0, max_new_tokens=None, **kwargs): diff --git a/air_llm/airllm/auto_model.py b/air_llm/airllm/auto_model.py index f6608dfd..af44e461 100644 --- a/air_llm/airllm/auto_model.py +++ b/air_llm/airllm/auto_model.py @@ -2,6 +2,8 @@ from transformers import AutoConfig from sys import platform +from .utils import load_prefer_no_remote_code + is_on_mac_os = False if platform == "darwin": @@ -36,11 +38,11 @@ def __init__(self): @classmethod def get_module_class(cls, pretrained_model_name_or_path, *inputs, **kwargs): - if 'hf_token' in kwargs: - config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True, - token=kwargs['hf_token']) - else: - config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True) + token_kwargs = {'token': kwargs['hf_token']} if 'hf_token' in kwargs else {} + # Prefer transformers' native config; only run the repo's remote code if it's actually + # required to parse the config (custom architectures). See load_prefer_no_remote_code. + config = load_prefer_no_remote_code( + AutoConfig.from_pretrained, pretrained_model_name_or_path, **token_kwargs) architectures = getattr(config, "architectures", None) or [] arch = architectures[0] if architectures else "" diff --git a/air_llm/airllm/utils.py b/air_llm/airllm/utils.py index 75740db4..5824d153 100644 --- a/air_llm/airllm/utils.py +++ b/air_llm/airllm/utils.py @@ -85,6 +85,24 @@ def clean_memory(): torch.cuda.empty_cache() +def load_prefer_no_remote_code(loader, *args, **kwargs): + """ + Call a transformers ``from_pretrained``-style loader, preferring transformers' native + implementation and only trusting the model's bundled remote code when the native load fails. + + AirLLM's advertised entry point is an arbitrary Hugging Face repo id, so a normal load of a + standard architecture should not execute Python shipped in the repo. Standard models load fine + with ``trust_remote_code=False``; only models that genuinely ship custom code (ChatGLM, Baichuan, + some Qwen) raise without it, so we fall back to ``trust_remote_code=True`` for those. This mirrors + how ``AirLLMBaseModel`` already loads the model config/weights, and keeps every load boundary + (config, tokenizer, model) consistent instead of hard-coding ``trust_remote_code=True``. + """ + try: + return loader(*args, trust_remote_code=False, **kwargs) + except Exception: + return loader(*args, trust_remote_code=True, **kwargs) + + def uncompress_layer_state_dict(layer_state_dict): uncompressed_layer_state_dict = None if any(['4bit' in k for k in layer_state_dict.keys()]): diff --git a/air_llm/tests/test_trust_remote_code.py b/air_llm/tests/test_trust_remote_code.py new file mode 100644 index 00000000..c242a205 --- /dev/null +++ b/air_llm/tests/test_trust_remote_code.py @@ -0,0 +1,95 @@ +import unittest + +import airllm.auto_model as auto_model_mod +from airllm.auto_model import AutoModel +from airllm.utils import load_prefer_no_remote_code + + +class _FakeConfig: + def __init__(self, architectures): + self.architectures = architectures + + +class TestLoadPreferNoRemoteCode(unittest.TestCase): + def test_uses_false_and_does_not_retry_when_native_load_succeeds(self): + calls = [] + + def loader(path, trust_remote_code=None, **kwargs): + calls.append(trust_remote_code) + return 'ok' + + result = load_prefer_no_remote_code(loader, 'some/repo') + + self.assertEqual(result, 'ok') + self.assertEqual(calls, [False]) # tried False, never fell back to True + + def test_falls_back_to_true_when_native_load_raises(self): + calls = [] + + def loader(path, trust_remote_code=None, **kwargs): + calls.append(trust_remote_code) + if trust_remote_code is False: + raise ValueError('requires trust_remote_code=True') + return 'ok-remote' + + result = load_prefer_no_remote_code(loader, 'custom/repo') + + self.assertEqual(result, 'ok-remote') + self.assertEqual(calls, [False, True]) # tried False first, then True + + def test_forwards_extra_kwargs(self): + seen = {} + + def loader(path, trust_remote_code=None, **kwargs): + seen.update(kwargs) + seen['path'] = path + return 'ok' + + load_prefer_no_remote_code(loader, 'some/repo', token='abc') + + self.assertEqual(seen['path'], 'some/repo') + self.assertEqual(seen['token'], 'abc') + + +class TestGetModuleClassTrust(unittest.TestCase): + """AutoModel.get_module_class must not enable remote code unless it is actually required.""" + + def setUp(self): + self._orig = auto_model_mod.AutoConfig.from_pretrained + + def tearDown(self): + auto_model_mod.AutoConfig.from_pretrained = self._orig + + def test_standard_arch_loads_config_without_remote_code(self): + calls = [] + + def fake(path, trust_remote_code=None, **kwargs): + calls.append(trust_remote_code) + return _FakeConfig(['LlamaForCausalLM']) + + auto_model_mod.AutoConfig.from_pretrained = staticmethod(fake) + + module, cls = AutoModel.get_module_class('some/llama-repo') + + self.assertEqual((module, cls), ('airllm', 'AirLLMBaseModel')) + self.assertEqual(calls, [False]) # native path only; remote code never enabled + + def test_custom_arch_falls_back_to_remote_code(self): + calls = [] + + def fake(path, trust_remote_code=None, **kwargs): + calls.append(trust_remote_code) + if trust_remote_code is False: + raise ValueError('requires trust_remote_code=True') + return _FakeConfig(['ChatGLMModel']) + + auto_model_mod.AutoConfig.from_pretrained = staticmethod(fake) + + module, cls = AutoModel.get_module_class('some/chatglm-repo') + + self.assertEqual((module, cls), ('airllm', 'AirLLMChatGLM')) + self.assertEqual(calls, [False, True]) # only trusts remote code as a fallback + + +if __name__ == '__main__': + unittest.main()