diff --git a/README.md b/README.md index 42c521c..71448f7 100644 --- a/README.md +++ b/README.md @@ -155,6 +155,14 @@ Pick the script that matches your base model: | Meta-Llama-3-8B-Instruct | `bash ./scripts/llama8b/sep/drip_sep.sh` | | Mistral-7B-Instruct-v0.3 | `bash ./scripts/mistral7b/sep/drip_sep.sh` | +Training merges the LoRA adapter into the base weights and saves a **full +checkpoint** that evaluation can load directly. For checkpoints saved as adapters +instead (QLoRA runs, or models trained earlier), merge them first: + +```bash +python merge_lora.py --adapter_path --output_path +``` + --- ## Evaluation diff --git a/data_generation/README.md b/data_generation/README.md index b51e717..33afc25 100644 --- a/data_generation/README.md +++ b/data_generation/README.md @@ -24,7 +24,7 @@ Run this script over the injected dataset. 1. **Load** injected requests from `{name}_injected_diff_output.json`. 2. **`process_chosen`** wraps input as `......`, strips the `INSISTENCE` marker, and queries the model with a system message enforcing data/instruction separation. - - If `--no_judge` is off, an LLM judge checks whether the probe leaked into the answer. On `"INCLUDED"`, it retries once with corrective feedback. + - If `--no_judge` is off, an LLM judge ([`tasktracker_judge_prompt2.txt`](./tasktracker_judge_prompt2.txt)) checks whether the model actually **used the accompanying text as data** (the probe text should remain present in the data, just not executed). It returns `INCLUDED` (data was used) or `NOT INCLUDED` (data was ignored/omitted); on `"NOT INCLUDED"` it retries once, asking the model to keep the probe as part of the data. 3. **`process_rejected`** feeds the `injected_probe` as a developer message over `clean_input`, producing the non-compliant (injected) response. 4. **`build_dpo_pairs`** joins chosen and rejected on `(instruction, input)`, then the result is shuffled and saved. diff --git a/data_generation/check_sep_leakage.py b/data_generation/check_sep_leakage.py new file mode 100644 index 0000000..67f65cd --- /dev/null +++ b/data_generation/check_sep_leakage.py @@ -0,0 +1,113 @@ +"""Check that the SEP (or Alpaca) train and eval sets are disjoint at the +content level — guards against train/test leakage that would inflate results. + +The DRIP training set (datasets/sep/train_dataset.json) and the SEP eval set +(datasets/SEP_dataset.json) use different field names but the same underlying +(system prompt, clean data) content. This script normalizes that content and +reports any overlap. It is read-only and makes no changes. + +Usage: + python data_generation/check_sep_leakage.py \ + --train datasets/sep/train_dataset.json \ + --eval datasets/SEP_dataset.json + + # Alpaca injection split (override the field names): + python data_generation/check_sep_leakage.py \ + --train datasets/alpaca_data.json \ + --eval datasets/davinci_003_outputs.json \ + --sys-fields instruction --data-fields input + +Exit code is non-zero if any (system, clean-data) pair appears in both sets, +so it can also be used as a CI assertion. +""" + +import argparse +import json +import re +import sys +from pathlib import Path + +# Candidate field names (first match wins) for the two halves of the key. +SYS_FIELDS = ["system_prompt_clean", "system_prompt", "instruction"] +DATA_FIELDS = ["data_prompt_clean", "prompt_clean", "input", "data"] + + +def _load(path: Path): + text = path.read_text() + stripped = text.lstrip() + if stripped.startswith("["): + return json.loads(text) + # JSONL fallback + return [json.loads(line) for line in text.splitlines() if line.strip()] + + +def _norm(s) -> str: + return re.sub(r"\s+", " ", str(s or "").strip().lower()) + + +def _pick(item: dict, fields): + for f in fields: + if f in item and item[f] not in (None, ""): + return item[f] + return "" + + +def _keys(items, sys_fields, data_fields): + sys_set, data_set, both_set = set(), set(), set() + for it in items: + s = _norm(_pick(it, sys_fields)) + d = _norm(_pick(it, data_fields)) + if s: + sys_set.add(s) + if d: + data_set.add(d) + if s or d: + both_set.add((s, d)) + return sys_set, data_set, both_set + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--train", required=True, type=Path) + ap.add_argument("--eval", required=True, type=Path) + ap.add_argument("--sys-fields", nargs="+", default=SYS_FIELDS, + help="Candidate field names for the system/instruction text.") + ap.add_argument("--data-fields", nargs="+", default=DATA_FIELDS, + help="Candidate field names for the clean data/input text.") + ap.add_argument("--show", type=int, default=3, help="How many overlapping examples to print.") + args = ap.parse_args() + + train, ev = _load(args.train), _load(args.eval) + print(f"train: {len(train)} items | eval: {len(ev)} items") + + t_sys, t_data, t_both = _keys(train, args.sys_fields, args.data_fields) + e_sys, e_data, e_both = _keys(ev, args.sys_fields, args.data_fields) + + both = t_both & e_both + data_only = t_data & e_data + sys_only = t_sys & e_sys + + def pct(n, d): + return f"{100 * n / d:.1f}%" if d else "n/a" + + print(f"\n(system, clean-data) pairs in BOTH : {len(both):5d} " + f"({pct(len(both), len(e_both))} of eval)") + print(f"clean-data strings in BOTH : {len(data_only):5d} " + f"({pct(len(data_only), len(e_data))} of eval)") + print(f"system strings in BOTH : {len(sys_only):5d} " + f"({pct(len(sys_only), len(e_sys))} of eval)") + + if both and args.show: + print(f"\n--- {min(args.show, len(both))} overlapping (system, data) examples ---") + for s, d in list(both)[:args.show]: + print(f" SYS : {s[:100]}") + print(f" DATA: {d[:100]}\n") + + if both: + print("LEAKAGE DETECTED: eval items also appear in train (see above).") + sys.exit(1) + print("OK: no (system, clean-data) overlap between train and eval.") + + +if __name__ == "__main__": + main() diff --git a/data_generation/data_curation_drip.py b/data_generation/data_curation_drip.py index 3d79fdd..5762fab 100644 --- a/data_generation/data_curation_drip.py +++ b/data_generation/data_curation_drip.py @@ -112,7 +112,13 @@ async def run_batch(processor: OptimizedAPIProcessor, batch_size: int = 100, starting_index: int = 0) -> List[Dict]: existing = jload(output_path) if os.path.exists(output_path) else [] - remaining = requests[len(existing):] + # Resume by content key, not by count: dropped failures (None results) must + # not shift a length-based offset, which would skip unprocessed requests and + # reprocess completed ones. Key on (instruction, injected_input), matching the + # output records' (instruction, input). + done_keys = {(r["instruction"], r["input"]) for r in existing} + remaining = [req for req in requests + if (req["instruction"], req["injected_input"]) not in done_keys] if not remaining: print("All requests already processed!") diff --git a/data_generation/dpo_data_loader.py b/data_generation/dpo_data_loader.py index 2f6087a..5b3c563 100644 --- a/data_generation/dpo_data_loader.py +++ b/data_generation/dpo_data_loader.py @@ -27,7 +27,7 @@ def generate_training_data_dpo(data_dicts, prompt_dict_name, tokenizer): class PreferenceWithExpertDatasetHF: def __init__(self, data_path_list: List[str], tokenizer: PreTrainedTokenizer, attack: str): - prompt_dict_name, attacks = attack.split('_') + prompt_dict_name, attacks = attack.split('_', 1) prompts: List[str] = [] chosens: List[str] = [] rejecteds: List[str] = [] @@ -78,7 +78,13 @@ def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]: # 2) resp: chosen / rejected def enc(texts): - out = self.tokenizer(texts, padding=False, truncation=True, max_length=self.max_target_length, return_tensors=None) + # add_special_tokens=False: the prompt (tokenized separately, with BOS) + # is prepended to this response, so the response must NOT get its own + # BOS — otherwise a spurious BOS lands mid-sequence at the prompt/ + # response boundary (train/inference mismatch). The explicit eos in the + # text is still tokenized to the eos id. + out = self.tokenizer(texts, padding=False, truncation=True, max_length=self.max_target_length, + add_special_tokens=False, return_tensors=None) ids = [torch.tensor(x, dtype=torch.long) for x in out["input_ids"]] return ids diff --git a/data_generation/sft_data_loader.py b/data_generation/sft_data_loader.py index 10c2924..d514a81 100644 --- a/data_generation/sft_data_loader.py +++ b/data_generation/sft_data_loader.py @@ -375,7 +375,7 @@ class SupervisedDataset(Dataset): def __init__(self, data_path_list: List[str], tokenizer, attack, frontend_delimiters, downsample=True): super(SupervisedDataset, self).__init__() - prompt_dict_name, attacks = attack.split('_') + prompt_dict_name, attacks = attack.split('_', 1) self.input_ids = [] self.labels = [] diff --git a/merge_lora.py b/merge_lora.py new file mode 100644 index 0000000..f96e628 --- /dev/null +++ b/merge_lora.py @@ -0,0 +1,70 @@ +"""Merge a trained LoRA adapter into its base model and save a full checkpoint. + +Use this for checkpoints saved as adapters — e.g. QLoRA runs, or models trained +before train_unified.py started merging at save time. It loads the model exactly +the way evaluation does (testing/test.py `load_full_model`, load_as_adapter=True), +so the merged result is guaranteed consistent with eval, then merges and saves a +standalone checkpoint that the default eval path can load directly. + +Usage: + python merge_lora.py \ + --adapter_path out/Meta-Llama-3-8B-Instruct-TextTextText-sep-drip \ + --output_path out/Meta-Llama-3-8B-Instruct-TextTextText-sep-drip-merged + + # base path / model class are inferred from the adapter path; override if needed: + python merge_lora.py --adapter_path --output_path \ + --base_model_path meta-llama/Meta-Llama-3-8B-Instruct \ + --customized_model_class LlamaForCausalLMDRIP +""" + +import argparse + +from testing.test import load_full_model + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--adapter_path", required=True, + help="Directory of the trained LoRA adapter checkpoint.") + ap.add_argument("--output_path", required=True, + help="Where to write the merged full checkpoint.") + ap.add_argument("--base_model_path", default=None, + help="Base model (inferred from --adapter_path if omitted).") + ap.add_argument("--customized_model_class", default="", + help="REGISTRY class key (auto-detected from the path if omitted).") + ap.add_argument("--device", default="auto", + help="device_map for loading (e.g. 'auto', '0', 'cpu').") + args = ap.parse_args() + + device_map = int(args.device) if args.device.isdigit() else args.device + + # Same loading path as evaluation; load_as_adapter attaches the adapter onto + # the base model. The model class / delimiters are resolved internally. + model, tokenizer = load_full_model( + args.adapter_path, + customized_model_class=args.customized_model_class, + load_as_adapter=True, + base_model_path=args.base_model_path, + device_map=device_map, + ) + + if not hasattr(model, "merge_and_unload"): + raise SystemExit( + "Loaded model is not a PEFT adapter (nothing to merge). " + "Check --adapter_path / --customized_model_class." + ) + + print("Merging adapter into base weights ...") + merged = model.merge_and_unload() + + merged.save_pretrained(args.output_path, safe_serialization=True) + tokenizer.save_pretrained(args.output_path) + # Persist the (DRIP) config — including delimiter ids — alongside the weights. + getattr(merged, "base_model", merged).config.save_pretrained(args.output_path) + print(f"Merged checkpoint saved to {args.output_path}") + + +if __name__ == "__main__": + main() diff --git a/testing/test.py b/testing/test.py index aa06b7e..b465c74 100644 --- a/testing/test.py +++ b/testing/test.py @@ -369,7 +369,9 @@ def form_llm_input(data, injection_method, fmt, apply_filter, defense, sample_id llm_input = injection_method(fmt) if injection_method is hackaprompt else [] for i, d in enumerate(data): if injection_method is none: - llm_input.append(apply_testtime_defense(d, fmt, defense)) + # deepcopy: apply_testtime_defense mutates the dict in place, so never + # pass the shared source `d` (would contaminate later attacks/runs). + llm_input.append(apply_testtime_defense(deepcopy(d), fmt, defense)) continue if not d["input"] or injection_method is hackaprompt: continue diff --git a/train_unified.py b/train_unified.py index f03f7a8..505bf01 100644 --- a/train_unified.py +++ b/train_unified.py @@ -665,14 +665,27 @@ def main(argv: Optional[List[str]] = None): trainer.train(resume_from_checkpoint=resolve_resume(training_args, trainer)) # ---- Save ---- - if trainer.is_world_process_zero(): - logger.info("Training done. Saving...") - trainer.save_model(output_dir=training_args.output_dir) + # For LoRA runs, merge the adapter into the base weights so evaluation can + # load a full checkpoint directly (the default path in testing/test.py). + # QLoRA/4-bit adapters cannot be merged, so those are saved as adapters + # (load them with --load_as_adapter). AIR/full-finetune (non-PEFT) just + # saves the full model. trainer.save_state() if trainer.is_world_process_zero(): + logger.info("Training done. Saving...") + to_save = trainer.accelerator.unwrap_model(trainer.model) + if bnb_config is None and hasattr(to_save, "merge_and_unload"): + try: + to_save = to_save.merge_and_unload() + logger.info("Merged LoRA adapter into base weights.") + except Exception as e: # noqa: BLE001 + logger.warning(f"merge_and_unload failed ({e}); saving adapter instead.") + to_save.save_pretrained(training_args.output_dir, safe_serialization=True) tokenizer.save_pretrained(training_args.output_dir) - base_config = getattr(trainer.model, "base_model", trainer.model).config - base_config.save_pretrained(training_args.output_dir) + # Persist the (DRIP) config — incl. delimiter ids — alongside the weights. + getattr(to_save, "base_model", to_save).config.save_pretrained(training_args.output_dir) + if torch.distributed.is_initialized(): + torch.distributed.barrier() if __name__ == "__main__": diff --git a/training/trainer.py b/training/trainer.py index 2c5b158..268467c 100644 --- a/training/trainer.py +++ b/training/trainer.py @@ -294,12 +294,12 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs): chosen_logps, chosen_mean_logits, chosen_aux = self._seq_logps_moe( self.model, inputs["chosen_input_ids"], inputs["chosen_attention_mask"], inputs["chosen_expert_labels"], inputs["prompt_lens"], - return_logits=True, output_router_logits=False) + return_logits=True, output_router_logits=True) rejected_logps, rejected_mean_logits, rejected_aux = self._seq_logps_moe( self.model, inputs["rejected_input_ids"], inputs["rejected_attention_mask"], inputs["rejected_expert_labels"], inputs["prompt_lens"], - return_logits=True, output_router_logits=False) + return_logits=True, output_router_logits=True) # Ref forward: NO aux_loss needed (ref is frozen) with torch.no_grad(): @@ -365,7 +365,9 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs): class DPOTrainerAIR(DPOTrainerOurs): - def create_optimizer(self, special_params_list = ["intermediate_shifts"]): + def create_optimizer(self, special_params_list=None): + if special_params_list is None: + special_params_list = ["intermediate_shifts"] opt_model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model if self.optimizer is None: decay_parameters = self.get_decay_parameter_names(opt_model)