add SimPO

2025-09-01 17:17:05 +00:00 · 2024-06-24 02:12:20 +00:00
parent 84eab13078
commit 82aecd6374
14 changed files with 128 additions and 70 deletions
--- a/applications/ColossalChat/examples/training_scripts/hostfile
+++ b/applications/ColossalChat/examples/training_scripts/hostfile
@@ -1,5 +1 @@
-XXX.XX.XXX.XXX   # Your master IP
-XXX.XX.XXX.XXX   # Your slave IPs
-XXX.XX.XXX.XXX   # Your slave IPs
-XXX.XX.XXX.XXX   # Your slave IPs
-XXX.XX.XXX.XXX   # Your slave IPs
+localhost
--- a/applications/ColossalChat/examples/training_scripts/train_dpo.py
+++ b/applications/ColossalChat/examples/training_scripts/train_dpo.py
@@ -116,7 +116,7 @@ def train(args):
        else:
            model = AutoModelForCausalLM.from_pretrained(args.pretrain)
        disable_dropout(model)
-        if args.enable_reference_model:
+        if not args.disable_reference_model:
            if args.use_flash_attn:
                ref_model = AutoModelForCausalLM.from_pretrained(
                    args.pretrain,
@@ -128,7 +128,7 @@ def train(args):
            disable_dropout(ref_model)
        else:
            ref_model = None
-
+        print("ref_model is None", args.disable_reference_model, ref_model is None)
        if args.lora_rank > 0:
            model = convert_to_lora_module(model, args.lora_rank, lora_train_bias=args.lora_train_bias)

@@ -255,6 +255,9 @@ def train(args):
        save_interval=args.save_interval,
        save_dir=args.save_dir,
        coordinator=coordinator,
+        beta=args.beta,
+        gamma=args.gamma,
+        length_normalization=args.length_normalization,
    )

    trainer.fit(
@@ -296,6 +299,9 @@ if __name__ == "__main__":
    parser.add_argument("--tp", type=int, default=1)
    parser.add_argument("--pp", type=int, default=1)
    parser.add_argument("--sp", type=int, default=1)
+    parser.add_argument("--beta", type=float, default=0.1, help="beta in DPO loss")
+    parser.add_argument("--gamma", type=float, default=0.0, help="gamma in SimPO loss")
+    parser.add_argument("--length_normalization", default=False, action="store_true")
    parser.add_argument("--enable_sequence_parallelism", default=False, action="store_true")
    parser.add_argument("--zero_stage", type=int, default=0, help="Zero stage", choices=[0, 1, 2])
    parser.add_argument("--zero_cpu_offload", default=False, action="store_true")
@@ -312,7 +318,12 @@ if __name__ == "__main__":
    parser.add_argument("--max_length", type=int, default=2048, help="Model max length")
    parser.add_argument("--max_epochs", type=int, default=3)
    parser.add_argument("--batch_size", type=int, default=4)
-    parser.add_argument("--enable_reference_model", type=bool, default=True)
+    parser.add_argument(
+        "--disable_reference_model",
+        action="store_true",
+        default=False,
+        help="Disable the reference model (enabled by default)",
+    )
    parser.add_argument("--mixed_precision", type=str, default="fp16", choices=["fp16", "bf16"], help="Mixed precision")
    parser.add_argument("--lora_rank", type=int, default=0, help="low-rank adaptation matrices rank")
    parser.add_argument(
--- a/applications/ColossalChat/examples/training_scripts/train_dpo.sh
+++ b/applications/ColossalChat/examples/training_scripts/train_dpo.sh
@@ -13,7 +13,7 @@ set_n_least_used_CUDA_VISIBLE_DEVICES() {
    echo "Now CUDA_VISIBLE_DEVICES is set to:"
    echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"
 }
-set_n_least_used_CUDA_VISIBLE_DEVICES 8
+set_n_least_used_CUDA_VISIBLE_DEVICES 4
 # export CUDA_VISIBLE_DEVICES=6

 PROJECT_NAME="dpo"
@@ -24,16 +24,16 @@ PRETRAINED_MODEL_PATH="" # huggingface or local model path
 PRETRAINED_TOKENIZER_PATH="" # huggingface or local tokenizer path

 declare -a dataset=(
-    YOUR/DATA/DIR/arrow/part-00000
-    YOUR/DATA/DIR/arrow/part-00001
-    YOUR/DATA/DIR/arrow/part-00002
-    YOUR/DATA/DIR/arrow/part-00003
-    YOUR/DATA/DIR/arrow/part-00004
-    YOUR/DATA/DIR/arrow/part-00005
-    YOUR/DATA/DIR/arrow/part-00006
-    YOUR/DATA/DIR/arrow/part-00007
-    YOUR/DATA/DIR/arrow/part-00008
-    YOUR/DATA/DIR/arrow/part-00009
+    /Your/Preference/Data/arrow/part-00000
+    /Your/Preference/Data/arrow/part-00001
+    /Your/Preference/Data/arrow/part-00002
+    /Your/Preference/Data/arrow/part-00003
+    /Your/Preference/Data/arrow/part-00004
+    /Your/Preference/Data/arrow/part-00005
+    /Your/Preference/Data/arrow/part-00006
+    /Your/Preference/Data/arrow/part-00007
+    /Your/Preference/Data/arrow/part-00008
+    /Your/Preference/Data/arrow/part-00009
 )

 TIMESTAMP=$(date +%Y-%m-%d-%H-%M-%S)
@@ -41,7 +41,7 @@ FULL_PROJECT_NAME="${PROJECT_NAME}-${TIMESTAMP}"
 SAVE_DIR="${PARENT_SAVE_DIR}${FULL_PROJECT_NAME}"
 CONFIG_FILE="${PARENT_CONFIG_FILE}-${FULL_PROJECT_NAME}.json"

-colossalai run --nproc_per_node 8 --hostfile hostfile --master_port 31312 train_dpo.py \
+colossalai run --nproc_per_node 4 --hostfile hostfile --master_port 31313 train_dpo.py \
    --pretrain $PRETRAINED_MODEL_PATH \
    --checkpoint_path $PRETRAINED_MODEL_PATH \
    --tokenizer_dir $PRETRAINED_TOKENIZER_PATH \
@@ -51,12 +51,14 @@ colossalai run --nproc_per_node 8 --hostfile hostfile --master_port 31312 train_
    --save_dir $SAVE_DIR \
    --config_file $CONFIG_FILE \
    --max_epochs 1 \
-    --accumulation_steps 4 \
-    --batch_size 2 \
+    --accumulation_steps 2 \
+    --batch_size 16 \
    --lr 1e-6 \
+    --beta 0.1 \
    --mixed_precision "bf16" \
    --grad_clip 1.0 \
+    --max_length 1024 \
    --weight_decay 0.01 \
-    --warmup_steps 100 \
+    --warmup_steps 60 \
    --grad_checkpoint \
    --use_wandb
--- a/applications/ColossalChat/examples/training_scripts/train_sft.py
+++ b/applications/ColossalChat/examples/training_scripts/train_sft.py
@@ -271,7 +271,7 @@ def train(args):
    # save model checkpoint after fitting on only rank0
    coordinator.print_on_master("Start saving final model checkpoint")

-    # booster.save_model(model, os.path.join(args.save_path, "modeling"), shard=True)
+    booster.save_model(model, os.path.join(args.save_path, "modeling"), shard=True)
    coordinator.print_on_master(f"Saved final model checkpoint at epoch {args.max_epochs} at folder {args.save_path}")

    coordinator.print_on_master(f"Max CUDA memory usage: {torch.cuda.max_memory_allocated()/1024**2:.2f} MB")
--- a/applications/ColossalChat/examples/training_scripts/train_sft.sh
+++ b/applications/ColossalChat/examples/training_scripts/train_sft.sh
@@ -17,22 +17,22 @@ set_n_least_used_CUDA_VISIBLE_DEVICES() {
 # export CUDA_VISIBLE_DEVICES=4,5,6
 set_n_least_used_CUDA_VISIBLE_DEVICES 2
 PROJECT_NAME="sft"
-PARENT_SAVE_DIR="" # Path to a folder to save checkpoints
-PARENT_TENSORBOARD_DIR="" # Path to a folder to save logs
-PARENT_CONFIG_FILE="" # Path to a folder to save training config logs
-PRETRAINED_MODEL_PATH="" # huggingface or local model path
-PRETRAINED_TOKENIZER_PATH="" # huggingface or local tokenizer path
+PARENT_SAVE_DIR="/home/yeanbang/data/experiment/rlhf_cont/dpo/ckpt" # Path to a folder to save checkpoints
+PARENT_TENSORBOARD_DIR="/home/yeanbang/data/experiment/rlhf_cont/dpo/log" # Path to a folder to save logs
+PARENT_CONFIG_FILE="/home/yeanbang/data/experiment/rlhf_cont/dpo/log" # Path to a folder to save training config logs
+PRETRAINED_MODEL_PATH="/home/yeanbang/data/models/Sheared-LLaMA-1.3B" # huggingface or local model path
+PRETRAINED_TOKENIZER_PATH="/home/yeanbang/data/models/Sheared-LLaMA-1.3B" # huggingface or local tokenizer path
 declare -a dataset=(
-    YOUR/SFT/DATA/DIR/arrow/part-00000
-    YOUR/SFT/DATA/DIR/arrow/part-00001
-    YOUR/SFT/DATA/DIR/arrow/part-00002
-    YOUR/SFT/DATA/DIR/arrow/part-00003
-    YOUR/SFT/DATA/DIR/arrow/part-00004
-    YOUR/SFT/DATA/DIR/arrow/part-00005
-    YOUR/SFT/DATA/DIR/arrow/part-00006
-    YOUR/SFT/DATA/DIR/arrow/part-00007
-    YOUR/SFT/DATA/DIR/arrow/part-00008
-    YOUR/SFT/DATA/DIR/arrow/part-00009
+    /home/yeanbang/data/experiment/rlhf_cont/dpo/dataset_tokenized/sft/arrow/part-00000
+    /home/yeanbang/data/experiment/rlhf_cont/dpo/dataset_tokenized/sft/arrow/part-00001
+    /home/yeanbang/data/experiment/rlhf_cont/dpo/dataset_tokenized/sft/arrow/part-00002
+    /home/yeanbang/data/experiment/rlhf_cont/dpo/dataset_tokenized/sft/arrow/part-00003
+    /home/yeanbang/data/experiment/rlhf_cont/dpo/dataset_tokenized/sft/arrow/part-00004
+    /home/yeanbang/data/experiment/rlhf_cont/dpo/dataset_tokenized/sft/arrow/part-00005
+    /home/yeanbang/data/experiment/rlhf_cont/dpo/dataset_tokenized/sft/arrow/part-00006
+    /home/yeanbang/data/experiment/rlhf_cont/dpo/dataset_tokenized/sft/arrow/part-00007
+    /home/yeanbang/data/experiment/rlhf_cont/dpo/dataset_tokenized/sft/arrow/part-00008
+    /home/yeanbang/data/experiment/rlhf_cont/dpo/dataset_tokenized/sft/arrow/part-00009
 )

 TIMESTAMP=$(date +%Y-%m-%d-%H-%M-%S)
@@ -43,7 +43,7 @@ CONFIG_FILE="${PARENT_CONFIG_FILE}-${FULL_PROJECT_NAME}.json"
 echo $(which colossalai)
 echo $(which python)
 # the real batch size for gradient descent is number_of_node_in_hostfile * nproc_per_node * train_batch_size
-colossalai run --nproc_per_node 2 --master_port 31312 --hostfile ./hostfile train_sft.py \
+colossalai run --nproc_per_node 1 --master_port 31312 --hostfile ./hostfile train_sft.py \
    --pretrain $PRETRAINED_MODEL_PATH \
    --tokenizer_dir $PRETRAINED_TOKENIZER_PATH \
    --save_interval 4000 \
@@ -51,15 +51,12 @@ colossalai run --nproc_per_node 2 --master_port 31312 --hostfile ./hostfile trai
    --save_path $SAVE_DIR \
    --config_file $CONFIG_FILE \
    --lora_rank 0 \
-    --plugin 3d \
-    --tp 2 \
-    --pp 1 \
-    --zero_stage 0 \
-    --batch_size 2 \
-    --max_epochs 3 \
-    --accumulation_steps 1 \
+    --plugin zero2 \
+    --batch_size 4 \
+    --max_epochs 1 \
+    --accumulation_steps 4 \
    --lr 5e-5 \
-    --max_len 400 \
+    --max_len 1000 \
    --grad_checkpoint \
    --use_wandb \
    --use_flash_attn