From ac3a72fbc991f5167bc5885fa6dcfeb0960ec512 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 3 Jun 2026 13:12:58 +0000 Subject: [PATCH 1/7] Add SEP/Alpaca train-vs-eval leakage checker Read-only integrity tool: normalizes the (system prompt, clean data) content of a training set and an eval set (handling the different field-name schemas of datasets/sep/train_dataset.json vs datasets/SEP_dataset.json) and reports any overlap, exiting non-zero on leakage. Guards against train/test contamination that would inflate results. Verified on synthetic data. --- data_generation/check_sep_leakage.py | 113 +++++++++++++++++++++++++++ 1 file changed, 113 insertions(+) create mode 100644 data_generation/check_sep_leakage.py 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() From c0c57ee036085679cd5fb5c8b1a28d2f3f38a255 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 3 Jun 2026 13:25:22 +0000 Subject: [PATCH 2/7] Fix data-loader bugs: spurious DPO response BOS, shared-dict mutation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - (MAIN, result-affecting) data_generation/dpo_data_loader.py: the DPO collator tokenized prompt and response separately, and enc() used the default add_special_tokens=True, prepending a BOS to every chosen/rejected response. Concatenated with the prompt (which already has BOS), this put a spurious BOS mid-sequence at the prompt/response boundary — a train/inference mismatch (at inference the answer follows the response delimiter with no BOS). The BOS term cancels in the DPO chosen-vs-rejected margin so the preference gradient is largely unaffected, but the conditioning of the answer tokens is not. Fix: add_special_tokens=False (the explicit eos string still maps to the eos id). NOTE: changes training data -> re-train DRIP/DPO models to fully apply. - testing/test.py form_llm_input: the no-attack branch passed the shared source dict into apply_testtime_defense, which mutates in place; deepcopy it so a non-'none' defense run cannot contaminate later attacks on the same data. - Harden attack.split('_') -> split('_', 1) in both data loaders. Verified: compiles, pytest 25 passed/1 skipped, ruff F821/F811 clean. --- data_generation/dpo_data_loader.py | 10 ++++++++-- data_generation/sft_data_loader.py | 2 +- testing/test.py | 4 +++- 3 files changed, 12 insertions(+), 4 deletions(-) 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/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 From 1a9a7d15e327a5ef41b8d7c12c529c73d3cf13e9 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 3 Jun 2026 13:32:38 +0000 Subject: [PATCH 3/7] Fix: MoE load-balancing aux loss was never applied in DPO training DPOTrainerMOE.compute_loss fetched the router aux loss with output_router_logits=False on the policy chosen/rejected forwards, so _seq_logps_moe always returned aux=None and the 'if chosen_aux is not None and rejected_aux is not None' guard was never true -- the load-balancing term was silently dropped (loss = dpo_loss only). The MoE forward only computes load_balancing_loss_func when output_router_logits=True (qwen_moe_drip.py:305). Enable it on the two POLICY forwards; the reference forwards stay False (frozen, no aux needed). NOTE: affects only the Qwen3-MoE variant; re-train it to apply. Verified: compiles, pytest 25 passed/1 skipped, ruff clean. --- training/trainer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/training/trainer.py b/training/trainer.py index 2c5b158..511a8dd 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(): From 87abf7e658e8463e2325a6a6316606f171386848 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 3 Jun 2026 13:40:38 +0000 Subject: [PATCH 4/7] Merge LoRA adapter at save time; fix AIR mutable default arg - train_unified.py: training saved a LoRA adapter via trainer.save_model, but testing/test.py's default load path expects a fully merged checkpoint (and the eval scripts don't pass --load_as_adapter) -- so evaluating a freshly trained model would load base/untrained weights. Now, on the main process, unwrap the model and (for non-QLoRA LoRA runs) merge_and_unload() before save_pretrained, producing a full checkpoint the default eval path can load. QLoRA/4-bit and non-PEFT (AIR/full-finetune) paths save as before; merge failures fall back to adapter save. Adds a post-save distributed barrier. - trainer.py: DPOTrainerAIR.create_optimizer used a mutable default arg (['intermediate_shifts']); switch to the None pattern (behavior unchanged). Verified: compiles, pytest 25 passed/1 skipped, ruff clean. NOTE: the save-path change can't be runtime-tested here (no GPU/distributed) -- smoke-test one short training run to confirm the merged checkpoint loads in eval. --- train_unified.py | 23 ++++++++++++++++++----- training/trainer.py | 4 +++- 2 files changed, 21 insertions(+), 6 deletions(-) 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 511a8dd..268467c 100644 --- a/training/trainer.py +++ b/training/trainer.py @@ -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) From 962106f5821145cabc4742750c9934d0e99b9cd8 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 3 Jun 2026 13:50:24 +0000 Subject: [PATCH 5/7] Add standalone merge_lora.py to merge LoRA adapters into full checkpoints For checkpoints saved as adapters (QLoRA runs, or models trained before train_unified started merging at save time). It reuses testing/test.py's load_full_model(load_as_adapter=True) so loading matches evaluation exactly, then merge_and_unload() and saves a full checkpoint (weights + tokenizer + delimiter config) that the default eval path loads directly. Base path and model class are auto-detected from the adapter path (overridable). Document it in the README Training section. --- README.md | 8 ++++++ merge_lora.py | 70 +++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 78 insertions(+) create mode 100644 merge_lora.py 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/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() From 74415cb1ed4f089c476c4c3d67c3b223a330f498 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 3 Jun 2026 14:25:37 +0000 Subject: [PATCH 6/7] Fix inverted judge description in data_generation/README.md The DPO chosen-side judge (tasktracker_judge_prompt2.txt) returns INCLUDED when the accompanying text was used as data and NOT INCLUDED when it was ignored; data_curation_drip.py correctly retries on NOT INCLUDED. The README described it backwards ('checks whether the probe leaked... on INCLUDED it retries'). Correct the description to match the code; no code change (the logic was right). --- data_generation/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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. From 42044a93baabd8e75f71aabc20f35b8cc68e0a95 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 3 Jun 2026 14:27:14 +0000 Subject: [PATCH 7/7] Fix run_batch resume offset in data_curation_drip.py Resume skipped already-done requests by COUNT (requests[len(existing):]). Since failed requests (None results) are dropped, len(existing) under-counts attempted items, so the length offset would skip some unprocessed requests and reprocess completed ones. Resume by content key (instruction, injected_input) against the existing output records' (instruction, input) instead, so failures are retried and completed items are skipped correctly. --- data_generation/data_curation_drip.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) 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!")