From f5ef0abbbc75c1c462a5852236027df1a2709409 Mon Sep 17 00:00:00 2001 From: shaahji Date: Thu, 16 Jul 2026 15:01:01 -0700 Subject: [PATCH 1/2] Parameterize trust_remote_code for quark_visitai trust_remote_code was being defaulted to True raising security concerns. Replacing the default with user driven value thru' the config. --- .../quark_quantization_vitisai.py | 12 +++++++++ .../llm_ptq/quantize_quark.py | 12 +++++++-- .../llm_utils/model_preparation.py | 25 +++++++++++++------ 3 files changed, 40 insertions(+), 9 deletions(-) diff --git a/olive/passes/quark_vitisai/quark_quantization_vitisai.py b/olive/passes/quark_vitisai/quark_quantization_vitisai.py index 652f4d3145..71d4884bd9 100644 --- a/olive/passes/quark_vitisai/quark_quantization_vitisai.py +++ b/olive/passes/quark_vitisai/quark_quantization_vitisai.py @@ -98,6 +98,11 @@ def _default_config(cls, accelerator_spec=None): "extra_options": PassConfigParam( type_=dict, default_value=None, description="Extra options for quantization. Default is {}." ), + "trust_remote_code": PassConfigParam( + type_=Union[bool, None], + default_value=None, + description="Trust remote code when loading models. If None (default), uses Olive's safe-by-default mechanism which must be explicitly enabled by user.", + ), **get_external_data_config(), } @@ -169,6 +174,12 @@ def _run_quark_torch(self, model: HfModelHandler, config: BasePassConfig, output device = "cuda" else: device = "cpu" + # Extract trust_remote_code from model's load kwargs, defaulting to False for safety + trust_remote_code = model.get_load_kwargs().get("trust_remote_code", False) + # Override with explicit config value if provided + if config.trust_remote_code is not None: + trust_remote_code = config.trust_remote_code + quant_algo_config_file_path = None if config.quant_config: with tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".json") as tmp_file: @@ -188,6 +199,7 @@ def _run_quark_torch(self, model: HfModelHandler, config: BasePassConfig, output exclude_layers=config.exclude_layers, device=device, quant_algo_config_file_path=quant_algo_config_file_path, + trust_remote_code=trust_remote_code, # Other args multi_gpu=False, model_attn_implementation="eager", diff --git a/olive/passes/quark_vitisai/torch/language_modeling/llm_ptq/quantize_quark.py b/olive/passes/quark_vitisai/torch/language_modeling/llm_ptq/quantize_quark.py index 6270508c54..1c4e6479ef 100644 --- a/olive/passes/quark_vitisai/torch/language_modeling/llm_ptq/quantize_quark.py +++ b/olive/passes/quark_vitisai/torch/language_modeling/llm_ptq/quantize_quark.py @@ -44,12 +44,20 @@ def run_quark_quantization(args: argparse.Namespace) -> None: device = args.device model, _ = get_model( - args.model_dir, args.data_type, device, args.multi_gpu, args.multi_device, args.model_attn_implementation + args.model_dir, + args.data_type, + device, + args.multi_gpu, + args.multi_device, + args.model_attn_implementation, + args.trust_remote_code, ) prepare_for_moe_quant(model) model_type = get_model_type(model) - tokenizer = get_tokenizer(args.model_dir, max_seq_len=args.seq_len, model_type=model_type) + tokenizer = get_tokenizer( + args.model_dir, max_seq_len=args.seq_len, model_type=model_type, trust_remote_code=args.trust_remote_code + ) multimodal = model_type in ["mllama"] if multimodal: processor = AutoProcessor.from_pretrained(args.model_dir) diff --git a/olive/passes/quark_vitisai/torch/language_modeling/llm_utils/model_preparation.py b/olive/passes/quark_vitisai/torch/language_modeling/llm_utils/model_preparation.py index 3cb5f51a89..9604c2171b 100644 --- a/olive/passes/quark_vitisai/torch/language_modeling/llm_utils/model_preparation.py +++ b/olive/passes/quark_vitisai/torch/language_modeling/llm_utils/model_preparation.py @@ -110,11 +110,17 @@ } -def get_tokenizer(ckpt_path: str, max_seq_len: int = 2048, model_type: Optional[str] = None) -> AutoTokenizer: +def get_tokenizer( + ckpt_path: str, max_seq_len: int = 2048, model_type: Optional[str] = None, trust_remote_code: bool = False +) -> AutoTokenizer: logger.info("Initializing tokenizer from %s", ckpt_path) use_fast = model_type in ["grok", "cohere", "olmo", "instella", "deepseekv2v3"] tokenizer = AutoTokenizer.from_pretrained( - ckpt_path, model_max_length=max_seq_len, padding_side="left", trust_remote_code=True, use_fast=use_fast + ckpt_path, + model_max_length=max_seq_len, + padding_side="left", + trust_remote_code=trust_remote_code, + use_fast=use_fast, ) if model_type and model_type in ["qwen", "qwen2"]: # qwen2 use token id 151643 as pad and eos tokens @@ -151,6 +157,7 @@ def get_model( multi_gpu: bool = False, multi_device=False, attn_implementation: str = "eager", + trust_remote_code: bool = False, ) -> tuple[nn.Module, torch.dtype]: if data_type == "float16": model_dtype = torch.float16 @@ -183,7 +190,7 @@ def get_model( device_map=device, torch_dtype=model_dtype, max_memory=max_memory, - trust_remote_code=True, + trust_remote_code=trust_remote_code, attn_implementation=attn_implementation, ) else: @@ -193,12 +200,16 @@ def get_model( device_map=device, torch_dtype=model_dtype, max_memory=max_memory, - trust_remote_code=True, + trust_remote_code=trust_remote_code, attn_implementation=attn_implementation, ) except Exception: model = AutoModelForCausalLM.from_pretrained( - ckpt_path, device_map=device, torch_dtype=model_dtype, max_memory=max_memory, trust_remote_code=True + ckpt_path, + device_map=device, + torch_dtype=model_dtype, + max_memory=max_memory, + trust_remote_code=trust_remote_code, ) if multi_device and hasattr(model, "hf_device_map"): logger.info("device_map: %s", model.hf_device_map) @@ -224,13 +235,13 @@ def get_model_type(model: nn.Module) -> str: return "unknown" -def save_model(model: nn.Module, tokenizer: AutoTokenizer, save_dir: str) -> None: +def save_model(model: nn.Module, tokenizer: AutoTokenizer, save_dir: str, trust_remote_code: bool = False) -> None: model.save_pretrained(save_dir, safe_serialization=True) model_name_or_path = getattr(model.config, "name_or_path", None) if tokenizer is None and model_name_or_path: try: - tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=trust_remote_code) logger.info("Saved the tokenizer from pretrained: %s", model_name_or_path) except Exception as e: logger.info("An error occurred when loading tokenizer: %s", e) From a62554f4c4a0d36cf5e69b3e31a060081e71db9c Mon Sep 17 00:00:00 2001 From: shaahji <96227573+shaahji@users.noreply.github.com> Date: Thu, 16 Jul 2026 15:21:09 -0700 Subject: [PATCH 2/2] Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- olive/passes/quark_vitisai/quark_quantization_vitisai.py | 2 +- .../torch/language_modeling/llm_ptq/quantize_quark.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/olive/passes/quark_vitisai/quark_quantization_vitisai.py b/olive/passes/quark_vitisai/quark_quantization_vitisai.py index 71d4884bd9..3ee906360f 100644 --- a/olive/passes/quark_vitisai/quark_quantization_vitisai.py +++ b/olive/passes/quark_vitisai/quark_quantization_vitisai.py @@ -101,7 +101,7 @@ def _default_config(cls, accelerator_spec=None): "trust_remote_code": PassConfigParam( type_=Union[bool, None], default_value=None, - description="Trust remote code when loading models. If None (default), uses Olive's safe-by-default mechanism which must be explicitly enabled by user.", + description="Trust remote code when loading models. If None (default), inherits from the model's load kwargs; otherwise overrides.", ), **get_external_data_config(), } diff --git a/olive/passes/quark_vitisai/torch/language_modeling/llm_ptq/quantize_quark.py b/olive/passes/quark_vitisai/torch/language_modeling/llm_ptq/quantize_quark.py index 1c4e6479ef..2790c4087e 100644 --- a/olive/passes/quark_vitisai/torch/language_modeling/llm_ptq/quantize_quark.py +++ b/olive/passes/quark_vitisai/torch/language_modeling/llm_ptq/quantize_quark.py @@ -60,7 +60,7 @@ def run_quark_quantization(args: argparse.Namespace) -> None: ) multimodal = model_type in ["mllama"] if multimodal: - processor = AutoProcessor.from_pretrained(args.model_dir) + processor = AutoProcessor.from_pretrained(args.model_dir, trust_remote_code=args.trust_remote_code) if args.model_export is not None: export_dir = Path(args.output_dir) export_dir.mkdir(parents=True, exist_ok=True)