Merge 5b26153b75 into a3c7b01cc0

tests/test_reward_utils.py

@@ -0,0 +1,43 @@
+import unittest
+
+from trainer.reward_utils import rep_penalty
+
+
+class StubTokenizer:
+    def __init__(self, token_ids):
+        self.token_ids = token_ids
+        self.calls = []
+
+    def encode(self, text, add_special_tokens=True):
+        self.calls.append((text, add_special_tokens))
+        return self.token_ids
+
+
+class RepetitionPenaltyTests(unittest.TestCase):
+    def test_repeated_chinese_tokens_receive_penalty(self):
+        tokenizer = StubTokenizer([1, 2, 1, 2, 1, 2, 1, 2, 1, 2])
+
+        self.assertEqual(rep_penalty("你好你好你好你好你好", tokenizer), 0.5)
+        self.assertEqual(tokenizer.calls, [("你好你好你好你好你好", False)])
+
+    def test_repeated_english_tokens_receive_penalty_case_insensitively(self):
+        tokenizer = StubTokenizer([3, 3, 3, 3])
+
+        self.assertEqual(rep_penalty("Hello HELLO hello hello", tokenizer), 0.5)
+        self.assertEqual(tokenizer.calls, [("hello hello hello hello", False)])
+
+    def test_unique_and_short_sequences_have_no_penalty(self):
+        self.assertEqual(rep_penalty("unique", StubTokenizer([1, 2, 3, 4, 5])), 0.0)
+        self.assertEqual(rep_penalty("short", StubTokenizer([1, 2])), 0.0)
+
+    def test_invalid_configuration_is_rejected(self):
+        tokenizer = StubTokenizer([])
+
+        with self.assertRaisesRegex(ValueError, "n must be greater than zero"):
+            rep_penalty("text", tokenizer, n=0)
+        with self.assertRaisesRegex(ValueError, "cap must be non-negative"):
+            rep_penalty("text", tokenizer, cap=-0.1)
+
+
+if __name__ == "__main__":
+    unittest.main()

trainer/reward_utils.py

@@ -0,0 +1,19 @@
+def rep_penalty(text, tokenizer, n=3, cap=0.5):
+    r"""Return a capped repetition penalty over model-token n-grams.
+
+    Token IDs keep the calculation consistent across languages. In particular,
+    contiguous CJK text must not be treated as one word, which is what happens
+    with a ``\w+``-based regular expression.
+    """
+    if n <= 0:
+        raise ValueError("n must be greater than zero")
+    if cap < 0:
+        raise ValueError("cap must be non-negative")
+
+    token_ids = tokenizer.encode(text.lower(), add_special_tokens=False)
+    grams = [tuple(token_ids[i:i + n]) for i in range(len(token_ids) - n + 1)]
+    if not grams:
+        return 0.0
+
+    repeated = len(grams) - len(set(grams))
+    return min(cap, repeated * cap * 2 / len(grams))

trainer/rollout_engine.py

@@ -84,7 +84,7 @@ class TorchRolloutEngine(RolloutEngine):
             ).clone()  # [B*num_gen, P+R]
             prompt_len = prompt_ids.size(1)
             completion_ids = output_ids[:, prompt_len:]  # [B*num_gen, R]
-            full_mask = (output_ids != self.tokenizer.pad_token_id).long()
+            full_mask = torch.cat([attention_mask.repeat_interleave(num_generations, dim=0), attention_mask.new_ones(output_ids.size(0), completion_ids.size(1))], dim=1)
             per_token_logps = compute_per_token_logps(self.policy_model, output_ids, completion_ids.size(1), attention_mask=full_mask)
         completions = self.tokenizer.batch_decode(completion_ids, skip_special_tokens=True)
         return RolloutResult(output_ids, completion_ids, per_token_logps, completions,

trainer/train_agent.py

@@ -26,16 +26,12 @@ from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
 from dataset.lm_dataset import AgentRLDataset
 from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
 from trainer.rollout_engine import create_rollout_engine, compute_per_token_logps
+from trainer.reward_utils import rep_penalty
 
 warnings.filterwarnings('ignore')
 
 # ================================ 工具与 Reward = Start ================================
 
-def rep_penalty(text, n=3, cap=0.5):
-    toks = re.findall(r"\w+|[^\w\s]", text.lower())
-    grams = [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)]
-    return min(cap, (len(grams) - len(set(grams))) * cap * 2 / len(grams)) if grams else 0.0
-
 # ======== 工具定义 ========
 TOOLS = [
     {"type": "function", "function": {"name": "calculate_math", "description": "计算数学表达式", "parameters": {"type": "object", "properties": {"expression": {"type": "string"}}, "required": ["expression"]}}},
@@ -106,34 +102,33 @@ def rollout_single(rollout_engine, tokenizer, messages, tools, max_turns=3, max_
     open_thinking = random.random() < thinking_ratio
     for turn in range(max_turns):
         context = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, tools=tools, open_thinking=open_thinking)
-        inputs = tokenizer(context, return_tensors="pt", add_special_tokens=False).to(device)
-        context_ids = inputs["input_ids"][0].tolist()
         if prompt_ids is None:
-            prompt_ids = context_ids
+            prompt_ids = tokenizer(context, add_special_tokens=False)["input_ids"]
+        input_ids = torch.tensor([prompt_ids + response_ids], device=device)
         rollout_result = rollout_engine.rollout(
-            prompt_ids=inputs["input_ids"],
+            prompt_ids=input_ids,
-            attention_mask=inputs["attention_mask"],
+            attention_mask=torch.ones_like(input_ids),
             num_generations=1,
             max_new_tokens=max_new_tokens,
             temperature=0.8,
         )
-        new_ids = rollout_result.completion_ids[0].tolist()
+        valid_len = int(rollout_result.completion_mask[0].sum().item())
-        new_logps = rollout_result.per_token_logps[0].tolist()
+        new_ids = rollout_result.completion_ids[0, :valid_len].tolist()
-        if len(new_ids) != len(new_logps): Logger(f"rollout token/logprob length mismatch: {len(new_ids)} vs {len(new_logps)}")
+        new_logps = rollout_result.per_token_logps[0, :valid_len].tolist()
-        pairs = [(t, lp) for t, lp in zip(new_ids, new_logps) if t != tokenizer.pad_token_id and t != tokenizer.eos_token_id]
+        if len(new_ids) != len(new_logps):
-        new_ids = [t for t, _ in pairs]
+            raise RuntimeError(f"rollout token/logprob length mismatch: {len(new_ids)} vs {len(new_logps)}")
-        new_logps = [lp for _, lp in pairs]
         new_text = rollout_result.completions[0]
         all_outputs.append(new_text)
         response_ids.extend(new_ids)
-        response_mask.extend([1] * len(new_ids))
+        response_mask.extend([int(t != tokenizer.eos_token_id) for t in new_ids])
         response_old_logps.extend(new_logps)
         final_context = context + new_text
         calls = parse_tool_calls(new_text)
         if not calls:
             break
         unfinished = turn == max_turns - 1
-        messages.append({"role": "assistant", "content": new_text})
+        assistant_message = {"role": "assistant", "content": new_text}
+        messages.append(assistant_message)
         for call in calls:
             name, raw = call.get("name", ""), call.get("arguments", {})
             if isinstance(raw, str):
@@ -143,10 +138,17 @@ def rollout_single(rollout_engine, tokenizer, messages, tools, max_turns=3, max_
             result_str = (json.dumps(result, ensure_ascii=False) if result else '{"error": "tool not found"}')[:2048]  # 防止天文数字撑爆tokenizer
             messages.append({"role": "tool", "content": result_str})
 
+        marker = f"<|agent_observation_{id(messages)}_{len(response_ids)}|>"
+        assistant_message["content"] += marker
+        marked_context = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=not unfinished, tools=tools, open_thinking=open_thinking)
+        assistant_message["content"] = new_text
+        _, found, observation = marked_context.partition(marker)
+        if not found:
+            raise RuntimeError("chat template did not preserve the assistant content boundary")
+        obs_delta = tokenizer(observation, add_special_tokens=False)["input_ids"]
+        if new_ids and new_ids[-1] == tokenizer.eos_token_id and obs_delta[:1] == [tokenizer.eos_token_id]:
+            obs_delta = obs_delta[1:]
         observe_context = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=not unfinished, tools=tools, open_thinking=open_thinking)
-        observe_ids = tokenizer(observe_context, return_tensors="pt", add_special_tokens=False)["input_ids"][0].tolist()
-        current_len = len(prompt_ids) + len(response_ids)
-        obs_delta = observe_ids[current_len:]
         response_ids.extend(obs_delta)
         response_mask.extend([0] * len(obs_delta))
         response_old_logps.extend([0.0] * len(obs_delta))
@@ -214,7 +216,7 @@ def calculate_rewards(prompts, completions, gt_batch, tools_batch, num_gen, rewa
                 messages = [{"role": role, "content": content.strip()} for role, content in matches]
                 score = reward_model.get_score(messages, answer)
                 reward += score  # RM分
-            reward -= rep_penalty(answer)
+            reward -= rep_penalty(answer, tokenizer)
             rewards[idx] = max(min(reward, 3.0), -3.0)  # 总分Clip
         # -------- 有工具调用：执行结果奖励 --------
         else:
@@ -234,7 +236,7 @@ def calculate_rewards(prompts, completions, gt_batch, tools_batch, num_gen, rewa
             verified = validate_gt_in_text(final_text, gt) if gt else set()
             if gt: reward += 2.5 * len(verified) / len(gt)  # GT分
             if unfinished: reward -= 0.5  # 未完成扣分
-            reward -= rep_penalty(final_text if final_text else answer)
+            reward -= rep_penalty(final_text if final_text else answer, tokenizer)
             rewards[idx] = max(min(reward, 3.0), -3.0)  # 总分Clip
     return rewards
 
@@ -268,7 +270,7 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model
         prompt_lens = torch.tensor([prompt_len for _, _, prompt_len, _ in packed_samples], device=args.device)
         full_response_masks = torch.tensor([mask + [0] * (max_len - len(mask)) for _, mask, _, _ in packed_samples], device=args.device, dtype=torch.float32)
         old_per_token_logps = torch.tensor([old_logps + [0.0] * ((max_len - 1) - len(old_logps)) for _, _, _, old_logps in packed_samples], device=args.device, dtype=torch.float32)
-        full_mask = (input_ids != tokenizer.pad_token_id).long()
+        full_mask = (torch.arange(max_len, device=args.device).unsqueeze(0) < seq_lens.unsqueeze(1)).long()
 
         rewards = calculate_rewards(prompts, completions, gt_batch, tools_batch, args.num_generations, reward_model, device=args.device, turn_outputs_batch=turn_outputs_batch, unfinished_batch=unfinished_batch)
 

trainer/train_grpo.py

@@ -24,16 +24,11 @@ from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
 from dataset.lm_dataset import RLAIFDataset
 from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
 from trainer.rollout_engine import create_rollout_engine
+from trainer.reward_utils import rep_penalty
 
 warnings.filterwarnings('ignore')
 
 
-def rep_penalty(text, n=3, cap=0.5):
-    toks = re.findall(r"\w+|[^\w\s]", text.lower())
-    grams = [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)]
-    return min(cap, (len(grams) - len(set(grams))) * cap * 2 / len(grams)) if grams else 0.0
-
-
 def calculate_rewards(prompts, responses, reward_model):
     rewards = torch.zeros(len(responses), device=args.device)
 
@@ -57,7 +52,7 @@ def calculate_rewards(prompts, responses, reward_model):
                     rewards[response_idx] += 1.0 if 20 <= len(thinking_content.strip()) <= 300 else -0.5
                     rewards[response_idx] += 0.25 if response.count('</think>') == 1 else -0.25
                     answer = answer_content.strip()
-                rewards[response_idx] -= rep_penalty(answer)
+                rewards[response_idx] -= rep_penalty(answer, tokenizer)
 
                 score = reward_model.get_score(messages, answer)
                 reward_model_scores.append(score)
@@ -91,6 +86,7 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
         prompt_lens = rollout_result.prompt_lens.to(args.device)
         full_mask = (outputs != tokenizer.pad_token_id).long()
         logp_pos = prompt_lens.unsqueeze(1) - 1 + torch.arange(completion_ids.size(1), device=args.device).unsqueeze(0)
+        full_mask.scatter_(1, logp_pos + 1, rollout_result.completion_mask.to(args.device, dtype=full_mask.dtype))
 
         rewards = calculate_rewards(prompts, completions, reward_model).to(args.device)  # [B*num_gen]
 
@@ -331,4 +327,4 @@ if __name__ == "__main__":
     # ========== 9. 清理分布进程 ==========
     if dist.is_initialized():
         dist.barrier()
-        dist.destroy_process_group()
+        dist.destroy_process_group()

trainer/train_ppo.py

@@ -22,16 +22,11 @@ from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
 from dataset.lm_dataset import RLAIFDataset
 from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
 from trainer.rollout_engine import create_rollout_engine
+from trainer.reward_utils import rep_penalty
 
 warnings.filterwarnings('ignore')
 
 
-def rep_penalty(text, n=3, cap=0.5):
-    toks = re.findall(r"\w+|[^\w\s]", text.lower())
-    grams = [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)]
-    return min(cap, (len(grams) - len(set(grams))) * cap * 2 / len(grams)) if grams else 0.0
-
-
 # 自定义的Critic模型，继承自MiniMindLM
 class CriticModel(MiniMindForCausalLM):
     def __init__(self, params):
@@ -64,7 +59,7 @@ def calculate_rewards(prompts, responses, reward_model):
                 rewards[i] += 1.0 if 20 <= len(thinking_content.strip()) <= 300 else -0.5
                 rewards[i] += 0.25 if response.count('</think>') == 1 else -0.25
                 answer = answer_content.strip()
-            rewards[i] -= rep_penalty(answer)
+            rewards[i] -= rep_penalty(answer, tokenizer)
 
             score = reward_model.get_score(messages, answer)
             reward_model_scores.append(score)
@@ -120,6 +115,7 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched
         resp_idx = torch.arange(resp_labels.size(1), device=gen_out.device).unsqueeze(0)
         logp_pos = prompt_lens.unsqueeze(1) - 1 + resp_idx
         resp_pad_mask = rollout_result.completion_mask.to(args.device).bool()
+        full_mask.scatter_(1, logp_pos + 1, resp_pad_mask.to(full_mask.dtype))
         resp_lengths = resp_pad_mask.sum(dim=1); valid_resp = resp_lengths > 0; eos_mask = resp_labels.eq(tokenizer.eos_token_id) & resp_pad_mask
         has_eos = eos_mask.any(dim=1); eos_pos = torch.argmax(eos_mask.int(), dim=1)
         resp_lengths = torch.where(has_eos, eos_pos + 1, resp_lengths).long().clamp(min=1)
@@ -448,4 +444,4 @@ if __name__ == "__main__":
     # ========== 9. 清理分布进程 ==========
     if dist.is_initialized():
         dist.barrier()
-        dist.destroy_process_group()
+        dist.destroy_process_group()