Artifact-evaluation fixes: is_moe, judge comments, LICENSE, GPU params, manifest, checksums, appendix PDFs

.gitignore

@@ -38,6 +38,7 @@ core.*
 
 # Datasets and credentials (see README for setup)
 /datasets
+!/datasets/MANIFEST.md
 **/openai_configs.yaml
 
 # Downloaded model checkpoints

CHECKPOINTS.md

@@ -0,0 +1,54 @@
+# Pretrained checkpoints — pinned revisions & checksums
+
+The four DRIP LoRA adapters are hosted on the Hugging Face Hub. For reproducible
+evaluation, **pin an exact revision** (a commit SHA, not a moving branch) and
+**verify the downloaded files against the checksums** below.
+
+| # | Repo (`Kelsey98/…`) | Base model | Template |
+|---|---|---|---|
+| 1 | `Llama-3.1-8B-Instruct-TextTextText-4roles-toolcall-drip` | `meta-llama/Llama-3.1-8B-Instruct` | 4-role |
+| 2 | `Llama-3.1-8B-Instruct-TextTextText-drip` | `meta-llama/Llama-3.1-8B-Instruct` | 3-role |
+| 3 | `Meta-Llama-3-8B-Instruct-TextTextText-drip` | `meta-llama/Meta-Llama-3-8B-Instruct` | 3-role |
+| 4 | `Mistral-7B-Instruct-v0.3-TextTextTextMistral-drip` | `mistralai/Mistral-7B-Instruct-v0.3` | 3-role |
+
+## 1. Download a pinned revision
+
+```bash
+REPO=Llama-3.1-8B-Instruct-TextTextText-4roles-toolcall-drip
+REV=<commit-sha>          # the exact revision from the table below
+huggingface-cli download "Kelsey98/$REPO" --revision "$REV" --local-dir "$REPO"
+```
+
+Find a repo's revisions with `huggingface-cli scan-cache` after download, or on the
+Hub under *Files and versions → History* (each commit has a full SHA).
+
+## 2. Verify checksums
+
+After downloading, recompute SHA-256 over the adapter files and compare to the
+manifest:
+
+```bash
+# from inside the downloaded adapter directory
+find . -type f \( -name '*.safetensors' -o -name '*.json' -o -name '*.model' \) \
+  | sort | xargs sha256sum
+```
+
+## 3. Checksum manifest
+
+SHA-256 of the weights `.safetensors` file in each repo. On the Hub, open the file
+under *Files and versions → `<file>.safetensors`*; the **`SHA256:`** line on that
+blob page is the value below (this equals `sha256sum` of the downloaded file — the
+separate "Xet hash" is Hugging Face's internal dedup hash, not used for
+verification).
+
+| # | Repo (`Kelsey98/…`) | Revision (commit SHA) | Weights SHA-256 | Size |
+|---|---|---|---|---|
+| 1 | `Llama-3.1-8B-Instruct-TextTextText-4roles-toolcall-drip` | `d97febc67553b2d5bd2e14bd6ea15d2294837b67` | `5abaaed207cb2f8eed129852a77d583203feca27108dfada7c084ff52b8a3d95` | 4.44 GB |
+| 2 | `Llama-3.1-8B-Instruct-TextTextText-drip` | `314c5930d141da2c9d2e2b2a907f830e07e519fc` | `4a0b7e9e228b042537661d541437c081eb9f88139c380bb066e68b2a9531ad53` | 4.44 GB |
+| 3 | `Meta-Llama-3-8B-Instruct-TextTextText-drip` | `1497185b5c93333dddf3892117af9828516dc37b` | `e7b0b5f25960d80a08b163cf48151dc89988b68637bef65f0bb4fdc1004bc7a4` | 4.44 GB |
+| 4 | `Mistral-7B-Instruct-v0.3-TextTextTextMistral-drip` | `ef717b36c82284fea93cfe6434cf13a6eb8aebab` | `f57d83a7edafe9afeab2d7d0fa7b5eec821b06e6691738f175554cbaf6ad8c21` | 1.31 GB |
+
+Each revision is the archived commit on the Hub (`HfApi().model_info(repo).sha`);
+download with `--revision <sha>` (Section 1) to fetch that exact snapshot, then
+verify the weights against the SHA-256 above. This lets evaluators confirm they are
+running the exact artifact evaluated in the paper.

README.md

@@ -176,6 +176,38 @@ instead (QLoRA runs, or models trained earlier), merge them first:
 python -m training.merge_lora --adapter_path <adapter_dir> --output_path <merged_dir>
 ```
 
+> 📄 See [`datasets/MANIFEST.md`](./datasets/MANIFEST.md) for exactly which dataset
+> each launcher consumes, and [`CHECKPOINTS.md`](./CHECKPOINTS.md) for pinned
+> adapter revisions and checksums.
+
+### Hardware & GPU configuration
+
+The launchers are **not** pinned to a fixed GPU count — the reference values are
+just defaults you can override with environment variables:
+
+```bash
+# use 2 GPUs instead of the default, on specific device ids
+NPROC_PER_NODE=2 CUDA_VISIBLE_DEVICES=0,1 bash ./scripts/llama8b/sep/drip_sep.sh
+```
+
+- `NPROC_PER_NODE` — number of processes / GPUs for `torch.distributed.run`
+  (defaults to the value baked into each script: 6 for most, 8/4 for a couple).
+- `CUDA_VISIBLE_DEVICES` — which device ids to use.
+
+**Training on fewer / smaller GPUs.** Our reference run uses 6–8× 48 GB GPUs, but
+DRIP is a lightweight LoRA fine-tune and fits far smaller setups by trading batch
+size for gradient accumulation (the effective batch size is
+`per_device_train_batch_size × gradient_accumulation_steps × NPROC_PER_NODE`):
+
+- Lower `--per_device_train_batch_size` (e.g. `1`) and raise
+  `--gradient_accumulation_steps` to keep the effective batch size constant.
+- Set `NPROC_PER_NODE=1 CUDA_VISIBLE_DEVICES=0` to run single-GPU.
+- For **24 GB** cards, add QLoRA-style 4-bit base loading (`bitsandbytes`) and keep
+  `--model_max_length` at the script default; the DRIP modules (`deinstruction_shift`,
+  fusion) are tiny and stay in fp16/bf16.
+
+These knobs change only throughput/precision, not the DRIP method.
+
 ### 3-role vs 4-role: train separately
 
 DRIP supports two chat formats, and you train a **separate** model for each (they

datasets/MANIFEST.md

@@ -0,0 +1,63 @@
+# Dataset manifest
+
+Every dataset referenced by the training, baseline, and evaluation scripts, with
+its role, the script(s) that consume it, and its source.
+
+**Source of truth.** All curated files live in the Zenodo data archive
+(DOI [10.5281/zenodo.20603331](https://doi.org/10.5281/zenodo.20603331)). Download
+and extract it into the repository root so the paths below resolve (see the main
+[README → Download the data](../README.md#3-download-the-data)):
+
+```bash
+wget -O datasets.zip "https://zenodo.org/records/20603331/files/datasets.zip?download=1"
+unzip datasets.zip && mv datasets1/ datasets/
+```
+
+Every curated file can also be **regenerated from scratch** with the pipeline in
+[`data_generation/`](../data_generation/README.md) (needs an `OPENAI_API_KEY`).
+
+---
+
+## Training / baseline DPO datasets
+
+| File | Consumed by (launcher) | Role | Regenerate with |
+|---|---|---|---|
+| `sep/sep_data_cleaned_dpo_gpt.json` | `{llama8b,mistral7b}/sep/drip_sep.sh`, `llama8b/sep/drip_{nofusion,concatfusion,embeddingshift}.sh` | **DRIP** main SEP DPO pairs (3-role, cleaned, GPT-authored chosen) | `data_generation/SEP_to_DPO.py` + `data_curation_drip.py` |
+| `alpaca_data_cleaned_dpo_gpt.json` | `{llama8b,mistral7b}/alpaca/drip_alpaca.sh` | DRIP Alpaca DPO pairs (3-role) | `data_generation/CleanAlpaca_to_DPO.py` + `data_curation_drip.py` |
+| `alpaca_injecagent_dpo_combined.json` | `llama8b/alpaca/drip_alpaca_4roles.sh`, `llama8b/agentdojo/drip_4roles.sh` | 4-role tool-calling DPO (~20K Alpaca + ~1K InjecAgent) | `data_generation/data_curation_drip_toolcall.py` |
+| `alpaca_data_dpo_ablate_no_judge.json` | `llama8b/alpaca/drip_nojudge.sh` | Ablation: Alpaca DPO **without** the judge filter | `data_curation_drip.py --no_judge` |
+| `sep/sep_data_cleaned.json` | `{llama8b,mistral7b}/sep/{struq,pft,ise}_sep.sh` | Baselines (StruQ / PFT / ISE) — cleaned SEP data | `data_generation/SEP_to_DPO.py` |
+| `sep/sep_data_dpo.json` | `{llama8b,mistral7b}/sep/secalign_sep.sh` | Baseline (SecAlign) — SEP DPO pairs | `data_generation/SEP_to_DPO.py` |
+| `sep/sep_data_origdata_dpo.json` | `{llama8b,mistral7b}/sep/air_sep_dpo.sh` | Baseline (AIR) — SEP DPO on original data | `data_generation/data_curation_orig.py` |
+
+## Evaluation datasets
+
+| File | Used by | Role |
+|---|---|---|
+| `SEP_dataset.json` | `testing/sep/test_sep.py` | SEP role-separation benchmark (9,160 examples) |
+| `alpaca_data_cleaned.json` | `testing/test.py` (Alpaca injection / utility) | Cleaned AlpacaFarm instructions |
+| `ifeval/input_data.jsonl` | `testing/ifeval/` | IFEval prompts (541) |
+| `mtbench.jsonl` | `testing/mt_bench/` | MT-Bench questions |
+| `judge_prompts.jsonl` | `testing/mt_bench/gen_judgment.py` | MT-Bench LLM-judge prompts |
+| `davinci_003_outputs.json` | AlpacaEval 2.0 | Reference outputs (davinci-003) |
+| `gpt4o_outputs.json` | AlpacaEval 2.0 | Reference outputs (GPT-4o) |
+| `injecagent/tools.json` | `testing/injecagent/` | InjecAgent tool specifications |
+| `injecagent/attacker_param_cache.json` | `testing/injecagent/`, tool-call curation | Cached attacker tool arguments |
+| `injecagent/attacker_simulated_responses.json` | `testing/injecagent/` | Simulated attacker tool responses |
+
+## Curation source / intermediate files
+
+Produced and consumed inside [`data_generation/`](../data_generation/README.md);
+listed for completeness (also shipped in the Zenodo archive):
+
+`sep/train_dataset.json` (raw SEP training split), `alpaca_data.json`,
+`alpaca_data_injected_diff_output.json`, `sep/sep_data_full_withanswer.json`,
+`sep/sep_data_injected_diff_output.json`, `sep/sep_data_origdata_sft.json`,
+`sep/sep_data_air_dpo.json`, `sep/sep_dpo_retrieved.jsonl`,
+`sep/sep_dpo_submit.jsonl`, `injecagent_dpo.json`, `injecagent_ds_dpo.json`.
+
+---
+
+*Counts and field schemas for the SEP splits are documented in
+[`testing/sep/README.md`](../testing/sep/README.md); the 4-role tool-calling mix is
+documented in [`testing/agentdojo/README.md`](../testing/agentdojo/README.md).*

scripts/llama8b/agentdojo/drip_4roles.sh

@@ -18,7 +18,7 @@ export TORCH_NCCL_ASYNC_ERROR_HANDLING=1
 export TORCH_NCCL_TIMEOUT_MS=1800000
 export TORCH_NCCL_TRACE_BUFFER_SIZE=20480
 export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:512
-export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5}"
 
 SCRIPT_PATH="train_unified.py"
 BASELINE="drip"
@@ -38,7 +38,7 @@ OBJECTIVE="dpo"
 MODEL_FAMILY="llama"
 ARCH="fuse"
 
-python -m torch.distributed.run --nproc_per_node=6 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-6}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/llama8b/alpaca/drip_alpaca.sh

@@ -17,7 +17,7 @@ OBJECTIVE="dpo"
 MODEL_FAMILY="llama"
 ARCH="fuse"
 
-python -m torch.distributed.run --nproc_per_node=6 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-6}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/llama8b/alpaca/drip_alpaca_4roles.sh

@@ -18,7 +18,7 @@ OBJECTIVE="dpo"
 MODEL_FAMILY="llama"
 ARCH="fuse"
 
-python -m torch.distributed.run --nproc_per_node=6 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-6}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/llama8b/alpaca/drip_nojudge.sh

@@ -18,7 +18,7 @@ OBJECTIVE="dpo"
 MODEL_FAMILY="llama"
 ARCH="fuse"
 
-python -m torch.distributed.run --nproc_per_node=6 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-6}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/llama8b/sep/air_sep_dpo.sh

@@ -18,7 +18,7 @@ OBJECTIVE="dpo"
 MODEL_FAMILY="llama"
 ARCH="air"
 
-python -m torch.distributed.run --nproc_per_node=6 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-6}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/llama8b/sep/drip_concatfusion.sh

@@ -20,7 +20,7 @@ OBJECTIVE="dpo"
 MODEL_FAMILY="llama"
 ARCH="concatfuse"
 
-python -m torch.distributed.run --nproc_per_node=6 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-6}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/llama8b/sep/drip_embeddingshift.sh

@@ -18,7 +18,7 @@ OBJECTIVE="dpo"
 MODEL_FAMILY="llama"
 ARCH="embeddingshift"
 
-python -m torch.distributed.run --nproc_per_node=4 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-4}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/llama8b/sep/drip_nofusion.sh

@@ -18,7 +18,7 @@ OBJECTIVE="dpo"
 MODEL_FAMILY="llama"
 ARCH="nofuse"
 
-python -m torch.distributed.run --nproc_per_node=6 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-6}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/llama8b/sep/drip_sep.sh

@@ -17,7 +17,7 @@ OBJECTIVE="dpo"
 MODEL_FAMILY="llama"
 ARCH="fuse"
 
-python -m torch.distributed.run --nproc_per_node=6 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-6}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/llama8b/sep/ise_sep.sh

@@ -19,7 +19,7 @@ OBJECTIVE="sft"
 MODEL_FAMILY="llama"
 ARCH="ise"
 
-python -m torch.distributed.run --nproc_per_node=8 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-8}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/llama8b/sep/pft_sep.sh

@@ -18,7 +18,7 @@ OBJECTIVE="sft"
 MODEL_FAMILY="llama"
 ARCH="possep"
 
-python -m torch.distributed.run --nproc_per_node=6 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-6}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/llama8b/sep/secalign_sep.sh

@@ -18,9 +18,9 @@ OBJECTIVE="secalign_dpo"
 MODEL_FAMILY="llama"
 ARCH="base"
 
-export CUDA_VISIBLE_DEVICES=1,2,3,4,5,6
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-1,2,3,4,5,6}"
 
-python -m torch.distributed.run --nproc_per_node=6 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-6}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/llama8b/sep/struq_sep.sh

@@ -18,7 +18,7 @@ OBJECTIVE="struq_sft"
 MODEL_FAMILY="llama"
 ARCH="base"
 
-python -m torch.distributed.run --nproc_per_node=6 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-6}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/mistral7b/alpaca/drip_alpaca.sh

@@ -18,7 +18,7 @@ OBJECTIVE="dpo"
 MODEL_FAMILY="mistral"
 ARCH="fuse"
 
-python -m torch.distributed.run --nproc_per_node=6 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-6}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/mistral7b/sep/air_sep_dpo.sh

@@ -17,7 +17,7 @@ OBJECTIVE="dpo"
 MODEL_FAMILY="mistral"
 ARCH="air"
 
-python -m torch.distributed.run --nproc_per_node=6 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-6}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/mistral7b/sep/drip_sep.sh

@@ -18,7 +18,7 @@ OBJECTIVE="dpo"
 MODEL_FAMILY="mistral"
 ARCH="fuse"
 
-python -m torch.distributed.run --nproc_per_node=6 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-6}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/mistral7b/sep/ise_sep.sh

@@ -20,7 +20,7 @@ MODEL_FAMILY="mistral"
 ARCH="ise"
 
 
-python -m torch.distributed.run --nproc_per_node=6 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-6}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/mistral7b/sep/pft_sep.sh

@@ -19,7 +19,7 @@ MODEL_FAMILY="mistral"
 ARCH="possep"
 
 
-python -m torch.distributed.run --nproc_per_node=6 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-6}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/mistral7b/sep/secalign_sep.sh

@@ -19,7 +19,7 @@ MODEL_FAMILY="mistral"
 ARCH="base"
 
 
-python -m torch.distributed.run --nproc_per_node=6 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-6}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \

scripts/mistral7b/sep/struq_sep.sh

@@ -19,7 +19,7 @@ MODEL_FAMILY="mistral"
 ARCH="base"
 
 
-python -m torch.distributed.run --nproc_per_node=6 --master_port=29951 "$SCRIPT_PATH" \
+python -m torch.distributed.run --nproc_per_node="${NPROC_PER_NODE:-6}" --master_port=29951 "$SCRIPT_PATH" \
   --objective "${OBJECTIVE}" \
   --model-family "${MODEL_FAMILY}" \
   --arch "${ARCH}" \