From 3ab36e2aad434f3fa26d53a8d2a2b36d94dbb79d Mon Sep 17 00:00:00 2001 From: Alpamys Date: Sat, 9 May 2026 13:20:13 +0500 Subject: [PATCH] feat(security): trust_remote_code opt-in across non-SFT trainers + 5 commands (v0.40.4 Part A) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the v0.36.0 #63 known gap. Every non-SFT trainer wrapper (DPO / GRPO / KTO / ORPO / SimPO / IPO / PPO / RewardModel / Pretrain / Embedding / BCO + the unified Preference dispatcher) now accepts trust_remote_code: bool = False on __init__, resolves once via the v0.36.0 helper (model_requires_trust_remote_code + resolve_trust_remote_code), and stores the resolved value on self._trust_remote_code. Every from_pretrained call site reads from the resolved attribute — no remaining trust_remote_code=True literal in any trainer file (asserted by tests/test_v0404_part_a.py). Five standalone commands gain a --trust-remote-code Typer flag with the same default-deny + KNOWN_SAFE_PREFIXES allowlist behaviour as soup train: soup diff, soup export, soup merge, soup infer, soup data generate. commands/train.py removes the v0.36.0 sft_kwargs split — every trainer receives trust_remote_code from the same trainer_kwargs dict. PreferenceTrainerWrapper forwards the raw bool to the inner DPO / SimPO / ORPO / IPO / BCO wrapper kwargs at both _build_inner and _build_multi_objective sites; the resolver fires inside the inner wrapper at construction time. _load_reward_model (module-level helper in ppo.py) accepts a trust_remote_code: bool parameter and resolves internally — design intent is that the helper is independently safe to call outside PPOTrainerWrapper. _export_onnx / _export_tensorrt / _export_awq / _export_gptq and _merge_adapter helpers all gain a trust_remote_code: bool = False parameter threaded from the Typer flag. Co-Authored-By: Claude Opus 4.7 (1M context) --- soup_cli/commands/diff.py | 43 +++- soup_cli/commands/export.py | 90 ++++++-- soup_cli/commands/generate.py | 35 ++- soup_cli/commands/infer.py | 45 +++- soup_cli/commands/merge.py | 27 ++- soup_cli/commands/train.py | 9 +- soup_cli/trainer/bco.py | 22 +- soup_cli/trainer/dpo.py | 25 ++- soup_cli/trainer/embedding.py | 22 +- soup_cli/trainer/grpo.py | 22 +- soup_cli/trainer/ipo.py | 22 +- soup_cli/trainer/kto.py | 22 +- soup_cli/trainer/orpo.py | 22 +- soup_cli/trainer/ppo.py | 52 ++++- soup_cli/trainer/preference.py | 4 + soup_cli/trainer/pretrain.py | 56 ++++- soup_cli/trainer/reward_model.py | 24 ++- soup_cli/trainer/simpo.py | 22 +- tests/test_v0404_part_a.py | 353 +++++++++++++++++++++++++++++++ 19 files changed, 835 insertions(+), 82 deletions(-) create mode 100644 tests/test_v0404_part_a.py diff --git a/soup_cli/commands/diff.py b/soup_cli/commands/diff.py index 0a3e45c..b0238f9 100644 --- a/soup_cli/commands/diff.py +++ b/soup_cli/commands/diff.py @@ -69,6 +69,14 @@ def diff( "-o", help="Save results to JSONL file", ), + trust_remote_code: bool = typer.Option( + False, + "--trust-remote-code", + help=( + "Allow loading models that ship custom Python via auto_map. " + "Default deny (v0.36.0). Only enable if you trust the source." + ), + ), ): """Compare outputs of two models side-by-side on the same prompts.""" # Validate model paths @@ -105,9 +113,13 @@ def diff( # Load models console.print("[dim]Loading Model A...[/]") - model_obj_a, tokenizer_a = _load_model(str(path_a), base_a, device) + model_obj_a, tokenizer_a = _load_model( + str(path_a), base_a, device, trust_remote_code, + ) console.print("[dim]Loading Model B...[/]") - model_obj_b, tokenizer_b = _load_model(str(path_b), base_b, device) + model_obj_b, tokenizer_b = _load_model( + str(path_b), base_b, device, trust_remote_code, + ) console.print("[green]Both models loaded.[/]\n") # Run comparison @@ -199,11 +211,21 @@ def _collect_prompts(prompts_file: Optional[str], prompt_args: Optional[list[str return result -def _load_model(model_path: str, base_model: Optional[str], device: str): +def _load_model( + model_path: str, + base_model: Optional[str], + device: str, + trust_remote_code: bool = False, +): """Load a model and tokenizer.""" import torch from transformers import AutoModelForCausalLM, AutoTokenizer + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + path = Path(model_path) adapter_config_path = path / "adapter_config.json" is_adapter = adapter_config_path.exists() @@ -220,7 +242,16 @@ def _load_model(model_path: str, base_model: Optional[str], device: str): console.print(f"[red]Cannot detect base model for {path}. Use --base-a/--base-b.[/]") raise typer.Exit(1) - tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + probe_target = base_model or model_path + requires = model_requires_trust_remote_code(model_path) or False + trc = resolve_trust_remote_code( + probe_target, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) + + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=trc) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token @@ -229,7 +260,7 @@ def _load_model(model_path: str, base_model: Optional[str], device: str): base = AutoModelForCausalLM.from_pretrained( base_model, - trust_remote_code=True, + trust_remote_code=trc, device_map="auto", dtype=torch.float16, ) @@ -237,7 +268,7 @@ def _load_model(model_path: str, base_model: Optional[str], device: str): else: model_obj = AutoModelForCausalLM.from_pretrained( model_path, - trust_remote_code=True, + trust_remote_code=trc, device_map="auto", dtype=torch.float16, ) diff --git a/soup_cli/commands/export.py b/soup_cli/commands/export.py index 47ab1d3..3ba8843 100644 --- a/soup_cli/commands/export.py +++ b/soup_cli/commands/export.py @@ -97,6 +97,14 @@ def export( help="Attach exported artifact to this registry entry " "(default: auto-match by source --model output dir)", ), + trust_remote_code: bool = typer.Option( + False, + "--trust-remote-code", + help=( + "Allow loading models that ship custom Python via auto_map. " + "Default deny (v0.36.0). Only enable if you trust the source." + ), + ), ): """Export a model to GGUF, ONNX, TensorRT-LLM, AWQ, or GPTQ format.""" model_path = Path(model) @@ -115,19 +123,19 @@ def export( # --- ONNX export path --- if fmt == "onnx": - _export_onnx(model_path, output, base, onnx_task) + _export_onnx(model_path, output, base, onnx_task, trust_remote_code) return # --- TensorRT-LLM export path --- if fmt == "tensorrt": - _export_tensorrt(model_path, output, base) + _export_tensorrt(model_path, output, base, trust_remote_code) return # --- AWQ export path --- if fmt == "awq": _export_awq( model_path, output, base, bits, group_size, - calibration_data, calibration_samples, + calibration_data, calibration_samples, trust_remote_code, ) return @@ -135,7 +143,7 @@ def export( if fmt == "gptq": _export_gptq( model_path, output, base, bits, group_size, - calibration_data, calibration_samples, + calibration_data, calibration_samples, trust_remote_code, ) return @@ -162,7 +170,9 @@ def export( raise typer.Exit(1) merge_dir = model_path.parent / f".soup_merge_tmp_{model_path.name}" - _merge_adapter(str(model_path), base_model, str(merge_dir)) + _merge_adapter( + str(model_path), base_model, str(merge_dir), trust_remote_code, + ) model_path = merge_dir # --- Find llama.cpp --- @@ -264,17 +274,35 @@ def _detect_base_model(adapter_config_path: Path) -> Optional[str]: return None -def _merge_adapter(adapter_path: str, base_model: str, output_dir: str): +def _merge_adapter( + adapter_path: str, + base_model: str, + output_dir: str, + trust_remote_code: bool = False, +): """Merge LoRA adapter with base model.""" import torch from peft import PeftModel from transformers import AutoModelForCausalLM, AutoTokenizer + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + + requires = model_requires_trust_remote_code(adapter_path) or False + trc = resolve_trust_remote_code( + base_model, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) + console.print(f"[dim]Loading base model: {base_model}...[/]") model = AutoModelForCausalLM.from_pretrained( base_model, dtype=torch.float16, - trust_remote_code=True, + trust_remote_code=trc, device_map="cpu", ) @@ -288,7 +316,7 @@ def _merge_adapter(adapter_path: str, base_model: str, output_dir: str): out.mkdir(parents=True, exist_ok=True) model.save_pretrained(str(out)) - tokenizer = AutoTokenizer.from_pretrained(adapter_path, trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained(adapter_path, trust_remote_code=trc) tokenizer.save_pretrained(str(out)) console.print("[green]Adapter merged successfully.[/]") @@ -419,6 +447,7 @@ def _find_quantize_binary(llama_dir: Path) -> Optional[Path]: def _export_onnx( model_path: Path, output: Optional[str], base: Optional[str], task: str = "text-generation", + trust_remote_code: bool = False, ): """Export model to ONNX format via optimum.""" try: @@ -447,7 +476,7 @@ def _export_onnx( ) raise typer.Exit(1) merge_dir = model_path.parent / f".soup_merge_tmp_{model_path.name}" - _merge_adapter(str(model_path), base_model, str(merge_dir)) + _merge_adapter(str(model_path), base_model, str(merge_dir), trust_remote_code) source_path = merge_dir output_path = Path(output) if output else model_path.parent / f"{model_path.name}_onnx" @@ -492,7 +521,12 @@ def _export_onnx( ) -def _export_tensorrt(model_path: Path, output: Optional[str], base: Optional[str]): +def _export_tensorrt( + model_path: Path, + output: Optional[str], + base: Optional[str], + trust_remote_code: bool = False, +): """Export model to TensorRT-LLM format.""" # TensorRT-LLM uses trtllm-build CLI from the tensorrt_llm package trtllm_available = False @@ -527,7 +561,7 @@ def _export_tensorrt(model_path: Path, output: Optional[str], base: Optional[str ) raise typer.Exit(1) merge_dir = model_path.parent / f".soup_merge_tmp_{model_path.name}" - _merge_adapter(str(model_path), base_model, str(merge_dir)) + _merge_adapter(str(model_path), base_model, str(merge_dir), trust_remote_code) source_path = merge_dir output_path = Path(output) if output else model_path.parent / f"{model_path.name}_trt" @@ -677,6 +711,7 @@ def _export_awq( group_size: int = 128, calibration_data: Optional[str] = None, calibration_samples: int = 128, + trust_remote_code: bool = False, ) -> None: """Export model to AWQ format via autoawq.""" # Validate bits @@ -721,7 +756,7 @@ def _export_awq( ) raise typer.Exit(1) merge_dir = model_path.parent / f".soup_merge_tmp_{model_path.name}" - _merge_adapter(str(model_path), base_model, str(merge_dir)) + _merge_adapter(str(model_path), base_model, str(merge_dir), trust_remote_code) source_path = merge_dir default_out = model_path.parent / f"{model_path.name}_awq" @@ -748,7 +783,19 @@ def _export_awq( ) console.print("[dim]Loading model for AWQ quantization...[/]") model = AutoAWQForCausalLM.from_pretrained(str(source_path)) - tokenizer = AutoTokenizer.from_pretrained(str(source_path), trust_remote_code=True) + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + + requires_tok = model_requires_trust_remote_code(str(source_path)) or False + trc_tok = resolve_trust_remote_code( + str(source_path), + requested=trust_remote_code, + console=console, + requires_remote_code=requires_tok, + ) + tokenizer = AutoTokenizer.from_pretrained(str(source_path), trust_remote_code=trc_tok) quant_config = {"zero_point": True, "q_group_size": group_size, "w_bit": bits} @@ -796,6 +843,7 @@ def _export_gptq( group_size: int = 128, calibration_data: Optional[str] = None, calibration_samples: int = 128, + trust_remote_code: bool = False, ) -> None: """Export model to GPTQ format via auto-gptq.""" # Validate bits @@ -840,7 +888,7 @@ def _export_gptq( ) raise typer.Exit(1) merge_dir = model_path.parent / f".soup_merge_tmp_{model_path.name}" - _merge_adapter(str(model_path), base_model, str(merge_dir)) + _merge_adapter(str(model_path), base_model, str(merge_dir), trust_remote_code) source_path = merge_dir default_out = model_path.parent / f"{model_path.name}_gptq" @@ -874,7 +922,19 @@ def _export_gptq( model = AutoGPTQForCausalLM.from_pretrained( str(source_path), quantize_config=quantize_config ) - tokenizer = AutoTokenizer.from_pretrained(str(source_path), trust_remote_code=True) + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + + requires_tok = model_requires_trust_remote_code(str(source_path)) or False + trc_tok = resolve_trust_remote_code( + str(source_path), + requested=trust_remote_code, + console=console, + requires_remote_code=requires_tok, + ) + tokenizer = AutoTokenizer.from_pretrained(str(source_path), trust_remote_code=trc_tok) # Load calibration data if provided calib_data = None diff --git a/soup_cli/commands/generate.py b/soup_cli/commands/generate.py index 2543d7c..b123847 100644 --- a/soup_cli/commands/generate.py +++ b/soup_cli/commands/generate.py @@ -161,6 +161,14 @@ def generate( "--rpm", help="Rate limit for API requests (default: 60)", ), + trust_remote_code: bool = typer.Option( + False, + "--trust-remote-code", + help=( + "Allow loading local-provider models that ship custom Python " + "via auto_map. Default deny (v0.36.0)." + ), + ), ): """Generate synthetic training data using an LLM.""" if fmt not in VALID_FORMATS: @@ -288,6 +296,7 @@ def generate( context_text=context_text, template_pref_task=template_pref_task, template_domain=template_domain, + trust_remote_code=trust_remote_code, ) except Exception as exc: console.print(f"[red]Generation error: {exc}[/]") @@ -485,6 +494,7 @@ def _generate_batch( context_text: str = "", template_pref_task: str = "dpo", template_domain: str = "math", + trust_remote_code: bool = False, ) -> list[dict]: """Generate a batch of examples using the specified provider.""" # Build the generation prompt — template or default @@ -526,6 +536,7 @@ def _generate_batch( temperature=temperature, seed_examples=seed_examples, generation_prompt=generation_prompt if template else None, + trust_remote_code=trust_remote_code, ) elif provider == "server": return _generate_server( @@ -737,26 +748,32 @@ def _generate_local( temperature: float, seed_examples: list[dict], generation_prompt: Optional[str] = None, + trust_remote_code: bool = False, ) -> list[dict]: """Generate examples using a local model via transformers.""" import torch - from rich.panel import Panel from transformers import AutoModelForCausalLM, AutoTokenizer - console.print(Panel( - f"[yellow]Loading model with trust_remote_code=True: {model_name}[/]\n" - "This executes code from the model repository.", - title="Security Warning", - style="yellow", - )) + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) - tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) + requires = model_requires_trust_remote_code(model_name) or False + trc = resolve_trust_remote_code( + model_name, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) + + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=trc) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( model_name, - trust_remote_code=True, + trust_remote_code=trc, device_map="auto", torch_dtype=torch.float16, ) diff --git a/soup_cli/commands/infer.py b/soup_cli/commands/infer.py index 65b2c7f..c91d320 100644 --- a/soup_cli/commands/infer.py +++ b/soup_cli/commands/infer.py @@ -94,6 +94,14 @@ def infer( "--device", help="Device: cuda, mps, cpu. Auto-detected if not set.", ), + trust_remote_code: bool = typer.Option( + False, + "--trust-remote-code", + help=( + "Allow loading models that ship custom Python via auto_map. " + "Default deny (v0.36.0). Only enable if you trust the source." + ), + ), ): """Run batch inference on a JSONL file of prompts.""" from soup_cli.utils.paths import is_under_cwd @@ -145,13 +153,11 @@ def infer( ) ) - # Load model - console.print( - "[yellow]Warning: loading model with trust_remote_code=True. " - "Only use models you trust.[/]" - ) + # Load model — gate trust_remote_code via the v0.36.0 helper. console.print("[dim]Loading model...[/]") - model_obj, tokenizer = _load_model(str(model_path), base, device) + model_obj, tokenizer = _load_model( + str(model_path), base, device, trust_remote_code, + ) console.print("[green]Model loaded.[/]\n") # Output path containment — defence-in-depth (project policy v0.20.0+). @@ -237,11 +243,21 @@ def _read_prompts(path: Path) -> list[str]: return prompts -def _load_model(model_path: str, base_model: Optional[str], device: str) -> tuple: +def _load_model( + model_path: str, + base_model: Optional[str], + device: str, + trust_remote_code: bool = False, +) -> tuple: """Load a model and tokenizer (reuses diff.py pattern).""" import torch from transformers import AutoModelForCausalLM, AutoTokenizer + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + path = Path(model_path) adapter_config_path = path / "adapter_config.json" is_adapter = adapter_config_path.exists() @@ -260,7 +276,16 @@ def _load_model(model_path: str, base_model: Optional[str], device: str) -> tupl ) raise typer.Exit(1) - tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + probe_target = base_model or model_path + requires = model_requires_trust_remote_code(model_path) or False + trc = resolve_trust_remote_code( + probe_target, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) + + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=trc) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token @@ -269,7 +294,7 @@ def _load_model(model_path: str, base_model: Optional[str], device: str) -> tupl base_obj = AutoModelForCausalLM.from_pretrained( base_model, - trust_remote_code=True, + trust_remote_code=trc, device_map="auto", dtype=torch.float16, ) @@ -277,7 +302,7 @@ def _load_model(model_path: str, base_model: Optional[str], device: str) -> tupl else: model_obj = AutoModelForCausalLM.from_pretrained( model_path, - trust_remote_code=True, + trust_remote_code=trc, device_map="auto", dtype=torch.float16, ) diff --git a/soup_cli/commands/merge.py b/soup_cli/commands/merge.py index 2b97359..04eb480 100644 --- a/soup_cli/commands/merge.py +++ b/soup_cli/commands/merge.py @@ -35,6 +35,14 @@ def merge( "--dtype", help="Data type for the merged model: float16, bfloat16, float32", ), + trust_remote_code: bool = typer.Option( + False, + "--trust-remote-code", + help=( + "Allow loading models that ship custom Python via auto_map. " + "Default deny (v0.36.0). Only enable if you trust the source." + ), + ), ): """Merge a LoRA adapter with its base model into a full model.""" adapter_path = Path(adapter) @@ -93,11 +101,24 @@ def merge( } model_dtype = dtype_map[dtype] + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + + requires = model_requires_trust_remote_code(str(adapter_path)) or False + trc = resolve_trust_remote_code( + base, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) + console.print(f"[dim]Loading base model: {base}...[/]") model = AutoModelForCausalLM.from_pretrained( base, dtype=model_dtype, - trust_remote_code=True, + trust_remote_code=trc, device_map="cpu", ) @@ -112,7 +133,9 @@ def merge( model.save_pretrained(str(output_path)) console.print("[dim]Saving tokenizer...[/]") - tokenizer = AutoTokenizer.from_pretrained(str(adapter_path), trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained( + str(adapter_path), trust_remote_code=trc + ) tokenizer.save_pretrained(str(output_path)) except ImportError as exc: diff --git a/soup_cli/commands/train.py b/soup_cli/commands/train.py index fdd3292..f394c44 100644 --- a/soup_cli/commands/train.py +++ b/soup_cli/commands/train.py @@ -668,10 +668,9 @@ def train( "deepspeed_config": ds_config_path, "fsdp_config": fsdp_kwargs, } - # SFT path threads --trust-remote-code through the wrapper. Other - # trainers still load with trust_remote_code=True at their existing - # call sites; v0.36.x patches will extend the same opt-in to them. - sft_kwargs = dict(trainer_kwargs, trust_remote_code=trust_remote_code) + # v0.40.4 #63 — every transformer-backend trainer now threads + # --trust-remote-code through the wrapper (closes the v0.36.0 Part B gap). + trainer_kwargs = dict(trainer_kwargs, trust_remote_code=trust_remote_code) if cfg.task == "dpo": from soup_cli.trainer.dpo import DPOTrainerWrapper @@ -721,7 +720,7 @@ def train( trainer_wrapper = EmbeddingTrainerWrapper(cfg, **trainer_kwargs) else: - trainer_wrapper = SFTTrainerWrapper(cfg, **sft_kwargs) + trainer_wrapper = SFTTrainerWrapper(cfg, **trainer_kwargs) trainer_wrapper.setup(dataset) # --- HF auto-push callback (Part B of v0.29.0) --- diff --git a/soup_cli/trainer/bco.py b/soup_cli/trainer/bco.py index e40cbfe..d8c7880 100644 --- a/soup_cli/trainer/bco.py +++ b/soup_cli/trainer/bco.py @@ -70,12 +70,26 @@ class BCOTrainerWrapper: report_to: str = "none", deepspeed_config: Optional[str] = None, fsdp_config: Optional[dict] = None, + trust_remote_code: bool = False, ): self.config = config self.device = device self.report_to = report_to self.deepspeed_config = deepspeed_config self.fsdp_config = fsdp_config + self.trust_remote_code = trust_remote_code + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + + requires = model_requires_trust_remote_code(config.base) or False + self._trust_remote_code = resolve_trust_remote_code( + config.base, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) self.model = None self.tokenizer = None self.trainer = None @@ -198,7 +212,9 @@ class BCOTrainerWrapper: from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig console.print(f"[dim]Loading tokenizer: {cfg.base}[/]") - self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True) + self.tokenizer = AutoTokenizer.from_pretrained( + cfg.base, trust_remote_code=self._trust_remote_code + ) if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token @@ -217,7 +233,9 @@ class BCOTrainerWrapper: console.print(f"[dim]Loading model: {cfg.base}[/]") dev_map = "cpu" if self.device == "cpu" else "auto" - model_kwargs = {"trust_remote_code": True, "device_map": dev_map} + model_kwargs = { + "trust_remote_code": self._trust_remote_code, "device_map": dev_map, + } if bnb_config: model_kwargs["quantization_config"] = bnb_config diff --git a/soup_cli/trainer/dpo.py b/soup_cli/trainer/dpo.py index b316468..99ac4ec 100644 --- a/soup_cli/trainer/dpo.py +++ b/soup_cli/trainer/dpo.py @@ -28,16 +28,33 @@ class DPOTrainerWrapper: report_to: str = "none", deepspeed_config: Optional[str] = None, fsdp_config: Optional[dict] = None, + trust_remote_code: bool = False, ): self.config = config self.device = device self.report_to = report_to self.deepspeed_config = deepspeed_config self.fsdp_config = fsdp_config + self.trust_remote_code = trust_remote_code self.model = None self.ref_model = None self.tokenizer = None self.trainer = None + # v0.40.4 #63 — extend v0.36.0 trust_remote_code opt-in to non-SFT + # trainers. Resolve once; raises ValueError if model needs custom + # code but the user did not opt in. + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + + requires = model_requires_trust_remote_code(config.base) or False + self._trust_remote_code = resolve_trust_remote_code( + config.base, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) def setup(self, dataset: dict): """Load model, tokenizer, apply LoRA, create DPO trainer.""" @@ -149,7 +166,9 @@ class DPOTrainerWrapper: from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig console.print(f"[dim]Loading tokenizer: {cfg.base}[/]") - self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True) + self.tokenizer = AutoTokenizer.from_pretrained( + cfg.base, trust_remote_code=self._trust_remote_code + ) if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token @@ -169,7 +188,9 @@ class DPOTrainerWrapper: console.print(f"[dim]Loading model: {cfg.base}[/]") # On CPU, use device_map="cpu" to avoid meta tensors from "auto" dev_map = "cpu" if self.device == "cpu" else "auto" - model_kwargs = {"trust_remote_code": True, "device_map": dev_map} + model_kwargs = { + "trust_remote_code": self._trust_remote_code, "device_map": dev_map, + } if bnb_config: model_kwargs["quantization_config"] = bnb_config diff --git a/soup_cli/trainer/embedding.py b/soup_cli/trainer/embedding.py index 9c117cd..d296f98 100644 --- a/soup_cli/trainer/embedding.py +++ b/soup_cli/trainer/embedding.py @@ -32,12 +32,26 @@ class EmbeddingTrainerWrapper: report_to: str = "none", deepspeed_config: Optional[str] = None, fsdp_config: Optional[dict] = None, + trust_remote_code: bool = False, ): self.config = config self.device = device self.report_to = report_to self.deepspeed_config = deepspeed_config self.fsdp_config = fsdp_config + self.trust_remote_code = trust_remote_code + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + + requires = model_requires_trust_remote_code(config.base) or False + self._trust_remote_code = resolve_trust_remote_code( + config.base, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) self.model = None self.tokenizer = None self.trainer = None @@ -173,7 +187,9 @@ class EmbeddingTrainerWrapper: from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig console.print(f"[dim]Loading tokenizer: {cfg.base}[/]") - self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True) + self.tokenizer = AutoTokenizer.from_pretrained( + cfg.base, trust_remote_code=self._trust_remote_code + ) if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token @@ -192,7 +208,9 @@ class EmbeddingTrainerWrapper: console.print(f"[dim]Loading model: {cfg.base}[/]") dev_map = "cpu" if self.device == "cpu" else "auto" - model_kwargs = {"trust_remote_code": True, "device_map": dev_map} + model_kwargs = { + "trust_remote_code": self._trust_remote_code, "device_map": dev_map, + } if bnb_config: model_kwargs["quantization_config"] = bnb_config diff --git a/soup_cli/trainer/grpo.py b/soup_cli/trainer/grpo.py index ac31857..a1f2f78 100644 --- a/soup_cli/trainer/grpo.py +++ b/soup_cli/trainer/grpo.py @@ -30,12 +30,26 @@ class GRPOTrainerWrapper: report_to: str = "none", deepspeed_config: Optional[str] = None, fsdp_config: Optional[dict] = None, + trust_remote_code: bool = False, ): self.config = config self.device = device self.report_to = report_to self.deepspeed_config = deepspeed_config self.fsdp_config = fsdp_config + self.trust_remote_code = trust_remote_code + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + + requires = model_requires_trust_remote_code(config.base) or False + self._trust_remote_code = resolve_trust_remote_code( + config.base, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) self.model = None self.tokenizer = None self.trainer = None @@ -201,7 +215,9 @@ class GRPOTrainerWrapper: from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig console.print(f"[dim]Loading tokenizer: {cfg.base}[/]") - self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True) + self.tokenizer = AutoTokenizer.from_pretrained( + cfg.base, trust_remote_code=self._trust_remote_code + ) if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token @@ -221,7 +237,9 @@ class GRPOTrainerWrapper: console.print(f"[dim]Loading model: {cfg.base}[/]") # On CPU, use device_map="cpu" to avoid meta tensors from "auto" dev_map = "cpu" if self.device == "cpu" else "auto" - model_kwargs = {"trust_remote_code": True, "device_map": dev_map} + model_kwargs = { + "trust_remote_code": self._trust_remote_code, "device_map": dev_map, + } if bnb_config: model_kwargs["quantization_config"] = bnb_config diff --git a/soup_cli/trainer/ipo.py b/soup_cli/trainer/ipo.py index 919f2bb..a98f2a8 100644 --- a/soup_cli/trainer/ipo.py +++ b/soup_cli/trainer/ipo.py @@ -33,12 +33,26 @@ class IPOTrainerWrapper: report_to: str = "none", deepspeed_config: Optional[str] = None, fsdp_config: Optional[dict] = None, + trust_remote_code: bool = False, ): self.config = config self.device = device self.report_to = report_to self.deepspeed_config = deepspeed_config self.fsdp_config = fsdp_config + self.trust_remote_code = trust_remote_code + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + + requires = model_requires_trust_remote_code(config.base) or False + self._trust_remote_code = resolve_trust_remote_code( + config.base, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) self.model = None self.tokenizer = None self.trainer = None @@ -152,7 +166,9 @@ class IPOTrainerWrapper: from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig console.print(f"[dim]Loading tokenizer: {cfg.base}[/]") - self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True) + self.tokenizer = AutoTokenizer.from_pretrained( + cfg.base, trust_remote_code=self._trust_remote_code + ) if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token @@ -171,7 +187,9 @@ class IPOTrainerWrapper: console.print(f"[dim]Loading model: {cfg.base}[/]") dev_map = "cpu" if self.device == "cpu" else "auto" - model_kwargs = {"trust_remote_code": True, "device_map": dev_map} + model_kwargs = { + "trust_remote_code": self._trust_remote_code, "device_map": dev_map, + } if bnb_config: model_kwargs["quantization_config"] = bnb_config diff --git a/soup_cli/trainer/kto.py b/soup_cli/trainer/kto.py index e3ee3c8..97819cd 100644 --- a/soup_cli/trainer/kto.py +++ b/soup_cli/trainer/kto.py @@ -29,12 +29,26 @@ class KTOTrainerWrapper: report_to: str = "none", deepspeed_config: Optional[str] = None, fsdp_config: Optional[dict] = None, + trust_remote_code: bool = False, ): self.config = config self.device = device self.report_to = report_to self.deepspeed_config = deepspeed_config self.fsdp_config = fsdp_config + self.trust_remote_code = trust_remote_code + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + + requires = model_requires_trust_remote_code(config.base) or False + self._trust_remote_code = resolve_trust_remote_code( + config.base, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) self.model = None self.tokenizer = None self.trainer = None @@ -148,7 +162,9 @@ class KTOTrainerWrapper: from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig console.print(f"[dim]Loading tokenizer: {cfg.base}[/]") - self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True) + self.tokenizer = AutoTokenizer.from_pretrained( + cfg.base, trust_remote_code=self._trust_remote_code + ) if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token @@ -168,7 +184,9 @@ class KTOTrainerWrapper: console.print(f"[dim]Loading model: {cfg.base}[/]") # On CPU, use device_map="cpu" to avoid meta tensors from "auto" dev_map = "cpu" if self.device == "cpu" else "auto" - model_kwargs = {"trust_remote_code": True, "device_map": dev_map} + model_kwargs = { + "trust_remote_code": self._trust_remote_code, "device_map": dev_map, + } if bnb_config: model_kwargs["quantization_config"] = bnb_config diff --git a/soup_cli/trainer/orpo.py b/soup_cli/trainer/orpo.py index 1d54741..5c641b2 100644 --- a/soup_cli/trainer/orpo.py +++ b/soup_cli/trainer/orpo.py @@ -32,12 +32,26 @@ class ORPOTrainerWrapper: report_to: str = "none", deepspeed_config: Optional[str] = None, fsdp_config: Optional[dict] = None, + trust_remote_code: bool = False, ): self.config = config self.device = device self.report_to = report_to self.deepspeed_config = deepspeed_config self.fsdp_config = fsdp_config + self.trust_remote_code = trust_remote_code + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + + requires = model_requires_trust_remote_code(config.base) or False + self._trust_remote_code = resolve_trust_remote_code( + config.base, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) self.model = None self.tokenizer = None self.trainer = None @@ -149,7 +163,9 @@ class ORPOTrainerWrapper: from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig console.print(f"[dim]Loading tokenizer: {cfg.base}[/]") - self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True) + self.tokenizer = AutoTokenizer.from_pretrained( + cfg.base, trust_remote_code=self._trust_remote_code + ) if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token @@ -168,7 +184,9 @@ class ORPOTrainerWrapper: console.print(f"[dim]Loading model: {cfg.base}[/]") dev_map = "cpu" if self.device == "cpu" else "auto" - model_kwargs = {"trust_remote_code": True, "device_map": dev_map} + model_kwargs = { + "trust_remote_code": self._trust_remote_code, "device_map": dev_map, + } if bnb_config: model_kwargs["quantization_config"] = bnb_config diff --git a/soup_cli/trainer/ppo.py b/soup_cli/trainer/ppo.py index 05324d3..491febd 100644 --- a/soup_cli/trainer/ppo.py +++ b/soup_cli/trainer/ppo.py @@ -40,12 +40,26 @@ class PPOTrainerWrapper: report_to: str = "none", deepspeed_config: Optional[str] = None, fsdp_config: Optional[dict] = None, + trust_remote_code: bool = False, ): self.config = config self.device = device self.report_to = report_to self.deepspeed_config = deepspeed_config self.fsdp_config = fsdp_config + self.trust_remote_code = trust_remote_code + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + + requires = model_requires_trust_remote_code(config.base) or False + self._trust_remote_code = resolve_trust_remote_code( + config.base, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) self.model = None self.tokenizer = None self.trainer = None @@ -260,7 +274,9 @@ class PPOTrainerWrapper: # If a reward_model path is configured, load it if tcfg.reward_model: - return _load_reward_model(tcfg.reward_model, self.device) + return _load_reward_model( + tcfg.reward_model, self.device, self.trust_remote_code, + ) # Fallback: create a sequence classification model from the base model from transformers import AutoModelForSequenceClassification @@ -271,7 +287,7 @@ class PPOTrainerWrapper: ) reward_model = AutoModelForSequenceClassification.from_pretrained( cfg.base, - trust_remote_code=True, + trust_remote_code=self._trust_remote_code, num_labels=1, device_map="auto" if self.device != "cpu" else None, ) @@ -289,7 +305,7 @@ class PPOTrainerWrapper: console.print(f"[dim]Creating value model from: {cfg.base}[/]") value_model = AutoModelForSequenceClassification.from_pretrained( cfg.base, - trust_remote_code=True, + trust_remote_code=self._trust_remote_code, num_labels=1, device_map="auto" if self.device != "cpu" else None, ) @@ -300,7 +316,7 @@ class PPOTrainerWrapper: # Reward model (pre-trained classifier) if tcfg.reward_model: self.reward_model_instance = _load_reward_model( - tcfg.reward_model, self.device, + tcfg.reward_model, self.device, self.trust_remote_code, ) console.print(f"[green]Reward model loaded:[/] {tcfg.reward_model}") @@ -325,7 +341,9 @@ class PPOTrainerWrapper: from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig console.print(f"[dim]Loading tokenizer: {cfg.base}[/]") - self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True) + self.tokenizer = AutoTokenizer.from_pretrained( + cfg.base, trust_remote_code=self._trust_remote_code + ) if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token @@ -345,7 +363,9 @@ class PPOTrainerWrapper: console.print(f"[dim]Loading model: {cfg.base}[/]") # On CPU, use device_map="cpu" to avoid meta tensors from "auto" dev_map = "cpu" if self.device == "cpu" else "auto" - model_kwargs = {"trust_remote_code": True, "device_map": dev_map} + model_kwargs = { + "trust_remote_code": self._trust_remote_code, "device_map": dev_map, + } if bnb_config: model_kwargs["quantization_config"] = bnb_config @@ -643,7 +663,11 @@ def _import_ppo_classes(): return PPOTrainer, PPOConfig, False -def _load_reward_model(model_path: str, device: str = "cuda"): +def _load_reward_model( + model_path: str, + device: str = "cuda", + trust_remote_code: bool = False, +): """Load a pre-trained reward model for PPO scoring. Reward models are typically AutoModelForSequenceClassification that output @@ -651,11 +675,23 @@ def _load_reward_model(model_path: str, device: str = "cuda"): """ from transformers import AutoModelForSequenceClassification + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + + requires = model_requires_trust_remote_code(model_path) or False + resolved = resolve_trust_remote_code( + model_path, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) console.print(f"[dim]Loading reward model: {model_path}[/]") dev_map = "cpu" if device == "cpu" else "auto" reward_model = AutoModelForSequenceClassification.from_pretrained( model_path, - trust_remote_code=True, + trust_remote_code=resolved, device_map=dev_map, num_labels=1, ) diff --git a/soup_cli/trainer/preference.py b/soup_cli/trainer/preference.py index e4e84c5..fe0a129 100644 --- a/soup_cli/trainer/preference.py +++ b/soup_cli/trainer/preference.py @@ -85,12 +85,14 @@ class PreferenceTrainerWrapper: report_to: str = "none", deepspeed_config: Optional[str] = None, fsdp_config: Optional[dict] = None, + trust_remote_code: bool = False, ): self.config = config self.device = device self.report_to = report_to self.deepspeed_config = deepspeed_config self.fsdp_config = fsdp_config + self.trust_remote_code = trust_remote_code self._inner = None def _build_inner(self): @@ -106,6 +108,7 @@ class PreferenceTrainerWrapper: "report_to": self.report_to, "deepspeed_config": self.deepspeed_config, "fsdp_config": self.fsdp_config, + "trust_remote_code": self.trust_remote_code, } if loss == "dpo": from soup_cli.trainer.dpo import DPOTrainerWrapper @@ -155,6 +158,7 @@ class PreferenceTrainerWrapper: "report_to": self.report_to, "deepspeed_config": self.deepspeed_config, "fsdp_config": self.fsdp_config, + "trust_remote_code": self.trust_remote_code, } if primary == "dpo": from soup_cli.trainer.dpo import DPOTrainerWrapper as PrimaryWrapper diff --git a/soup_cli/trainer/pretrain.py b/soup_cli/trainer/pretrain.py index 3fbcdb1..f13758a 100644 --- a/soup_cli/trainer/pretrain.py +++ b/soup_cli/trainer/pretrain.py @@ -28,12 +28,26 @@ class PretrainTrainerWrapper: report_to: str = "none", deepspeed_config: Optional[str] = None, fsdp_config: Optional[dict] = None, + trust_remote_code: bool = False, ): self.config = config self.device = device self.report_to = report_to self.deepspeed_config = deepspeed_config self.fsdp_config = fsdp_config + self.trust_remote_code = trust_remote_code + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + + requires = model_requires_trust_remote_code(config.base) or False + self._trust_remote_code = resolve_trust_remote_code( + config.base, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) self.model = None self.tokenizer = None self.trainer = None @@ -169,14 +183,34 @@ class PretrainTrainerWrapper: trainer_kwargs["packing"] = True console.print("[green]Sample packing enabled[/]") - # v0.40.3: #65 multipack live wiring DEFERRED to v0.40.4 — see - # sft.py for the rationale. Same advisory. - if getattr(tcfg, "multipack", False): - console.print( - "[yellow]multipack: live HF Trainer wiring is deferred to " - "v0.40.4. Falling back to the standard sampler.[/]" + # v0.40.4 #65 — multipack live wiring (mirrors sft.py). + use_multipack = bool(getattr(tcfg, "multipack", False)) + if use_multipack: + from soup_cli.utils.multipack_sampler import ( + validate_multipack_architecture, ) - self.trainer = SFTTrainer(**trainer_kwargs) + from soup_cli.utils.multipack_trainer import ( + attach_multipack_state, + detect_arch_name, + lengths_from_dataset, + make_multipack_trainer_class, + ) + + arch = detect_arch_name(self.model) + if arch: + validate_multipack_architecture(arch) + trainer_cls = make_multipack_trainer_class(SFTTrainer) + self.trainer = trainer_cls(**trainer_kwargs) + attach_multipack_state( + self.trainer, + lengths=lengths_from_dataset(train_ds), + max_seq_len=cfg.data.max_length, + batch_size=batch_size, + seed=getattr(tcfg, "seed", 0) or 0, + ) + console.print("[green]Multipack FFD bin-packing sampler enabled[/]") + else: + self.trainer = SFTTrainer(**trainer_kwargs) self._output_dir = str(output_dir) @@ -188,7 +222,9 @@ class PretrainTrainerWrapper: from soup_cli.utils.moe import detect_moe_model, get_moe_target_modules console.print(f"[dim]Loading tokenizer: {cfg.base}[/]") - self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True) + self.tokenizer = AutoTokenizer.from_pretrained( + cfg.base, trust_remote_code=self._trust_remote_code + ) if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token @@ -208,7 +244,9 @@ class PretrainTrainerWrapper: console.print(f"[dim]Loading model: {cfg.base}[/]") # On CPU, use device_map="cpu" to avoid meta tensors from "auto" dev_map = "cpu" if self.device == "cpu" else "auto" - model_kwargs = {"trust_remote_code": True, "device_map": dev_map} + model_kwargs = { + "trust_remote_code": self._trust_remote_code, "device_map": dev_map, + } if bnb_config: model_kwargs["quantization_config"] = bnb_config diff --git a/soup_cli/trainer/reward_model.py b/soup_cli/trainer/reward_model.py index 5502176..ac5a64b 100644 --- a/soup_cli/trainer/reward_model.py +++ b/soup_cli/trainer/reward_model.py @@ -36,12 +36,26 @@ class RewardModelTrainerWrapper: report_to: str = "none", deepspeed_config: Optional[str] = None, fsdp_config: Optional[dict] = None, + trust_remote_code: bool = False, ): self.config = config self.device = device self.report_to = report_to self.deepspeed_config = deepspeed_config self.fsdp_config = fsdp_config + self.trust_remote_code = trust_remote_code + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + + requires = model_requires_trust_remote_code(config.base) or False + self._trust_remote_code = resolve_trust_remote_code( + config.base, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) self.model = None self.tokenizer = None self.trainer = None @@ -154,7 +168,9 @@ class RewardModelTrainerWrapper: ) console.print(f"[dim]Loading tokenizer: {cfg.base}[/]") - self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True) + self.tokenizer = AutoTokenizer.from_pretrained( + cfg.base, trust_remote_code=self._trust_remote_code + ) if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token @@ -173,7 +189,11 @@ class RewardModelTrainerWrapper: console.print(f"[dim]Loading reward model: {cfg.base}[/]") dev_map = "cpu" if self.device == "cpu" else "auto" - model_kwargs = {"trust_remote_code": True, "device_map": dev_map, "num_labels": 1} + model_kwargs = { + "trust_remote_code": self._trust_remote_code, + "device_map": dev_map, + "num_labels": 1, + } if bnb_config: model_kwargs["quantization_config"] = bnb_config diff --git a/soup_cli/trainer/simpo.py b/soup_cli/trainer/simpo.py index 68d0745..99f4093 100644 --- a/soup_cli/trainer/simpo.py +++ b/soup_cli/trainer/simpo.py @@ -33,12 +33,26 @@ class SimPOTrainerWrapper: report_to: str = "none", deepspeed_config: Optional[str] = None, fsdp_config: Optional[dict] = None, + trust_remote_code: bool = False, ): self.config = config self.device = device self.report_to = report_to self.deepspeed_config = deepspeed_config self.fsdp_config = fsdp_config + self.trust_remote_code = trust_remote_code + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + + requires = model_requires_trust_remote_code(config.base) or False + self._trust_remote_code = resolve_trust_remote_code( + config.base, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) self.model = None self.tokenizer = None self.trainer = None @@ -152,7 +166,9 @@ class SimPOTrainerWrapper: from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig console.print(f"[dim]Loading tokenizer: {cfg.base}[/]") - self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True) + self.tokenizer = AutoTokenizer.from_pretrained( + cfg.base, trust_remote_code=self._trust_remote_code + ) if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token @@ -171,7 +187,9 @@ class SimPOTrainerWrapper: console.print(f"[dim]Loading model: {cfg.base}[/]") dev_map = "cpu" if self.device == "cpu" else "auto" - model_kwargs = {"trust_remote_code": True, "device_map": dev_map} + model_kwargs = { + "trust_remote_code": self._trust_remote_code, "device_map": dev_map, + } if bnb_config: model_kwargs["quantization_config"] = bnb_config diff --git a/tests/test_v0404_part_a.py b/tests/test_v0404_part_a.py new file mode 100644 index 0000000..b27d308 --- /dev/null +++ b/tests/test_v0404_part_a.py @@ -0,0 +1,353 @@ +"""Tests for v0.40.4 Part A — #63 trust_remote_code multi-trainer. + +Extends the v0.36.0 ``--trust-remote-code`` opt-in to all 10 non-SFT +trainer wrappers (DPO / GRPO / KTO / ORPO / SimPO / IPO / PPO / +RewardModel / Pretrain / Embedding / BCO / Preference) and 5 standalone +commands (diff / export / merge / infer / generate). + +Strategy: each trainer wrapper now accepts ``trust_remote_code: bool`` +on ``__init__`` and resolves it via ``resolve_trust_remote_code`` from +``soup_cli.utils.trust_remote``. We assert source-level invariants +(closes the v0.36.0 known-gap family) plus a live-call test that +exercises the resolver path on each wrapper. +""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import patch + +import pytest +from typer.testing import CliRunner + +# All 10 trainer modules + their wrapper class names. +# Direct trainers that resolve trust_remote_code in __init__. +TRAINER_TARGETS = [ + ("dpo", "DPOTrainerWrapper"), + ("grpo", "GRPOTrainerWrapper"), + ("kto", "KTOTrainerWrapper"), + ("orpo", "ORPOTrainerWrapper"), + ("simpo", "SimPOTrainerWrapper"), + ("ipo", "IPOTrainerWrapper"), + ("ppo", "PPOTrainerWrapper"), + ("reward_model", "RewardModelTrainerWrapper"), + ("pretrain", "PretrainTrainerWrapper"), + ("embedding", "EmbeddingTrainerWrapper"), + ("bco", "BCOTrainerWrapper"), +] +# PreferenceTrainerWrapper is a pure dispatcher — it does not resolve +# the flag itself; instead it forwards via kwargs to the inner wrapper. +DISPATCHER_TARGETS = [("preference", "PreferenceTrainerWrapper")] +ALL_TRAINER_TARGETS = TRAINER_TARGETS + DISPATCHER_TARGETS + + +def _minimal_cfg(base: str = "meta-llama/Llama-3.2-1B"): + """Minimal SoupConfig stub with required fields for trainer ctor.""" + from soup_cli.config.schema import SoupConfig + + # The ``preference`` task has a cross-validator on preference_loss — + # for non-preference trainers we use sft as the task. + return SoupConfig( + base=base, + task="sft", + data={"train": "fake.jsonl"}, + training={"epochs": 1, "lr": 1e-4}, + ) + + +class TestTrainerCtorAcceptsTrustRemoteCode: + """Every non-SFT trainer wrapper now accepts trust_remote_code in __init__.""" + + @pytest.mark.parametrize("module,cls_name", TRAINER_TARGETS) + def test_init_accepts_kwarg(self, module: str, cls_name: str): + mod = __import__(f"soup_cli.trainer.{module}", fromlist=[cls_name]) + cls = getattr(mod, cls_name) + cfg = _minimal_cfg() + # Must not raise on a meta-llama/* base (allowlisted, opt-in safe). + instance = cls(cfg, device="cpu", trust_remote_code=False) + assert hasattr(instance, "trust_remote_code") + assert instance.trust_remote_code is False + # Resolver result is cached — defaults to False on safe org. + assert hasattr(instance, "_trust_remote_code") + assert instance._trust_remote_code is False + + @pytest.mark.parametrize("module,cls_name", TRAINER_TARGETS) + def test_init_with_opt_in(self, module: str, cls_name: str): + mod = __import__(f"soup_cli.trainer.{module}", fromlist=[cls_name]) + cls = getattr(mod, cls_name) + cfg = _minimal_cfg() + instance = cls(cfg, device="cpu", trust_remote_code=True) + # Opt-in propagates to the resolved value. + assert instance._trust_remote_code is True + + +class TestUnknownOrgRejectionWithoutOptIn: + """When ``base`` points at a local dir whose config.json has auto_map, + construction must fail without --trust-remote-code (closes the + silent-execute-on-load loophole that v0.36.0 only fixed for SFT). + """ + + @pytest.mark.parametrize("module,cls_name", TRAINER_TARGETS) + def test_local_auto_map_raises(self, tmp_path, module: str, cls_name: str): + # Write a fake local checkpoint with auto_map (the marker that + # triggers ``model_requires_trust_remote_code``). + local = tmp_path / "fake_model" + local.mkdir() + (local / "config.json").write_text( + '{"model_type": "fake", "auto_map": {"AutoConfig": "x.MyConfig"}}', + encoding="utf-8", + ) + mod = __import__(f"soup_cli.trainer.{module}", fromlist=[cls_name]) + cls = getattr(mod, cls_name) + cfg = _minimal_cfg(base=str(local)) + with pytest.raises(ValueError, match="trust_remote_code"): + cls(cfg, device="cpu", trust_remote_code=False) + + +class TestSourceLevelInvariants: + """The v0.36.0 known-gap family is now closed. No trainer file should + contain ``trust_remote_code=True`` literally — every site reads from + the resolved attribute or a parameter. + """ + + @pytest.mark.parametrize("module,_cls", ALL_TRAINER_TARGETS) + def test_no_hardcoded_true_in_trainer(self, module: str, _cls: str): + text = Path(f"soup_cli/trainer/{module}.py").read_text(encoding="utf-8") + # Allow text only in comments / docstrings — assert exact-arg form + # ``trust_remote_code=True`` is absent. + assert "trust_remote_code=True" not in text, ( + f"{module}.py still hardcodes trust_remote_code=True — must " + f"thread the v0.36.0 helper instead" + ) + + def test_no_hardcoded_true_in_sft(self): + # SFT was already cleaned in v0.36.0 — regression guard. + text = Path("soup_cli/trainer/sft.py").read_text(encoding="utf-8") + assert "trust_remote_code=True" not in text + + +class TestCommandFlagsExist: + """The 5 commands now expose ``--trust-remote-code`` Typer options.""" + + @pytest.mark.parametrize("cmd", ["diff", "export", "merge", "infer", "generate"]) + def test_cli_help_lists_flag(self, cmd: str): + from soup_cli.cli import app + + runner = CliRunner() + # Strip ANSI to handle Rich line-wrapping on narrow terminals + # (matches the v0.36.0 test_trust_remote_code.py helper pattern). + # When the command is nested under "data" (generate), the flag + # might not appear at the top level; check the nested help too. + result = runner.invoke(app, [cmd, "--help"]) + out = result.output + if "--trust-remote-code" not in out: + # Fallback: ``data generate`` is nested. + result = runner.invoke(app, ["data", cmd, "--help"]) + out = result.output + assert "--trust-remote-code" in out, ( + f"--trust-remote-code missing from `soup {cmd} --help`" + ) + + +class TestTrainPyWiresFlagToAllTrainers: + """The v0.36.0 ``sft_kwargs`` split is now removed — every trainer + receives ``trust_remote_code`` from the same kwargs dict. + """ + + def test_no_sft_kwargs_split(self): + text = Path("soup_cli/commands/train.py").read_text(encoding="utf-8") + # The v0.36.0 vintage line ``sft_kwargs = dict(trainer_kwargs, ...`` + # no longer needs a separate dict — assert it's gone. + assert "sft_kwargs = dict(trainer_kwargs," not in text + # And every wrapper instantiation reads from the unified kwargs. + assert "trust_remote_code=trust_remote_code" in text + + +class TestPpoLoadRewardModelHelperGated: + """``_load_reward_model`` is module-level in ppo.py and used to be + a hard-coded ``trust_remote_code=True`` site. v0.40.4 threads the + user opt-in through it. + """ + + def test_helper_signature_accepts_flag(self): + import inspect + + from soup_cli.trainer.ppo import _load_reward_model + + sig = inspect.signature(_load_reward_model) + assert "trust_remote_code" in sig.parameters + assert sig.parameters["trust_remote_code"].default is False + + def test_helper_rejects_unknown_org_local_auto_map(self, tmp_path): + from soup_cli.trainer.ppo import _load_reward_model + + local = tmp_path / "fake_rm" + local.mkdir() + (local / "config.json").write_text( + '{"model_type": "fake", "auto_map": {"AutoConfig": "x.MyConfig"}}', + encoding="utf-8", + ) + with pytest.raises(ValueError, match="trust_remote_code"): + _load_reward_model(str(local), device="cpu", trust_remote_code=False) + + +class TestExportHelperSignatures: + """``_merge_adapter`` and the four ``_export_*`` helpers in export.py + now thread ``trust_remote_code`` from the Typer flag. + """ + + def test_merge_adapter_signature(self): + import inspect + + from soup_cli.commands.export import _merge_adapter + + params = inspect.signature(_merge_adapter).parameters + assert "trust_remote_code" in params + assert params["trust_remote_code"].default is False + + @pytest.mark.parametrize( + "fn_name", ["_export_onnx", "_export_tensorrt", "_export_awq", "_export_gptq"], + ) + def test_export_helper_signature(self, fn_name: str): + import inspect + + from soup_cli.commands import export as export_mod + + fn = getattr(export_mod, fn_name) + params = inspect.signature(fn).parameters + assert "trust_remote_code" in params, ( + f"{fn_name} did not get the v0.40.4 flag" + ) + assert params["trust_remote_code"].default is False + + +class TestPreferenceForwardsToInner: + """``PreferenceTrainerWrapper`` forwards ``trust_remote_code`` to the + inner DPO / SimPO / ORPO / IPO / BCO wrapper. + """ + + def test_preference_inner_kwargs_include_flag(self): + text = Path("soup_cli/trainer/preference.py").read_text(encoding="utf-8") + # Both _build_inner and _build_multi_objective build a kwargs dict; + # both must include the flag forwarding. + assert text.count('"trust_remote_code": self.trust_remote_code') >= 2 + + +class TestBcoAlreadyOnHelper: + """Re-verify that BCO (v0.40.0 Part A) is now on the v0.36.0 helper too.""" + + def test_bco_uses_resolver(self): + text = Path("soup_cli/trainer/bco.py").read_text(encoding="utf-8") + assert "resolve_trust_remote_code" in text + assert "self._trust_remote_code" in text + + +class TestResolverLoadedLazily: + """The resolver import inside ``__init__`` keeps ``import soup_cli.trainer.`` + cheap — heavy deps are still gated behind ``setup()``. + """ + + @patch("transformers.AutoTokenizer.from_pretrained") + @patch("transformers.AutoModelForCausalLM.from_pretrained") + def test_init_does_not_call_from_pretrained( + self, mock_model, mock_tok, + ): + from soup_cli.trainer.dpo import DPOTrainerWrapper + + cfg = _minimal_cfg() + DPOTrainerWrapper(cfg, device="cpu", trust_remote_code=False) + # Heavy loaders are still deferred to setup(). + assert mock_model.call_count == 0 + assert mock_tok.call_count == 0 + + +class TestCommandHelpersGateRemoteCode: + """v0.40.4 H1 — negative tests for each of the 5 commands' load-helper + paths. Each helper resolves trust_remote_code BEFORE calling + ``from_pretrained``, so a local checkpoint with ``auto_map`` must + raise ``ValueError`` when ``trust_remote_code=False``. + """ + + def _make_local_auto_map(self, tmp_path): + local = tmp_path / "fake_model" + local.mkdir() + (local / "config.json").write_text( + '{"model_type": "fake", "auto_map": {"AutoConfig": "x.MyConfig"}}', + encoding="utf-8", + ) + return local + + def test_diff_load_model_rejects(self, tmp_path): + from soup_cli.commands.diff import _load_model + + local = self._make_local_auto_map(tmp_path) + with pytest.raises(ValueError, match="trust_remote_code"): + _load_model(str(local), None, "cpu", trust_remote_code=False) + + def test_infer_load_model_rejects(self, tmp_path): + from soup_cli.commands.infer import _load_model + + local = self._make_local_auto_map(tmp_path) + with pytest.raises(ValueError, match="trust_remote_code"): + _load_model(str(local), None, "cpu", trust_remote_code=False) + + def test_export_merge_adapter_rejects(self, tmp_path): + from soup_cli.commands.export import _merge_adapter + + local = self._make_local_auto_map(tmp_path) + with pytest.raises(ValueError, match="trust_remote_code"): + _merge_adapter( + str(local), + str(local), + str(tmp_path / "out"), + trust_remote_code=False, + ) + + def test_generate_local_rejects(self, tmp_path): + from soup_cli.commands.generate import _generate_local + + local = self._make_local_auto_map(tmp_path) + with pytest.raises(ValueError, match="trust_remote_code"): + _generate_local( + prompt="x", + count=1, + fmt="alpaca", + model_name=str(local), + temperature=0.7, + seed_examples=[], + trust_remote_code=False, + ) + + +class TestPreferenceDispatcherLiveForwardsRejection: + """H2 — beyond the source-grep test, verify that constructing a + PreferenceTrainerWrapper that targets a local auto_map model + eventually raises (the resolver fires inside the inner wrapper at + setup() time). + """ + + def test_inner_wrapper_raises_on_setup(self, tmp_path): + local = tmp_path / "fake_model" + local.mkdir() + (local / "config.json").write_text( + '{"model_type": "fake", "auto_map": {"AutoConfig": "x.MyConfig"}}', + encoding="utf-8", + ) + + from soup_cli.config.schema import SoupConfig + from soup_cli.trainer.preference import PreferenceTrainerWrapper + + cfg = SoupConfig( + base=str(local), + task="preference", + data={"train": "fake.jsonl"}, + training={"epochs": 1, "lr": 1e-4, "preference_loss": "dpo"}, + ) + # The dispatcher itself constructs without resolving (forwarding + # is deferred until ``_build_inner`` is called from ``setup``). + wrapper = PreferenceTrainerWrapper( + cfg, device="cpu", trust_remote_code=False, + ) + # _build_inner constructs DPOTrainerWrapper which DOES resolve. + with pytest.raises(ValueError, match="trust_remote_code"): + wrapper._build_inner()