TRL 单卡GRPO

由于Verl安装一直报错，docker都解决不了的Vllm不能初始化，看了官方issue以及源码，实在无法解决，随放弃。改用TRL库

• Model：Qwen-3-0.6B-Base和Qwen-3-0.6B-Instruct

数据准备

数据:K-and-K/knights-and-knaves · Datasets at Hugging Face

• 该数据集为骑士和骗子，给定一段prompt，推理出谁是骑士谁是骗子。骑士用不说谎，骗子反之。

• 数据集切割为了7个子集，安装总人数区分。可是使用split=numppl进行load，当前使用2ppl。

• 数据集会提供solution，names的列表，前者是布尔数组，True代表骑士，False代表Lier。names与solution下标对应，是真实人名。

格式源自verl库，我这里额外处理了SFT的格式，因为R1是先SFT冷启动，R1-Zero是纯粹直接RL

def make_map_fn(split):
    instruction_following = (
        "You are a logical reasoning expert. "
        "Always think step by step inside <think> tags, showing your complete reasoning chain. "
        "After finishing your thinking, enclose your final conclusion in <answer> tags. "
        "Do not put any extra text outside these tags. "
        "Example format:\n"
        "<think>\nStep 1: Analyze Zoey's statement...\nStep 2: Consider possibilities...\n</think>\n"
        "<answer>\nZoey is a knave, Oliver is a knight.\n</answer>"
    )

    def process_fn(example, idx):
        question_raw = example.pop("quiz")
        names = example.pop("names")

        # 这个是Lets think step by step
        cot_head = example.pop("cot_head")
        # 核心思维链
        cot_core = example.pop("cot_repeat_steps")[0]
        # This leads to a feasible solution.
        cot_foot = example.pop("cot_foot")

        answer_raw = example.pop("solution_text")
        question = instruction_following + " " + question_raw + " " + cot_head
        solution = answer_raw

        reference_response = (
            "<think>\n"
            + cot_head
            + "\n"
            + cot_core
            + "\n"
            + cot_foot
            + "\n</think>\n"
            "<answer>\n"
            + answer_raw
            + "\n</answer>"  # 用 solution_text 作为最终答案
        )
        # True,False 数组
        solution = example["solution"]
        names = [name.lower() for name in names]
        ground_truth_roles = {}
        for i in range(len(names)):
            if solution[i]:
                ground_truth_roles[names[i]] = "knight"
            else:
                ground_truth_roles[names[i]] = "knave"
        data = {
            "data_source": split,  # 数据集名称
            "prompt": [
                {
                    "role": "user",
                    "content": question,
                }
            ],
            "response": reference_response,  # sft用的答案
            "ability": "logic",
            "cot": [cot_head, cot_core, cot_foot],
            "names": names,
            "ground_truth": ground_truth_roles,
            # "reward_model": {"style": "rule", "ground_truth": ground_truth_roles},
            # "extra_info": {
            #     "split": split,
            #     "index": idx,
            #     "answer": answer_raw,
            #     "question": question_raw,
            #     "names": names,
            # },
        }
        return data

    return process_fn

奖励函数

该奖励函数是照着verl库的自定义奖励格式编写的，我直接搬过来用力，TRL格式稍有不同，我向外再次封装了一层，去调用下方的compute_score

def compute_score(
    data_source: str,
    solution_str: str,  # 模型的完整回复
    ground_truth: str,
    extra_info: dict = None,
) -> float:
    """
    score 的 Docstring
    输入模型回复，人物名，以及对应骑士的True/False，返回分数

    data_source: 源数据，如 "gsm8k"
    solution_str: 模型生成的答案，纯字符串，仅回复
    ground_truth: 源数据中的答案，纯字符串，仅回复
    extra_info: 额外的信息
    """
    # [False, True, True]

    names = extra_info["names"]
    format, think, answer = check_format(solution_str)

    # 格式正确
    if format[0]:
        parsed = parse_answer(answer, names)
        if not format[1]:  # 格式对了一半
            return 1
        print(parsed)
        print(ground_truth)
        same = is_correct(parsed, ground_truth)
        if same:  # 答案一摸一样
            return 5
        else:
            return 3
    else:
        return -5

• check_format会检查格式，如果是<think> ..<think/> <answer> .. </answer>那返回的布尔列表均为True。如果只有一对标签，则返回一个True，给定一个标签奖励分，因为Base模型能力稍差，一刀切死，很难摸索出奖励。

• 如果标签完全完整，则会对比答案。这里答案处理方法是，规则处理匹配，只要人名与答案对的上即可，忽略大小写。

参数

tranning_args = GRPOConfig(
        output_dir=save_path(),
        per_device_train_batch_size=2,
        gradient_accumulation_steps=2,
        use_vllm=False,  # 为True时，需要自己启动vllm，端口8000
        learning_rate=1e-6,
        adam_beta1=0.9,
        adam_beta2=0.99,
        weight_decay=0.01,
        warmup_ratio=0.1,
        optim="adamw_8bit",
        report_to=["swanlab"],
        generation_batch_size=16,
        logging_steps=1,
        bf16=True,
        gradient_checkpointing=False,
        num_generations=8,
        loss_type='grpo',
        max_steps=500,
        epsilon=0.2,
        epsilon_high=0.28,  # one sided
        # SFTTrainer 有max_length这个，GRPO没有
        max_prompt_length=412,      # 输入（Prompt）的最大长度
        max_completion_length=1000,  # 输出（思维链+答案）的最大生成长度

    )

显存告急，不想梯度检查耗费时间，所以batch设置为2，group为16。在单卡A5000上，跑了一个小时多点左右。显存占用21GB。

实际是近两个小时，为什么后面变快了，因为出现了奖励破解

为什么不用FSDP。因为模型太小，通信时延估计会比多卡带来的加速进行抵消。

这里注意，max_completion_length不能调太小，否则会截断，即使是6亿参数的模型，输出也是大几百的长度。

V1-Res

可以看到模型从200步之后，出现奖励破解了。因为1分的奖励，让模型学会了输出那两个格式化标签，不再思考了，entropy也直接到了近0

推理结果

{"accuracy": 0.04, "has_answer": 0.08, "answer_accuracy": 0.4999993750007813, "total": 100, "answer_total": 8, "answer_acc_count": 4}

100条测试集，只对了4个，有answer标签的只有8个。但是奖励最大值我印象里没有出现过5

后来发现，不知道是什么问题，如果你处理数据那里，返回了"ground_truth": ground_truth_roles字典，这个字典会额外多许多字段，似乎把所有人名都加进去了。

一条本该只有两个人名，可我输出发现，全部人名都有，只是其他名字的值为None

即本该{A:K,B:K}，实际{A:K,B:K,C:K....}

现逻辑已修改

---

此外，我发现数据有不少是输出

Zoey is a knave, Oliver is a knight

3个epoch过拟合了，故改为1个，指令改为下方。

"<think>\nStep 1: Analyze Person A's statement...\n</think>\n"
"<answer>\nPerson A is a knave, Person B is a knight.\n</answer>"

---

解决方案：

• 最先想到的思路是进行长度惩罚，但不能多加损失。[2512.16649] JustRL: Scaling a 1.5B LLM with a Simple RL Recipe，去年年底的这篇论文指出，实际不需要那么多损失，只需要RL，就能激发模型的能力，很多不稳定，都是多个Tricks搞出来的，这篇论文很大道至简，后续会考虑复现一下。就是不知道有没有人蒸馏过0.6B这种小模型。

• 所以我调整了给分原则，之前从-5跳到1，差了6分，模型会认为学到了正确的方向

•     if format[0]:
          parsed = parse_answer(answer, names)
          if not format[1]:  # 格式对了一半
              return -0.3
          print(parsed)
          print(ground_truth)
          same = is_correct(parsed, ground_truth)
          if same:  # 答案一摸一样
              return 5
          else:
              return 0.3
      else:
          return -1
  

V2-Res

结果如下：

• 再次奖励破解。我猜测是两个原因

  • 一是模型能力差，毕竟6亿参数，如果两个标签都完整，模型很难摸索出。

  • 二是奖励不明确，如果think和answer的奖励不加以区分，即无论是有think，还是answer，都是一个分。模型会晕的，哪个先出现，模型就定死哪个了，不会再额外探索另一个。

• 数据

  • {"question": "A very special island is inhabited only by knights and knaves. Knights always tell the truth, and knaves always lie. You meet 2 inhabitants: Zoey, and Jacob. As Zoey put it, \"If Jacob is a knight then Zoey is a knight\". Jacob told you that Zoey is not a knave. So who is a knight and who is a knave?", "model_response": " \n\n1. Assume Zoey is a knave.\n2. If Zoey is a knave, then her statement \"If Jacob is a knight then Zoey is a knight\" would mean that Jacob is a knave (since a knave lies and says \"If A is a knave then B is a knave\" is false if A is a knave).\n3. However, if Zoey is a knave and Jacob is a knave, Jacob would have to lie about Zoey being a knave, which contradicts the fact that Zoey is not a knave.\n4. Therefore, Zoey cannot be a knave.\n5. Now consider Jacob.\n6. If Jacob is a knave, then his statement \"Zoey is not a knave\" would mean that Zoey is a knave (since a knave lies and says \"Zoey is not a knave\" is false if Zoey is a knave).\n7. But if Jacob is a knave and Zoey is a knave, Zoey would have to lie about Jacob being a knave, which again contradicts the fact that Zoey is not a knave.\n8. Therefore, Jacob cannot be a knave.\n9. Combining these two conclusions, we find that Zoey is a knave and Jacob is a knave.\n\n<answer>\nPerson A is a knave and Person B is a knave.\n</answer>", "names": ["zoey", "jacob"], "ground_truth": {"zoey": "knight", "jacob": "knight"}, "correct": 0}
    {"question": "A very special island is inhabited only by knights and knaves. Knights always tell the truth, and knaves always lie. You meet 2 inhabitants: David, and Jackson. David said that If Jackson is a knight then Jackson is a knave. Jackson stated, \"David is a knave or David is a knight\". So who is a knight and who is a knave?", "model_response": " \n\n1. Assume Jackson is a knave. \n    - If Jackson is a knave, then his statement \"David is a knave or David is a knight\" must be a lie. \n    - If David is a knave, then his statement \"If Jackson is a knave then Jackson is a knave\" must be a lie, which contradicts our assumption that Jackson is a knave. \n    - Therefore, Jackson cannot be a knave under this assumption. \n2. Assume Jackson is a knight. \n    - If Jackson is a knight, then his statement \"David is a knave or David is a knight\" must be a truth. \n    - If David is a knave, then his statement \"If Jackson is a knave then Jackson is a knave\" must be a lie, which contradicts our assumption that Jackson is a knight. \n    - Therefore, Jackson cannot be a knight under this assumption. \n\nSince both assumptions lead to contradictions, we conclude that Jackson is a knave and David is a knight. \n\n<answer>\nPerson A is a knave and Person B is a knight.\n</answer>", "names": ["david", "jackson"], "ground_truth": {"david": "knave", "jackson": "knight"}, "correct": 0}
    {"question": "A very special island is inhabited only by knights and knaves. Knights always tell the truth, and knaves always lie. You meet 2 inhabitants: Matthew, and Jackson. As Matthew put it, \"Jackson is a knave and Jackson is a knight\". \"Matthew is not a knight,\" Jackson claimed. So who is a knight and who is a knave?", "model_response": " \n\n1. **Assume Jackson is a knave and Matthew is a knight.**\n   - If Jackson is a knave, he would lie. However, this contradicts his claim that \"Jackson is a knave and Jackson is a knight.\"\n   - Therefore, this assumption is false.\n\n2. **Assume Jackson is a knight and Matthew is a knave.**\n   - If Jackson is a knight, he would tell the truth. Since Matthew claims \"Jackson is a knave and Jackson is a knight,\" this aligns with his claim.\n   - Therefore, this assumption is true.\n\nThus, the final conclusion is that Person B is a knight and Person A is a knave.", "names": ["matthew", "jackson"], "ground_truth": {"matthew": "knave", "jackson": "knight"}, "correct": -1}
    

  • 模型确实在推理，但是呢，强行拟合我要求的answer实例

• 更改思路

  • 取消think，只要求模型将最终结果放在answer中，并且指令取消实例。

V3-Res

虽然上面提到取消think，写到answer标签即可，但是效果仍然很差，和V2结果差不多

所以参考Gsm8K数据的思想，最后结果放到####里

此时已经训练了436步，部分输出如下：

存在不少问题，胡言乱语，重复，截断。

再加上我的验证是基于简单规则，提取内容后进行一个简单的和答案匹配

def parse_answer(answer, names):
	# 这里的answer是提取标签内容后的结果
    content = answer.lower()

    # 简单规则：找名字 + knight/knave
    roles = {}
    for name in names:

        if re.search(rf"{name}.*?knave", content):
            roles[name] = "knave"
        elif re.search(rf"{name}.*?knight", content):
            roles[name] = "knight"
    return roles
def is_correct(parsed, ground_truth_roles):
    # ground_truth_roles = {"zoey": "knave", "oliver": "knight"}
    return parsed == ground_truth_roles

所以我认为主要三个原因：

• 答案截断抑制了模型的探索步伐（算力受限，此时长度为1024）
• 规则与奖励的不完善。这和上者是相辅相成的，答案截断必然影响奖励的获得，并且后续我虽然设计了更复杂的正则提取，但模型输出仍不固定，可能会输出Both are，以及人名在后的情况
• 模型本身能力问题。不少技术报告指出，小模型最好还是蒸馏比较好，参考[2505.09388] Qwen3 Technical Report，DeepSeek-V3 Technical Report.它们都是对小模型做的蒸馏。0.6B的Base模型，能力还是不太好

Base模型的探索就到此为止了，打算采用Instruct模型进行下一步的测试。后续Base可能会用GSM8K再试试，这个规则验证比较简单，思维链也不需要很长

V4-Instruct

奖励函数

def check_format_instruct(content):
    """
    提取 <think> 和 <answer> 标签的内容。
    标准格式：<think>...</think> ... <answer>...</answer>

    Returns:
        is_valid (bool): 格式是否匹配
        think_content (str/None): think 标签内的内容
        answer_content (str/None): answer 标签内的内容
    """


    # <think>(.*?)    -> 匹配 <think> 开头，并捕获中间的所有内容 (group 1)
    # </think>        -> 匹配 think 结束标签
    # .*?             -> 非贪婪匹配中间可能存在的换行或废话
    # <answer>(.*?)   -> 匹配 <answer> 开头，并捕获中间的所有内容 (group 2)
    # </answer>       -> 匹配 answer 结束标签
    pattern = r"<think>(.*?)</think>.*?<answer>(.*?)</answer>"

    match = re.search(pattern, content.strip(), flags=re.DOTALL)

    if match:
        # group(1) 是 think 的内容，group(2) 是 answer 的内容
        think_content = match.group(1).strip()
        answer_content = match.group(2).strip()

        # 返回: True, think内容, answer内容
        return True, think_content, answer_content
    else:
        # 格式不符合标准
        return False, None, None


def compute_score(
    data_source: str,
    solution_str: str,  # 模型的完整回复
    ground_truth: str,
    extra_info: dict = None,
) -> float:
    """
    score 的 Docstring
    输入模型回复，人物名，以及对应骑士的True/False，返回分数

    data_source: 数据名
    solution_str: 模型生成的答案，纯字符串，仅回复
    ground_truth: 源数据中的答案，纯字符串，仅回复
    extra_info: 额外的信息
    """
    # [False, True, True]

    names = extra_info["names"]
    format, think, answer = check_format_instruct(solution_str)

    # 格式正确
    if format:
        parsed = parse_answer(answer, names)
        print("模型生成的答案:")
        print(parsed)
        print("真实答案:")
        print(ground_truth)
        same = is_correct(parsed, ground_truth)
        print("模型回答正确：")
        if same:  # 答案一摸一样
            return 5
        else:
            return -0.1
    else:
        return -1


def my_custom_reward_func(completions, solution, names, **kwargs):
    """
    TRL 会自动把 dataset['ground_truth'] 传给 ground_truth 参数，
    把 dataset['names'] 传给 names 参数，
    把模型生成的内容传给 completions 参数。

    completions: [{'role': 'assistant', 'content': ""}]
    """
    rewards = []
    count = len(completions)
    # 遍历 Batch 中的每一条数据
    for solution_str, gt, name_list in zip(completions, solution, names):

        name_list = [name.lower() for name in name_list]
        extra_info = {"names": name_list}
        groud_truth = name2dict(name_list, gt)
        count -= 1
        if count == 0:
            print(f"模型回复：{solution_str}")
            print(f"真实答案：{groud_truth}")
            print("==========")
        score = compute_score(
            data_source="unknown",
            solution_str=solution_str[0]["content"],
            ground_truth=groud_truth,  # 注意这里传进去的是 dict
            extra_info=extra_info
        )

        rewards.append(score)

    return rewards

数据指令

instruction_following = (
        "You are a logical reasoning expert. "
        "Always think step by step inside <think> tags, showing your complete reasoning chain. "
        "After finishing your thinking, enclose your final conclusion in <answer> tags. "
        "Do not put any extra text outside these tags. "
    )

    SYSTEM_PROMPT = """
    Answer the question using the following format:
    <think>
    Your thinking process
    </think>
    <answer>
    Your answer
    </answer>
    """
    
 question = instruction_following + " " + \
            SYSTEM_PROMPT + "  " + question_raw + " " + cot_head

由于Instruct的版本，指令遵循能力已经很强了，所以直接让它按照规则回答就可以了。

V4-Res

可以看到，虽然奖励值并不全是是负数了，但是曲线并不好看

原因在于，推理长度很多，再次被截断了，为了避免喜报

我将参数进行了调整：

num_generations=8,
max_completion_length=1500,

回答长度涨到了1500个token，但仍发生会截断，有些问题模型假设的情景很多，这里就不放图了。

所以经过分析，奖励均值没有递增的原因如下：

• 截断，导致被惩罚
• 规则泛化性不强。我注意到模型会常会回复，A and B are both knights这种，也有时人名会颠倒。这种正则可以处理，但只能是我观察到一个，打一次补丁，很难具有很强的泛化性。用模型去判断会更好些，但只有一张卡，实在没办法了

但整体奖励还是能看到，是逐渐往上增的，说明确实学到了。

其实也有优化思路，Qwen3的技术报告，有提到思考预算，我可以设置思考预算，用以解决上下文不够的问题

由于只有一张卡，每次跑要近两个小时，加上还有别的事情，所以这次的尝试就到这里了。

---

而且我还注意到一个有意思的bug，当你调用 dataset.map函数时，如果你返回了 dict的内容，会出现所有kv挤在一起

例如正常如下，因为一条数据只有两个人名

"ground_truth": {"zoey": "knight", "james": "knave"} // data 1
"ground_truth": {"Phoebe": "knight", "mike": "knave"} // data 2

可实际上你遍历返回的dataset后，取出的batch里，没条变成了这样：

"ground_truth": {"zoey": "knight", "james": "knave","Phoebe": None, "mike":None} // data 1
"ground_truth": {"Phoebe": "knight", "mike": "knave","zoey": None, "james": None} // data 2

所以不要返回字典。

后续可能会更新GRPO和PPO算法的torch实现，以及一些技术报告什么的。