update

examples/tutorials/dpo/ultrafeedback-dpo/step_1_prepare_data.py

@@ -47,11 +47,11 @@ def main():
         local_dir=args.local_dir,
     )
     #huggingface_hub
-    snapshot_download(
-        repo_type="model",
-        repo_id=args.repo_id,
-        local_dir=args.local_dir,
-    )
+    # snapshot_download(
+    #     repo_type="model",
+    #     repo_id=args.repo_id,
+    #     local_dir=args.local_dir,
+    # )
     return
 
 

examples/tutorials/dpo/ultrafeedback-dpo/step_2_train_dpo_model_ddp_qlora.py

@@ -14,7 +14,8 @@ DPO本来就是风格微调，用LoRA 训练更合理，更科学。
 ----------
 
 nohup torchrun --nproc_per_node=2 step_2_train_dpo_model_ddp_qlora.py \
---dpo_beta 0.5 \
+--learning_rate 5e-5
+--dpo_beta 0.05 \
 --lora_rank 32 \
 &
 
@@ -79,6 +80,7 @@ def get_args():
         type=str
     ),
 
+    parser.add_argument("--learning_rate", default=2e-5, type=float),
     parser.add_argument("--dpo_beta", default=0.5, type=float),
     parser.add_argument("--lora_rank", default=32, type=int),
 
@@ -229,7 +231,7 @@ def main():
         save_steps=100,
         save_total_limit=2,
         logging_steps=10,
-        learning_rate=2e-5,
+        learning_rate=args.learning_rate,
         warmup_steps=100,
         lr_scheduler_type="cosine",
         fp16=True,

examples/tutorials/grpo/requirements.txt

@@ -0,0 +1,2 @@
+trl==0.28.0
+transformers

examples/tutorials/grpo/step_1_download_model_ms.py

@@ -0,0 +1,49 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+或使用命令行
+
+python3 step_1_prepare_data.py \
+--repo_id qgyd2021/gpt2-for-sequence-classification-sst2-reward \
+--local_dir /root/autodl-tmp/OpenMiniMind/trained_models/gpt2-for-sequence-classification-sst2-reward
+
+"""
+import argparse
+import os
+from pathlib import Path
+import platform
+
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = os.path.abspath("../../../")
+    project_path = Path(project_path)
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+from modelscope import snapshot_download
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repo_id", default="Qwen/Qwen2.5-3B-Instruct", type=str)
+    parser.add_argument(
+        "--local_dir",
+        default=(temp_directory / "../pretrained_models/Qwen/Qwen2.5-3B-Instruct").as_posix(),
+        type=str
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    snapshot_download(
+        model_id=args.repo_id,
+        local_dir=args.local_dir,
+    )
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/tutorials/grpo/step_2_train_grpo_model.py

@@ -2,10 +2,301 @@
 # -*- coding: utf-8 -*-
 """
 
+https://huggingface.co/docs/trl/v0.28.0/en/grpo_trainer
+
+
 """
+import argparse
+import os
+from pathlib import Path
+import platform
+import re
+from typing import Any, Dict, List, Optional
+
+os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+
+# 路径配置
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = Path(os.path.abspath("../../../../"))
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+from datasets import load_dataset
+import torch
+from transformers import (
+    AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification,
+    GPT2LMHeadModel, GPT2ForSequenceClassification,
+    DataCollatorWithPadding
+)
+from trl import GRPOConfig, GRPOTrainer
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_name",
+        # default="Qwen/Qwen2.5-3B-Instruct",
+        # default=(project_path / "pretrained_models/Qwen/Qwen2.5-3B-Instruct").as_posix(),
+        default=(project_path / "pretrained_models/jingyaogong/MiniMind2").as_posix(),
+        type=str
+    )
+    parser.add_argument("--dataset_path", default="Jiayi-Pan/Countdown-Tasks-3to4", type=str)
+    parser.add_argument("--dataset_cache_dir",
+                        default=(temp_directory / "hub_datasets").as_posix(), type=str)
+    parser.add_argument("--model_cache_dir",
+                        default=(temp_directory / "hub_models").as_posix(), type=str)
+
+    # 训练参数
+    parser.add_argument("--valid_dataset_size", default=2000, type=int)
+
+    # 生成参数
+
+    parser.add_argument(
+        "--output_model_dir",
+        default=(project_path / "trained_models/qwen2_5-3B-Instruct-Countdown-GRPO").as_posix(),
+        type=str
+    ),
+
+    # 其他
+    parser.add_argument("--num_workers", default=0 if platform.system() == "Windows" else 2, type=int)
+    return parser.parse_args()
+
+
+SYSTEM_MESSAGE = """
+你是一个有用的助手。你首先在脑海中思考推理过程，然后为用户提供答案。
+""".strip()
+
+USER_TEMPLATE = (
+    "使用这些数字 {numbers}，创建一个等于 {target} 的等式。"
+    "你可以使用基本算术运算（+、-、*、/），每个数字只能使用一次。"
+    "在 <think> </think> 标签中展示你的解题过程。"
+    "并在 <answer> </answer> 标签中返回最终答案，例如 <answer> (1 + 2) / 3 </answer>。"
+)
+
+RESPONSE_PROMPT = "让我一步步来解决这个问题。\n<think>"
+
+
+def format_func(example, tokenizer):
+    numbers: List[int] = example["nums"]
+    target: int = example["target"]
+
+    user_message = USER_TEMPLATE.format(
+        numbers=numbers,
+        target=target
+    )
+
+    messages = [
+        {"role": "system", "content": SYSTEM_MESSAGE},
+        {"role": "user", "content": user_message},
+    ]
+
+    formatted_prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+    formatted_prompt = formatted_prompt + RESPONSE_PROMPT
+
+    tokenized = tokenizer(formatted_prompt,)
+    input_ids = tokenized["input_ids"]
+    attention_mask = tokenized["attention_mask"]
+    result = {
+        "prompt": formatted_prompt,
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+
+        "numbers": numbers,
+        "target": target,
+    }
+    return result
+
+
+def format_reward_function(
+        completions: str,
+        ) -> float:
+    """
+    检查模型的回复是否符合格式 <think>...</think><answer>...</answer>
+    """
+    think_regex = r"<think>.*?<\/think>"
+    answer_regex = r"<answer>.*?<\/answer>"
+    full_format_regex = r"^<think>.*?<\/think>\n<answer>.*?<\/answer>$"
+
+    think_match = re.search(think_regex, completions, re.DOTALL)
+    answer_match = re.search(answer_regex, completions, re.DOTALL)
+    full_format_match = re.match(
+        full_format_regex,
+        completions,
+        re.DOTALL
+    )
+    # 如果完全匹配，则给1分
+    if full_format_match:
+        return 1.0
+    reward = 0.0
+    # 如果有<think></think>标签对，则奖励加0.1分
+    if think_match:
+        reward += 0.1
+    # 如果有<answer></answer>标签对，则奖励加0.5分
+    if answer_match:
+        reward += 0.5
+    # 返回奖励
+    return reward
+
+
+def format_reward_func(
+        prompts: List[str],
+        completions: List[str],
+        completion_ids: List[List[int]],
+        # end_token: Optional[str] = None,
+        **kwargs) -> List[float]:
+    result = list()
+    for completion in completions:
+        reward = format_reward_function(completion)
+        result.append(reward)
+    return result
+
+
+def answer_reward_function(
+    response: str,
+    numbers: List[int] = None,
+    target: int = None
+) -> float:
+    """
+    检查答案中：
+    1. 是否使用了所有给的数字
+    2. 每个数字是否使用了一次
+    3. 答案中包含的表达式的求��结果是否等于目标数字
+    """
+    # 答案的正则表达式
+    answer_regex = r"<answer>(.*?)<\/answer>"
+    # 回答中是否有答案标签对
+    answer_match = re.search(answer_regex, response, re.DOTALL)
+    # 如果在回答中没有搜索到答案，那么给0分
+    if not answer_match:
+        return 0.0
+    # 提取出答案的文本
+    answer_content = answer_match.group(1)
+    # 如果答案标签内没有东西，给0分
+    if not answer_content:
+        return 0.0
+    # 如果答案标签中，除了表达式以外，还有其它内容，给0分
+    allowed_chars = r"^[0-9+\-*/() ]+$"
+    if not re.match(allowed_chars, answer_content):
+        return 0.0
+    # 检查答案中，每个数字是否只使用了一次
+    used_numbers = [
+        int(n) for n in re.findall(r"\d+", answer_content)
+    ]
+    if sorted(used_numbers) != sorted(numbers):
+        return 0.0
+    # 检查答案中包含的表达式的求值结果是否为目标数字
+    try:
+        result = eval(answer_content, {"__builtins__": None}, {})
+        if abs(float(result) - float(target)) < 1e-5:
+            return 1.0
+    except:
+        pass
+    return 0.0
+
+
+def answer_reward_func(
+        prompts: List[str],
+        completions: List[str],
+        completion_ids: List[List[int]],
+        **kwargs) -> List[float]:
+    target_list = kwargs["target"]
+    numbers_list = kwargs["numbers"]
+
+    result = list()
+    for completion, numbers, target in zip(completions, numbers_list, target_list):
+        reward = answer_reward_function(completion, numbers, target)
+        result.append(reward)
+    return result
+
+
+def main():
+    args = get_args()
+
+    model = AutoModelForCausalLM.from_pretrained(
+        pretrained_model_name_or_path=args.model_name,
+        trust_remote_code=True,
+        cache_dir=args.model_cache_dir,
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model_name,
+        padding_side="left",  # 对于生成任务很重要
+        cache_dir=args.model_cache_dir,
+    )
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+    print(f"eos_token: {tokenizer.eos_token}")
+    print(f"pad_token: {tokenizer.pad_token}")
+
+    dataset_dict = load_dataset(
+        path=args.dataset_path,
+        cache_dir=args.dataset_cache_dir,
+    )
+    dataset = dataset_dict["train"]
+    dataset = dataset.take(n=10000)
+
+    dataset = dataset.train_test_split(test_size=args.valid_dataset_size, seed=None)
+    train_dataset = dataset["train"]
+    valid_dataset = dataset["test"]
+
+    train_dataset = train_dataset.map(
+        lambda example: format_func(example, tokenizer),
+        batched=False,
+        remove_columns=train_dataset.column_names,
+    )
+    valid_dataset = valid_dataset.map(
+        lambda example: format_func(example, tokenizer),
+        batched=False,
+        remove_columns=valid_dataset.column_names,
+    )
+
+    grpo_config = GRPOConfig(
+        output_dir=args.output_model_dir,
+        num_train_epochs=3,
+        per_device_train_batch_size=4,
+        per_device_eval_batch_size=4,
+        gradient_accumulation_steps=4,
+        learning_rate=5e-6,  # GRPO通常使用较小的学习率
+        warmup_ratio=0.1,
+        logging_steps=10,
+        eval_strategy="steps",
+        eval_steps=100,
+        save_strategy="steps",
+        save_steps=100,
+        save_total_limit=3,
+        load_best_model_at_end=True,
+        metric_for_best_model="reward",  # 使用奖励作为评估指标
+        greater_is_better=True,
+        fp16=False,
+        bf16=False,
+        max_grad_norm=1.0,
+        report_to="none",  # 可根据需要改为"wandb"等
+        # GRPO特有参数
+        num_generations=4,  # 每个提示生成的响应数量
+        temperature=0.7,
+        max_completion_length=512,  # 生成的最大长度
+        reward_weights=[0.1, 1.0],
+    )
 
+    grpo_trainer = GRPOTrainer(
+        model=model,
+        processing_class=tokenizer,
+        args=grpo_config,
+        train_dataset=train_dataset,
+        eval_dataset=valid_dataset,
+        reward_funcs=[format_reward_func, answer_reward_func],
+    )
+    grpo_trainer.train()
+    grpo_trainer.save_model(args.output_model_dir)
 
+    return
 
 
 if __name__ == "__main__":
-    pass
+    main()