sheng22213/pku-safe-rlhf-masked-arm-round4-extreme-v2-training

PKU SafeRLHF Masked ARM Round4 Extreme V2 Training Bundle

This bundle contains the training data and code for the masked_round4_extreme_v2 ARM training runs.

V2 keeps the original round4 span-based loss mask, and adds a loser-side low-probability filter:

• —Chosen and rejected responses use the masked round4 extreme span mask.
• —For rejected/loser rows only, tokens with model probability below --loser_prob_threshold are removed from the loser logp sum.
• —These removed loser tokens still remain in the autoregressive context.
• —The default V2 threshold is 1e-3; the same code supports 1e-4 and 1e-5.

Directory Layout

• —code/: Python training code.
• —scripts/: PBS and submission scripts used on the source HPC cluster.
• —data/: Prepared JSONL train/eval data.

Main Files

• —code/train_arm_llama_masked_round4_v2.py: V2 training entrypoint.
• —code/masked_arm_round4_v2_trainer.py: V2 trainer implementing the loser-side threshold.
• —code/masked_arm_round4_trainer.py: Base masked ARM trainer used by V2.
• —code/train_arm_llama_masked_round4.py: Base masked training entrypoint; V2 imports debug helpers from this file.
• —code/prepare_masked_arm_round4_datasets.py: Dataset preparation/alignment script.

Data

• —data/helpfulness_train95_masked_round4_extreme.jsonl: 9500 examples.
• —data/helpfulness_eval5_masked_round4_extreme.jsonl: 500 examples.
• —data/harmlessness_train95_masked_round4_extreme.jsonl: 9500 examples.
• —data/harmlessness_eval5_masked_round4_extreme.jsonl: 500 examples.
• —data/*smoke*: 2-example smoke-test subsets.

Each row contains:

• —raw_prompt
• —prompt
• —chosen
• —rejected
• —dimension
• —chosen_nonrelated_reward_spans
• —rejected_nonrelated_reward_spans

Example Training Commands

Helpful V2, beta 0.5:

accelerate launch --gpu_ids 0,1 --num_processes 2 code/train_arm_llama_masked_round4_v2.py \
  --preference_dataset=PKU_SafeRLHF_Helpfulness \
  --preference_train_jsonl=data/helpfulness_train95_masked_round4_extreme.jsonl \
  --preference_eval_jsonl=data/helpfulness_eval5_masked_round4_extreme.jsonl \
  --algorithm=arm \
  --model_name_or_path=/path/to/alpaca-7b-reproduced \
  --beta=0.5 \
  --learning_rate=5e-4 \
  --num_train_epochs=1 \
  --output_dir=checkpoints/helpfulness/arm_beta_0p5_masked_round4_extreme_v2_loser1em4 \
  --per_device_train_batch_size=4 \
  --gradient_accumulation_steps=4 \
  --per_device_eval_batch_size=4 \
  --logging_steps=10 \
  --evaluation_strategy=steps \
  --eval_steps=20 \
  --save_strategy=steps \
  --save_steps=1000 \
  --lr_scheduler_type=cosine \
  --warmup_steps=20 \
  --weight_decay=0.05 \
  --gradient_checkpointing=True \
  --bf16=True \
  --max_prompt_length=512 \
  --max_length=1024 \
  --report_to=none \
  --remove_unused_columns=False \
  --loser_prob_threshold=1e-4

Harmless V2, beta 0.01:

accelerate launch --gpu_ids 0,1 --num_processes 2 code/train_arm_llama_masked_round4_v2.py \
  --preference_dataset=PKU_SafeRLHF_Harmlessness \
  --preference_train_jsonl=data/harmlessness_train95_masked_round4_extreme.jsonl \
  --preference_eval_jsonl=data/harmlessness_eval5_masked_round4_extreme.jsonl \
  --algorithm=arm \
  --model_name_or_path=/path/to/alpaca-7b-reproduced \
  --beta=0.01 \
  --learning_rate=5e-4 \
  --num_train_epochs=1 \
  --output_dir=checkpoints/harmlessness/arm_beta_0p01_masked_round4_extreme_v2_loser1em4 \
  --per_device_train_batch_size=4 \
  --gradient_accumulation_steps=4 \
  --per_device_eval_batch_size=4 \
  --logging_steps=10 \
  --evaluation_strategy=steps \
  --eval_steps=20 \
  --save_strategy=steps \
  --save_steps=1000 \
  --lr_scheduler_type=cosine \
  --warmup_steps=20 \
  --weight_decay=0.05 \
  --gradient_checkpointing=True \
  --bf16=True \
  --max_prompt_length=512 \
  --max_length=1024 \
  --report_to=none \
  --remove_unused_columns=False \
  --loser_prob_threshold=1e-4

Change --loser_prob_threshold to 1e-3, 1e-4, or 1e-5 for the threshold sweep.

Original HPC Runs

The source HPC scripts submitted four threshold-sweep jobs:

• —harmlessness, loser_prob_threshold=1e-4
• —harmlessness, loser_prob_threshold=1e-5
• —helpfulness, loser_prob_threshold=1e-4
• —helpfulness, loser_prob_threshold=1e-5

The PBS scripts are included as references, but paths should be edited for other machines.