Merge branch 'dev'

Author: TideDra

.pre-commit-config.yaml

@@ -8,7 +8,7 @@ ci:
 
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.5.0
+    rev: v5.0.0
     hooks:
       - id: check-yaml
       - id: check-case-conflict
@@ -18,20 +18,20 @@ repos:
       - id: requirements-txt-fixer
 
   - repo: https://github.com/PyCQA/autoflake
-    rev: v2.0.2
+    rev: v2.3.1
     hooks:
       - id: autoflake
         args: [--remove-all-unused-imports, --in-place]
 
   - repo: https://github.com/PyCQA/isort
-    rev: 5.13.2
+    rev: 6.0.1
     hooks:
       - id: isort
         name: Format imports
         exclude: docs/
 
   - repo: https://github.com/psf/black
-    rev: 24.3.0
+    rev: 25.1.0
     hooks:
       - id: black
         name: Format code

dockerfile/Dockerfile

@@ -15,7 +15,7 @@ RUN DEBIAN_FRONTEND=noninteractive apt install -y tzdata
 
 RUN apt-get -y install build-essential git python3-dev python3-pip libopenexr-dev libxi-dev libglfw3-dev libglew-dev libomp-dev libxinerama-dev libxcursor-dev gdb
 RUN pip uninstall xgboost transformer_engine flash_attn pynvml opencv-python-headless -y
-RUN pip install vllm==0.7.2
+RUN pip install vllm==0.8.3
 
 COPY docker-entrypoint.sh .
 RUN chmod a+x docker-entrypoint.sh

examples/scripts/experience_filter.py

@@ -0,0 +1,4 @@
+import torch
+
+def experience_filter(experience_maker,experiences):
+    return experiences

examples/scripts/train_grpo_ray_hybrid_engine.sh

@@ -26,8 +26,8 @@ ray job submit --address="http://127.0.0.1:8265" \
    --micro_train_batch_size 4 \
    --train_batch_size 128 \
    --micro_rollout_batch_size 8 \
-   --rollout_batch_size 1024 \
-   --n_samples_per_prompt 1 \
+   --rollout_batch_size 128 \
+   --n_samples_per_prompt 8 \
    --max_epochs 1 \
    --prompt_max_len 1024 \
    --max_samples 100000 \

examples/scripts/train_ppo_llama_with_dynamic_sampling.sh

@@ -0,0 +1,43 @@
+set -x
+
+read -r -d '' training_commands <<EOF
+openrlhf.cli.train_ppo \
+   --pretrain OpenRLHF/Llama-3-8b-sft-mixture \
+   --reward_pretrain OpenRLHF/Llama-3-8b-rm-mixture \
+   --custom_experience_filter examples/scripts/experience_filter.py \
+   --save_path ./checkpoint/llama-3-8b-rlhf \
+   --save_steps -1 \
+   --logging_steps 1 \
+   --eval_steps -1 \
+   --micro_train_batch_size 2 \
+   --train_batch_size 128 \
+   --micro_rollout_batch_size 4 \
+   --rollout_batch_size 1024 \
+   --max_steps 1000 \
+   --store_extra_buffers \
+   --max_epochs 1 \
+   --prompt_max_len 1024 \
+   --generate_max_len 1024 \
+   --zero_stage 2 \
+   --bf16 \
+   --actor_learning_rate 5e-7 \
+   --critic_learning_rate 9e-6 \
+   --init_kl_coef 0.01 \
+   --prompt_data OpenRLHF/prompt-collection-v0.1 \
+   --input_key context_messages \
+   --apply_chat_template \
+   --max_samples 100000 \
+   --normalize_reward \
+   --adam_offload \
+   --flash_attn \
+   --load_checkpoint \
+   --gradient_checkpointing
+EOF
+
+    # --packing_samples
+    # --use_wandb [WANDB_TOKENS] or True (use wandb login command)
+    # --remote_rm_url http://localhost:5000/get_reward
+
+if [[ ${1} != "slurm" ]]; then
+    deepspeed --module $training_commands
+fi

examples/scripts/train_reinforce_baseline_llama_ray_hybrid_engine.sh

@@ -0,0 +1,54 @@
+set -x
+
+ray job submit --address="http://127.0.0.1:8265" \
+   --runtime-env-json='{"working_dir": "/openrlhf"}' \
+   -- python3 -m openrlhf.cli.train_ppo_ray \
+   --ref_num_nodes 1 \
+   --ref_num_gpus_per_node 8 \
+   --reward_num_nodes 1 \
+   --reward_num_gpus_per_node 8 \
+   --actor_num_nodes 1 \
+   --actor_num_gpus_per_node 8 \
+   --vllm_num_engines 4 \
+   --vllm_tensor_parallel_size 2 \
+   --colocate_all_models \
+   --vllm_gpu_memory_utilization 0.6 \
+   --advantage_estimator reinforce_baseline \
+   --pretrain OpenRLHF/Llama-3-8b-sft-mixture \
+   --reward_pretrain OpenRLHF/Llama-3-8b-rm-700k \
+   --save_path /openrlhf/examples/test_scripts/final/llama3-8b-rlhf \
+   --ckpt_path /openrlhf/examples/test_scripts/ckpt/llama3-8b-rlhf \
+   --save_hf_ckpt \
+   --micro_train_batch_size 4 \
+   --train_batch_size 128 \
+   --micro_rollout_batch_size 8 \
+   --rollout_batch_size 128 \
+   --n_samples_per_prompt 8 \
+   --init_kl_coef 1e-3 \
+   --gamma 1.0 \
+   --use_kl_loss \
+   --kl_estimator k2 \
+   --max_epochs 1 \
+   --prompt_max_len 1024 \
+   --max_samples 100000 \
+   --generate_max_len 1024 \
+   --zero_stage 3 \
+   --bf16 \
+   --actor_learning_rate 5e-7 \
+   --critic_learning_rate 9e-6 \
+   --prompt_data OpenRLHF/prompt-collection-v0.1 \
+   --input_key context_messages \
+   --apply_chat_template \
+   --normalize_reward \
+   --gradient_checkpointing \
+   --packing_samples \
+   --vllm_sync_backend nccl \
+   --enforce_eager \
+   --vllm_enable_sleep \
+   --deepspeed_enable_sleep
+
+# You could also try
+#   --use_kl_loss \
+#   --kl_estimator k3 | k2 \
+
+# also supports --advantage_estimator rloo | reinforce_baseline

openrlhf/cli/batch_inference.py

@@ -54,9 +54,7 @@ class Empty:
         args.dataset_probs,
         dummy_strategy,
         args.seed,
-        return_eval=False,
         max_count=args.max_samples,
-        train_split=args.dataset_split,
     )
     if args.iter is None:
         prompts_data = prompts_data.select(range(min(args.max_samples, len(prompts_data))))
@@ -126,9 +124,7 @@ def tokenize_fn(texts):
         args.dataset_probs,
         strategy,
         args.seed,
-        return_eval=False,
         max_count=args.max_samples,
-        train_split=args.dataset_split,
     )
     if args.iter is None:
         prompts_data = prompts_data.select(range(min(args.max_samples, len(prompts_data))))
@@ -229,9 +225,7 @@ def batch_rm_inference(args):
         args.dataset_probs,
         strategy,
         args.seed,
-        return_eval=False,
         max_count=args.max_samples,
-        train_split=args.dataset_split,
     )
     dataset = dataset.select(range(min(args.max_samples, len(dataset))))
     dataset = SFTDataset(
@@ -316,8 +310,7 @@ def batch_rm_inference(args):
 
     # Custom dataset
     parser.add_argument("--dataset", type=str, default=None)
-    parser.add_argument("--dataset_probs", type=str, default="1.0")
-    parser.add_argument("--dataset_split", type=str, default="train")
+    parser.add_argument("--dataset_probs", type=str, default=None)
     parser.add_argument("--input_key", type=str, default="input", help="JSON dataset key")
     parser.add_argument("--output_key", type=str, default="output", help="JSON dataset key")
     parser.add_argument(

openrlhf/cli/lora_combiner.py

@@ -2,12 +2,13 @@
 
 import torch
 from peft import PeftModel
-from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer
-
+from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer, AutoConfig
+from openrlhf.models.lmm_kits.utils import get_generation_cls
 
 def apply_lora(model_name_or_path, lora_path, output_path, is_rm, bf16):
     print(f"Loading the base model from {model_name_or_path}")
-    model_cls = AutoModelForCausalLM if not is_rm else AutoModelForSequenceClassification
+    config = AutoConfig.from_pretrained(model_name_or_path)
+    model_cls = get_generation_cls(config) if not is_rm else AutoModelForSequenceClassification
     base = model_cls.from_pretrained(
         model_name_or_path, torch_dtype=torch.bfloat16 if bf16 else "auto", low_cpu_mem_usage=True
     )

openrlhf/cli/train_dpo.py

@@ -60,18 +60,15 @@ def train(args):
     optim = strategy.create_optimizer(model, lr=args.learning_rate, betas=args.adam_betas, weight_decay=args.l2)
 
     # prepare for data and dataset
-    train_data, eval_data = blending_datasets(
+    train_data = blending_datasets(
         args.dataset,
         args.dataset_probs,
         strategy,
         args.seed,
         max_count=args.max_samples,
-        stopping_strategy="all_exhausted",
-        train_split=args.train_split,
-        eval_split=args.eval_split,
     )
+
     train_data = train_data.select(range(min(args.max_samples, len(train_data))))
-    eval_data = eval_data.select(range(min(args.max_samples, len(eval_data))))
     train_dataset = RewardDataset(
         train_data,
         tokenizer,
@@ -81,32 +78,40 @@ def train(args):
         is_dpo=True,
         multiple_of=args.ring_attn_size,
     )
-    eval_dataset = RewardDataset(
-        eval_data,
-        tokenizer,
-        args.max_len,
-        strategy,
-        input_template=args.input_template,
-        is_dpo=True,
-        multiple_of=args.ring_attn_size,
-    )
 
     # prepare dataloader
     train_dataloader = strategy.setup_dataloader(
         train_dataset,
         args.micro_train_batch_size,
         True,
         True,
-        train_dataset.packing_collate_fn if args.packing_samples else train_dataset.collate_fn,
+        train_dataset.collate_fn,
     )
 
-    eval_dataloader = strategy.setup_dataloader(
-        eval_dataset,
-        args.micro_train_batch_size,
-        True,
-        False,
-        eval_dataset.packing_collate_fn if args.packing_samples else eval_dataset.collate_fn,
-    )
+    eval_dataset = None
+    eval_dataloader = None
+    if getattr(args, "eval_dataset", None):
+        eval_data = blending_datasets(
+            args.eval_dataset,
+            None,  # No probability sampling for eval datasets
+            strategy,
+        )
+        eval_dataset = RewardDataset(
+            eval_data,
+            tokenizer,
+            args.max_len,
+            strategy,
+            input_template=args.input_template,
+            is_dpo=True,
+            multiple_of=args.ring_attn_size,
+        )
+        eval_dataloader = strategy.setup_dataloader(
+            eval_dataset,
+            args.micro_train_batch_size,
+            True,
+            False,
+            eval_dataset.collate_fn,
+        )
 
     # scheduler
     num_update_steps_per_epoch = len(train_dataset) // args.train_batch_size
@@ -168,7 +173,7 @@ def train(args):
     parser.add_argument("--ckpt_path", type=str, default="./ckpt/checkpoints_dpo")
     parser.add_argument("--max_ckpt_num", type=int, default=3)
     parser.add_argument("--max_ckpt_mem", type=int, default=1e8)
-    parser.add_argument("--universal_ckpt", action="store_true", default=False)
+    parser.add_argument("--use_ds_universal_ckpt", action="store_true", default=False)
 
     # DeepSpeed
     parser.add_argument("--micro_train_batch_size", type=int, default=8, help="batch size per GPU")
@@ -236,10 +241,10 @@ def train(args):
     # Custom dataset
     parser.add_argument("--pretrain", type=str, default=None)
     parser.add_argument("--ref_pretrain", type=str, default=None)
-    parser.add_argument("--dataset", type=str, default=None)
-    parser.add_argument("--dataset_probs", type=str, default="1.0", help="sampling probs for datasets")
-    parser.add_argument("--train_split", type=str, default="train", help="train split of the HF dataset")
-    parser.add_argument("--eval_split", type=str, default="test", help="test split of the dataset")
+    parser.add_argument("--dataset", type=str, default=None, help="Path to the training dataset")
+    parser.add_argument("--dataset_probs", type=str, default=None, help="Sampling probabilities for training datasets")
+    parser.add_argument("--eval_dataset", type=str, default=None, help="Path to the evaluation dataset")
+    parser.add_argument("--max_samples", type=int, default=1000000, help="Maximum number of samples to use")
 
     parser.add_argument("--prompt_key", type=str, default=None)
     parser.add_argument("--chosen_key", type=str, default="chosen")
@@ -248,7 +253,6 @@ def train(args):
     parser.add_argument(
         "--apply_chat_template", action="store_true", default=False, help="Use HF tokenizer chat template"
     )
-    parser.add_argument("--max_samples", type=int, default=1e8, help="Max number of samples")
     parser.add_argument("--max_len", type=int, default=512)
 
     # wandb parameters