[Feature] reward model inferencer and dpov2 aligner

configs/accelerate_dsz3_config.yaml

@@ -0,0 +1,23 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  deepspeed_multinode_launcher: standard
+  offload_optimizer_device: none
+  offload_param_device: none
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8
+gpu_ids:
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

examples/dpov2_train.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2024 Statistics and Machine Learning Research Group. All rights reserved.
+import logging
+import os
+import sys
+import copy
+
+from transformers import (
+    HfArgumentParser
+)
+
+from lmflow.datasets import Dataset
+from lmflow.models.auto_model import AutoModel
+from lmflow.pipeline.auto_pipeline import AutoPipeline
+from lmflow.args import (
+    ModelArguments, 
+    DatasetArguments, 
+    AutoArguments,
+)
+from lmflow.utils.common import remove_dataclass_attr_prefix, create_copied_dataclass
+
+
+logger = logging.getLogger(__name__)
+
+
+ReferenceModelArguments = create_copied_dataclass(
+    original_dataclass=ModelArguments, 
+    field_prefix="reference_",
+    class_prefix="Reference"
+)
+
+
+def main():
+    # Parses arguments
+    pipeline_name = "dpov2_aligner"
+    PipelineArguments = AutoArguments.get_pipeline_args_class(pipeline_name)
+
+    parser = HfArgumentParser((
+        ModelArguments, 
+        ReferenceModelArguments,
+        DatasetArguments,
+        PipelineArguments
+    ))
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        # If we pass only one argument to the script and it's the path to a json file,
+        # let's parse it to get our arguments.
+        model_args, ref_model_args, data_args, pipeline_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, ref_model_args, data_args, pipeline_args = parser.parse_args_into_dataclasses()
+
+    ref_model_args_dict = remove_dataclass_attr_prefix(ref_model_args, "reference_")
+    ref_model_args = ModelArguments(**ref_model_args_dict)
+
+    train_dataset = Dataset(data_args)
+    eval_data_args = copy.deepcopy(data_args)
+    eval_data_args.dataset_path = pipeline_args.eval_dataset_path
+    eval_dataset = Dataset(eval_data_args)
+    model = AutoModel.get_model(model_args)
+    ref_model = AutoModel.get_model(ref_model_args)
+    aligner = AutoPipeline.get_pipeline(
+        pipeline_name=pipeline_name,
+        model_args=model_args,
+        data_args=data_args,
+        pipeline_args=pipeline_args,
+        ref_model_args=ref_model_args,
+    )
+
+    res = aligner.align(
+        model=model,
+        ref_model=ref_model,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/rm_inference.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2024 Statistics and Machine Learning Research Group. All rights reserved.
+import logging
+import os
+import sys
+
+from transformers import (
+    HfArgumentParser
+)
+
+from lmflow.datasets import Dataset
+from lmflow.models.auto_model import AutoModel
+from lmflow.pipeline.auto_pipeline import AutoPipeline
+from lmflow.args import (
+    ModelArguments, 
+    DatasetArguments, 
+    AutoArguments,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+def main():
+    # Parses arguments
+    pipeline_name = "rm_inferencer"
+    PipelineArguments = AutoArguments.get_pipeline_args_class(pipeline_name)
+
+    parser = HfArgumentParser((
+        ModelArguments, 
+        DatasetArguments,
+        PipelineArguments
+    ))
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        # If we pass only one argument to the script and it's the path to a json file,
+        # let's parse it to get our arguments.
+        model_args, data_args, pipeline_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, pipeline_args = parser.parse_args_into_dataclasses()
+
+    dataset = Dataset(data_args)
+    model = AutoModel.get_model(model_args, tune_strategy='none', use_accelerator=pipeline_args.use_accelerator)
+    inferencer = AutoPipeline.get_pipeline(
+        pipeline_name=pipeline_name,
+        model_args=model_args,
+        data_args=data_args,
+        pipeline_args=pipeline_args
+    )
+
+    res = inferencer.inference(
+        model,
+        dataset,
+    )
+
+    if pipeline_args.save_results:
+        res.save(pipeline_args.results_path)
+
+
+if __name__ == "__main__":
+    main()

examples/vllm_inference.py

@@ -10,7 +10,7 @@
 )
 
 from lmflow.datasets import Dataset
-from lmflow.models.hf_decoder_model import HFDecoderModel
+from lmflow.models.auto_model import AutoModel
 from lmflow.pipeline.auto_pipeline import AutoPipeline
 from lmflow.args import (
     ModelArguments, 
@@ -40,7 +40,7 @@ def main():
         model_args, data_args, pipeline_args = parser.parse_args_into_dataclasses()
 
     dataset = Dataset(data_args)
-    model = HFDecoderModel(model_args)
+    model = AutoModel.get_model(model_args, tune_strategy='none')
     inferencer = AutoPipeline.get_pipeline(
         pipeline_name=pipeline_name,
         model_args=model_args,

requirements.txt

@@ -6,7 +6,7 @@ peft>=0.10.0
 torch>=2.0.1
 wandb==0.14.0
 deepspeed<=0.14.0
-trl>=0.7.11
+trl==0.8.0
 sentencepiece
 transformers>=4.31.0
 flask

scripts/run_dpo_align.sh

@@ -20,7 +20,7 @@ while [[ $# -ge 1 ]]; do
       dataset_path="$2"
       shift
       ;;
-    -o|--output_lora_path)
+    -o|--output_dir)
       output_dir="$2"
       shift
       ;;

scripts/run_dpov2_align.sh

@@ -0,0 +1,81 @@
+#!/bin/bash
+
+# Parses arguments
+run_name=dpov2_align
+model_name_or_path=meta-llama/Meta-Llama-3-8B-Instruct
+reference_model_name_or_path=meta-llama/Meta-Llama-3-8B-Instruct
+dataset_path=data/iterative-prompt/train
+eval_dataset_path=data/iterative-prompt/eval
+output_dir=output_models/${run_name}
+
+while [[ $# -ge 1 ]]; do
+  key="$1"
+  case ${key} in
+    -r|--run_name)
+      run_name="$2"
+      shift
+      ;;
+    --model_name_or_path)
+      model_name_or_path="$2"
+      shift
+      ;;
+    --reference_model_name_or_path)
+      reference_model_name_or_path="$2"
+      shift
+      ;;
+    --dataset_path)
+      dataset_path="$2"
+      shift
+      ;;
+    --eval_dataset_path)
+      eval_dataset_path="$2"
+      shift
+      ;;
+    -o|--output_dir)
+      output_dir="$2"
+      shift
+      ;;
+    *)
+      echo "error: unknown option \"${key}\"" 1>&2
+      exit 1
+  esac
+  shift
+done
+
+project_dir=$(cd "$(dirname $0)"/..; pwd)
+log_dir=${project_dir}/log/${run_name}
+mkdir -p ${output_dir} ${log_dir}
+
+accelerate launch --config_file configs/accelerate_dsz3_config.yaml \
+  examples/dpov2_train.py \
+    --model_name_or_path ${model_name_or_path} \
+    --reference_model_name_or_path ${reference_model_name_or_path} \
+    --do_train True \
+    --dataset_path ${dataset_path} \
+    --eval_dataset_path ${eval_dataset_path} \
+    --bf16 True \
+    --learning_rate 5e-7 \
+    --lr_scheduler_type cosine \
+    --warmup_steps 100 \
+    --optim paged_adamw_32bit \
+    --per_device_train_batch_size 1 \
+    --per_device_eval_batch_size 1 \
+    --gradient_accumulation_steps 16 \
+    --gradient_checkpointing True \
+    --margin_scale 1.0 \
+    --max_prompt_length 1000 \
+    --num_train_epochs 2 \
+    --logging_steps 2 \
+    --save_strategy epoch \
+    --save_steps 5000 \
+    --evaluation_strategy steps \
+    --eval_steps 100 \
+    --loss_type sigmoid \
+    --output_dir ${output_dir} \
+    --run_name ${run_name} \
+    --sampling_paired_method max_min \
+    --report_to wandb \
+    --mask_prompt True \
+    --length_penalty 0 \
+    | tee ${log_dir}/train.log \
+    2> ${log_dir}/train.err

scripts/run_rm_inference.sh

@@ -0,0 +1,73 @@
+#!/bin/bash
+# Copyright 2024 Statistics and Machine Learning Research Group. All rights reserved.
+
+# Parses arguments
+run_name=rm_inference
+model_name_or_path=sfairXC/FsfairX-LLaMA3-RM-v0.1
+dataset_path=data/alpaca/test
+output_dir=data/rm_inference_results
+output_file_name=results.json
+conversation_template=llama3
+
+# Safety related arguments
+trust_remote_code=0
+
+while [[ $# -ge 1 ]]; do
+  key="$1"
+  case ${key} in
+    -r|--run_name)
+      run_name="$2"
+      shift
+      ;;
+    -m|--model_name_or_path)
+      model_name_or_path="$2"
+      shift
+      ;;
+    -d|--dataset_path)
+      dataset_path="$2"
+      shift
+      ;;
+    --conversation_template)
+      conversation_template="$2"
+      shift
+      ;;
+    --output_dir)
+      output_dir="$2"
+      shift
+      ;;
+    --output_file_name)
+      output_file_name="$2"
+      shift
+      ;;
+    --trust_remote_code)
+      trust_remote_code="$2"
+      shift
+      ;;
+    *)
+      echo "error: unknown option \"${key}\"" 1>&2
+      exit 1
+  esac
+  shift
+done
+
+# inference
+project_dir=$(cd "$(dirname $0)"/..; pwd)
+log_dir=${project_dir}/log/${run_name}
+output_file_path=${output_dir}/${run_name}/${output_file_name}
+mkdir -p ${output_dir}/${run_name} ${log_dir}
+
+accelerate launch --config_file configs/accelerator_multigpu_config.yaml \
+    examples/rm_inference.py \
+        --trust_remote_code ${trust_remote_code} \
+        --model_name_or_path ${model_name_or_path} \
+        --arch_type text_regression \
+        --use_accelerator True \
+        --block_size 4096 \
+        --inference_batch_size 16 \
+        --dataset_path ${dataset_path} \
+        --overwrite_cache True \
+        --conversation_template ${conversation_template} \
+        --preprocessing_num_workers 16 \
+        --save_results True \
+        --results_path ${output_file_path} \
+        2>&1 | tee ${log_dir}/rm_inference.log