diff --git a/.gitignore b/.gitignore index 6312753..a5bdcc1 100644 --- a/.gitignore +++ b/.gitignore @@ -38,6 +38,7 @@ core.* # Datasets and credentials (see README for setup) /datasets +!/datasets/MANIFEST.md **/openai_configs.yaml # Downloaded model checkpoints diff --git a/CHECKPOINTS.md b/CHECKPOINTS.md new file mode 100644 index 0000000..dee166c --- /dev/null +++ b/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= # 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 → `.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 ` (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. diff --git a/README.md b/README.md index 70ff581..7d84392 100644 --- a/README.md +++ b/README.md @@ -176,6 +176,38 @@ instead (QLoRA runs, or models trained earlier), merge them first: python -m training.merge_lora --adapter_path --output_path ``` +> 📄 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 diff --git a/datasets/MANIFEST.md b/datasets/MANIFEST.md new file mode 100644 index 0000000..a6020ad --- /dev/null +++ b/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).* diff --git a/scripts/llama8b/agentdojo/drip_4roles.sh b/scripts/llama8b/agentdojo/drip_4roles.sh index 36b8e8e..65bf22c 100644 --- a/scripts/llama8b/agentdojo/drip_4roles.sh +++ b/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}" \ diff --git a/scripts/llama8b/alpaca/drip_alpaca.sh b/scripts/llama8b/alpaca/drip_alpaca.sh index f108aa3..424597b 100644 --- a/scripts/llama8b/alpaca/drip_alpaca.sh +++ b/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}" \ diff --git a/scripts/llama8b/alpaca/drip_alpaca_4roles.sh b/scripts/llama8b/alpaca/drip_alpaca_4roles.sh index 356c18d..1b5a401 100644 --- a/scripts/llama8b/alpaca/drip_alpaca_4roles.sh +++ b/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}" \ diff --git a/scripts/llama8b/alpaca/drip_nojudge.sh b/scripts/llama8b/alpaca/drip_nojudge.sh index f986d48..3c8c03a 100644 --- a/scripts/llama8b/alpaca/drip_nojudge.sh +++ b/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}" \ diff --git a/scripts/llama8b/sep/air_sep_dpo.sh b/scripts/llama8b/sep/air_sep_dpo.sh index ae51be3..8761fca 100644 --- a/scripts/llama8b/sep/air_sep_dpo.sh +++ b/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}" \ diff --git a/scripts/llama8b/sep/drip_concatfusion.sh b/scripts/llama8b/sep/drip_concatfusion.sh index cf8788f..04c27eb 100644 --- a/scripts/llama8b/sep/drip_concatfusion.sh +++ b/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}" \ diff --git a/scripts/llama8b/sep/drip_embeddingshift.sh b/scripts/llama8b/sep/drip_embeddingshift.sh index 25d6b49..dc2ee0e 100644 --- a/scripts/llama8b/sep/drip_embeddingshift.sh +++ b/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}" \ diff --git a/scripts/llama8b/sep/drip_nofusion.sh b/scripts/llama8b/sep/drip_nofusion.sh index c15ab28..5b86a85 100644 --- a/scripts/llama8b/sep/drip_nofusion.sh +++ b/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}" \ diff --git a/scripts/llama8b/sep/drip_sep.sh b/scripts/llama8b/sep/drip_sep.sh index 5cf0229..5c3338d 100644 --- a/scripts/llama8b/sep/drip_sep.sh +++ b/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}" \ diff --git a/scripts/llama8b/sep/ise_sep.sh b/scripts/llama8b/sep/ise_sep.sh index be4c8dd..7fbc12e 100644 --- a/scripts/llama8b/sep/ise_sep.sh +++ b/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}" \ diff --git a/scripts/llama8b/sep/pft_sep.sh b/scripts/llama8b/sep/pft_sep.sh index 1e73796..14ea31d 100644 --- a/scripts/llama8b/sep/pft_sep.sh +++ b/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}" \ diff --git a/scripts/llama8b/sep/secalign_sep.sh b/scripts/llama8b/sep/secalign_sep.sh index 6d4c696..f035785 100644 --- a/scripts/llama8b/sep/secalign_sep.sh +++ b/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}" \ diff --git a/scripts/llama8b/sep/struq_sep.sh b/scripts/llama8b/sep/struq_sep.sh index 99a665c..ddd0541 100644 --- a/scripts/llama8b/sep/struq_sep.sh +++ b/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}" \ diff --git a/scripts/mistral7b/alpaca/drip_alpaca.sh b/scripts/mistral7b/alpaca/drip_alpaca.sh index 52e3850..64400eb 100644 --- a/scripts/mistral7b/alpaca/drip_alpaca.sh +++ b/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}" \ diff --git a/scripts/mistral7b/sep/air_sep_dpo.sh b/scripts/mistral7b/sep/air_sep_dpo.sh index 22dbb68..a02753d 100644 --- a/scripts/mistral7b/sep/air_sep_dpo.sh +++ b/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}" \ diff --git a/scripts/mistral7b/sep/drip_sep.sh b/scripts/mistral7b/sep/drip_sep.sh index 2a852b7..a88ea9d 100644 --- a/scripts/mistral7b/sep/drip_sep.sh +++ b/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}" \ diff --git a/scripts/mistral7b/sep/ise_sep.sh b/scripts/mistral7b/sep/ise_sep.sh index 67f65ae..754131c 100644 --- a/scripts/mistral7b/sep/ise_sep.sh +++ b/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}" \ diff --git a/scripts/mistral7b/sep/pft_sep.sh b/scripts/mistral7b/sep/pft_sep.sh index e2741fe..a294633 100644 --- a/scripts/mistral7b/sep/pft_sep.sh +++ b/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}" \ diff --git a/scripts/mistral7b/sep/secalign_sep.sh b/scripts/mistral7b/sep/secalign_sep.sh index 60feb10..ccc670a 100644 --- a/scripts/mistral7b/sep/secalign_sep.sh +++ b/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}" \ diff --git a/scripts/mistral7b/sep/struq_sep.sh b/scripts/mistral7b/sep/struq_sep.sh index 9ad11be..94fa72f 100644 --- a/scripts/mistral7b/sep/struq_sep.sh +++ b/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}" \