Merge branch 'PaddlePaddle:develop' into add_sharded_state_dict

Author: xingmingyyj

.github/workflows/lint.yml

@@ -72,8 +72,6 @@ jobs:
           else
             echo "local develop branch exist, skipping"
           fi
-
-          unset http_proxy && unset https_proxy
           '
           
       - name: Setup Environment
@@ -84,6 +82,7 @@ jobs:
           set -e
           python -m pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
           python -m pip config set global.cache-dir "/home/.cache/pip"
+          source $work_dir/../../../proxy
           python -m pip install --upgrade pip
           cd /workspace/PaddleNLP && git config --global --add safe.directory $PWD
           make install
@@ -93,6 +92,7 @@ jobs:
         run: |
           docker exec -t $container_name /bin/bash -c '
           set -e
+          source $work_dir/../../../proxy
           cd /workspace/PaddleNLP
           make lint
           '

.github/workflows/unittest-cpu.yml

@@ -11,6 +11,11 @@ env:
   COMMIT_ID: ${{ github.event.pull_request.head.sha }}
   BRANCH: ${{ github.event.pull_request.base.ref }}
   TASK: PaddleNLP-CI-${{ github.event.pull_request.number }}-unittest-cpu
+  HF_ENDPOINT: https://hf-mirror.com
+  STUDIO_GIT_HOST: http://git.prod.idc-to-cloud.aistudio.baidu-int.com
+  PPNLP_HOME: /home/disk1/cache
+  HF_DATASETS_CACHE: /home/disk1/cache/huggingface/datasets
+  TRANSFORMERS_CACHE: /home/disk1/cache/huggingface
 
 jobs:
   Test:
@@ -33,12 +38,18 @@ jobs:
             -v $work_dir/../../..:$work_dir/../../.. \
             -v $work_dir:/workspace \
             -v /home/.cache/pip:/home/.cache/pip \
+            -v /home/disk1/cache:/home/disk1/cache \
             -e BRANCH \
             -e PR_ID \
             -e COMMIT_ID \
             -e work_dir \
             -e no_proxy \
             -e python_version \
+            -e HF_ENDPOINT \
+            -e STUDIO_GIT_HOST \
+            -e PPNLP_HOME \
+            -e HF_DATASETS_CACHE \
+            -e TRANSFORMERS_CACHE \
             -w /workspace ${docker_image}
 
       - name: Download Code
@@ -68,7 +79,6 @@ jobs:
             echo "Not in a pull_request event. Skipping PR-specific operations."
           fi
           git log --pretty=oneline -10
-          unset http_proxy && unset https_proxy
           '
 
       - name: Setup Environment
@@ -79,6 +89,7 @@ jobs:
           set -e
           python -m pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
           python -m pip config set global.cache-dir "/home/.cache/pip"
+          source $work_dir/../../../proxy
           python -m pip install --upgrade pip
           cd /workspace/PaddleNLP && git config --global --add safe.directory $PWD
           pip install -r tests/requirements.txt

llm/alignment/rl/gsm8k_processor.py

@@ -0,0 +1,74 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Preprocess the GSM8k dataset to parquet format
+"""
+
+import argparse
+import os
+import re
+
+import datasets
+
+
+def extract_solution(solution_str):
+    solution = re.search("#### (\\-?[0-9\\.\\,]+)", solution_str)
+    assert solution is not None
+    final_solution = solution.group(0)
+    final_solution = final_solution.split("#### ")[1].replace(",", "")
+    return final_solution
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--local_dir", default="./gsm8k")
+
+    args = parser.parse_args()
+
+    data_source = "openai/gsm8k"
+
+    dataset = datasets.load_dataset(data_source, "main")
+
+    train_dataset = dataset["train"]
+    test_dataset = dataset["test"]
+
+    instruction_following = 'Let\'s think step by step and output the final answer after "####".'
+
+    # add a row to each data item that represents a unique id
+    def make_map_fn(split):
+        def process_fn(example, idx):
+            question_raw = "<|im_start|>user\n" + example.pop("question")
+
+            system_raw = (
+                "<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n"
+            )
+            question = system_raw + question_raw + " " + instruction_following + "<|im_end|>\n<|im_start|>assistant\n"
+
+            answer_raw = example.pop("answer")
+            solution = extract_solution(answer_raw)
+            data = {
+                "src": question,
+                "tgt": solution,
+            }
+            return data
+
+        return process_fn
+
+    train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True)
+    test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True)
+
+    local_dir = args.local_dir
+
+    train_dataset.to_json(os.path.join(local_dir, "train.jsonl"), orient="records", lines=True)
+    test_dataset.to_json(os.path.join(local_dir, "test.jsonl"), orient="records", lines=True)

llm/alignment/rl/run_rl.py

@@ -42,6 +42,7 @@
 from paddlenlp.transformers import (
     AutoConfig,
     AutoModelForCausalLM,
+    AutoModelForTokenClassification,
     AutoTokenizer,
     PretrainedConfig,
 )
@@ -134,7 +135,6 @@ def create_actor_models(
         )
         if not training_args.autotuner_benchmark:
             reference_model.set_state_dict(actor_model.state_dict())
-
     actor_tokenizer = AutoTokenizer.from_pretrained(
         model_args.actor_model_name_or_path,
         model_max_length=data_args.max_length,
@@ -210,46 +210,43 @@ def create_critic_models(
     data_args: DataArgument,
     training_args: TrainingArguments,
     common_config: Dict,
-    reward_model,
 ):
     with timers_scope_runtimer("Critic model loading time"):
-        reward_model_config = reward_model.config
-        if model_args.critic_model_name_or_path is None:
-            model_args.critic_model_name_or_path = model_args.reward_model_name_or_path
-            critic_model = AutoModelForScore.from_config(
-                reward_model_config,
-                dtype=training_args.model_dtype,
-                score_type="critic",
-                do_normalize=False,
-                clip_range_value=training_args.clip_range_value,
-                **common_config,
+        critic_model_config = AutoConfig.from_pretrained(
+            model_args.critic_model_name_or_path,
+            tensor_parallel_output=training_args.tensor_parallel_output,
+            tensor_parallel_degree=training_args.tensor_parallel_degree,
+            tensor_parallel_rank=training_args.tensor_parallel_rank,
+            dtype=training_args.model_dtype,
+            recompute=training_args.critic_recompute,
+            recompute_granularity=model_args.critic_recompute_granularity,
+            recompute_use_reentrant=training_args.recompute_use_reentrant,
+            **common_config,
+        )
+        LlmMetaConfig.set_llm_config(critic_model_config, training_args)
+
+        critic_model_config.max_position_embeddings = data_args.max_length
+        critic_model_config.use_sparse_head_and_loss_fn = False
+        critic_model_config.num_labels = 1
+        critic_model_config.classifier_dropout = 0.0
+        critic_model_config.hidden_dropout = 0.0
+        logger.info(f"Loading Critic model with config:\n\t{critic_model_config}\n")
+
+        if not training_args.autotuner_benchmark:
+            critic_model = AutoModelForTokenClassification.from_pretrained(
+                model_args.critic_model_name_or_path,
+                config=critic_model_config,
             )
-            if not training_args.autotuner_benchmark:
-                critic_model.set_state_dict(reward_model.state_dict())
         else:
-            if not training_args.autotuner_benchmark:
-                critic_model = AutoModelForScore.from_pretrained(
-                    model_args.critic_model_name_or_path,
-                    config=reward_model_config,
-                    score_type="critic",
-                    do_normalize=False,
-                    clip_range_value=training_args.clip_range_value,
-                    **common_config,
-                )
-            else:
-                critic_model = AutoModelForScore.from_config(
-                    reward_model_config,
-                    score_type="critic",
-                    do_normalize=False,
-                    clip_range_value=training_args.clip_range_value,
-                    **common_config,
-                )
+            critic_model = AutoModelForTokenClassification.from_config(
+                critic_model_config,
+            )
 
     critic_tokenizer = AutoTokenizer.from_pretrained(
         model_args.critic_model_name_or_path,
         model_max_length=data_args.max_length,
         padding_side="left",
-        tokenizer_alpha=model_args.reward_critic_tokenizer_alpha,
+        tokenizer_alpha=model_args.critic_tokenizer_alpha,
         use_fast=True,
     )
     if critic_tokenizer.pad_token_id is None:
@@ -261,16 +258,16 @@ def create_critic_models(
         if training_args.eval_mode == "single":
             config.tensor_parallel_degree = -1
             config.tensor_parallel_rank = 0
-        with timers_scope_runtimer("Reward critic eval model loading time"):
-            critic_eval_model = AutoModelForScore.from_config(config)
+        with timers_scope_runtimer("Critic eval model loading time"):
+            critic_eval_model = AutoModelForTokenClassification.from_config(config)
     else:
         critic_eval_model = None
 
     return critic_model, critic_eval_model, critic_tokenizer
 
 
 def create_rl_dataset(data_args, training_args, tokenizer):
-    requires_label = True if training_args.use_rm_server else False
+    requires_label = True if training_args.use_rm_server or training_args.use_rule_reward else False
     train_ds = RLHFDataset(
         dataset_name_or_path=data_args.train_datasets,
         tokenizer=tokenizer,
@@ -333,15 +330,16 @@ def main():
     actor_model, actor_eval_model, reference_model, actor_tokenizer = create_actor_models(
         model_args, data_args, training_args, common_config, reshard_controller
     )
-
-    if not training_args.use_rm_server and model_args.reward_model_name_or_path is not None:
+    if training_args.use_rule_reward:
+        reward_model, reward_tokenizer = None, actor_tokenizer
+    elif not training_args.use_rm_server and model_args.reward_model_name_or_path is not None:
         reward_model, reward_tokenizer = create_reward_models(model_args, data_args, training_args, common_config)
     else:
         reward_model, reward_tokenizer = model_args.reward_server, actor_tokenizer
 
     if training_args.rl_algorithm == "ppo":
         critic_model, critic_eval_model, critic_tokenizer = create_critic_models(
-            model_args, data_args, training_args, common_config, reward_model
+            model_args, data_args, training_args, common_config
         )
     else:
         critic_model, critic_eval_model, critic_tokenizer = None, None, None
@@ -355,15 +353,23 @@ def main():
         offload_tensor_to_cpu((reference_model, "freeze_model"))
 
         if training_args.rl_algorithm == "ppo":
-            offload_tensor_to_cpu((reward_model, "freeze_model"))
+            if not training_args.use_rm_server and not training_args.use_rule_reward:
+                offload_tensor_to_cpu((reward_model, "freeze_model"))
             if critic_eval_model is not None:
                 offload_tensor_to_cpu((critic_eval_model, "freeze_model"))
 
         # NOTE(gongenlei): release memory_reserved_size to equal to memory_allocated_size
         paddle.device.cuda.empty_cache()
 
     def compute_metrics(eval_preds):
-        accuracy = (eval_preds.predictions == 3).astype("float32").mean().item()
+        '''
+        If "use_rm_server" is TRUE, the score ranges from -3 to 3, with 3 being the only correct score (format + result).
+        If using the "Regularized Matching Function (use_rule_reward=True)" (currently only implemented for the gsm8k dataset), the score ranges from 0 to 1.
+        '''
+        if training_args.use_rule_reward:
+            accuracy = (eval_preds.predictions == 1).astype("float32").mean().item()
+        else:
+            accuracy = (eval_preds.predictions == 3).astype("float32").mean().item()
         return {"accuracy": accuracy}
 
     try:
@@ -389,7 +395,7 @@ def compute_metrics(eval_preds):
         data_collator=partial(
             collate_fn,
             pad_token_id=actor_tokenizer.pad_token_id,
-            requires_label=True if training_args.use_rm_server else False,
+            requires_label=True if training_args.use_rm_server or training_args.use_rule_reward else False,
             max_prompt_len=data_args.max_prompt_len if training_args.balance_batch else None,
         ),  # NOTE: enforce prompt padding to max_prompt_len when using balance_batch
         compute_metrics=compute_metrics,  # TODO: only used for grpo (kk datasets)