Parameterize trust_remote_code for quark_visitai

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), inherits from the model's load kwargs; otherwise overrides.",
+            ),
             **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",

olive/passes/quark_vitisai/torch/language_modeling/llm_ptq/quantize_quark.py

@@ -44,15 +44,23 @@ 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)
+        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)

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)