Merge 5b26153b75 into 7a9137d2e9

[perf] skip the per-forward RoPE buffer broadcast in the RL trainers

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/train_agent.py

@@ -25,16 +25,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, safe_math_eval
 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"]}}},
@@ -214,7 +210,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 +230,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
 
@@ -473,7 +469,8 @@ if __name__ == "__main__":
         Logger('torch.compile enabled')
         rollout_engine.update_policy(model)
     if dist.is_initialized():
-        model = DistributedDataParallel(model, device_ids=[local_rank])
+        # 同 train_ppo：RoPE buffer 各 rank 一致，每步广播纯属浪费
+        model = DistributedDataParallel(model, device_ids=[local_rank], broadcast_buffers=False)
     rollout_engine.update_policy(model)
 
     for epoch in range(start_epoch, args.epochs):

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)
@@ -314,7 +309,8 @@ if __name__ == "__main__":
         Logger('torch.compile enabled')
         rollout_engine.update_policy(model)
     if dist.is_initialized():
-        model = DistributedDataParallel(model, device_ids=[local_rank])
+        # 同 train_ppo：RoPE buffer 各 rank 一致，每步广播纯属浪费
+        model = DistributedDataParallel(model, device_ids=[local_rank], broadcast_buffers=False)
     rollout_engine.update_policy(model)
     
     # ========== 8. 开始训练 ==========

trainer/train_ppo.py

@@ -22,21 +22,16 @@ 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):
         super().__init__(params)
-        # 替换lm_head为输出单一价值的线性层
+        # lm_head 不参与 forward，仅靠 tie_word_embeddings 与 embed_tokens 共享权重，解绑后 DDP 会报未使用参数
         self.value_head = nn.Linear(params.hidden_size, 1)
 
     def forward(self, input_ids=None, attention_mask=None, **kwargs):
@@ -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)
@@ -429,8 +424,9 @@ if __name__ == "__main__":
         Logger('torch.compile enabled')
         rollout_engine.update_policy(actor_model)
     if dist.is_initialized():
-        actor_model = DistributedDataParallel(actor_model, device_ids=[local_rank])
-        critic_model = DistributedDataParallel(critic_model, device_ids=[local_rank])
+        # freqs_cos/freqs_sin 各 rank 由 config 确定性算出，默认每步广播一次纯属浪费
+        actor_model = DistributedDataParallel(actor_model, device_ids=[local_rank], broadcast_buffers=False)
+        critic_model = DistributedDataParallel(critic_model, device_ids=[local_rank], broadcast_buffers=False)
     rollout_engine.update_policy(actor_model)
     
     # ========== 8. 开始训练 ==========