Merge 450885c0d5 into 4bd6062ead

eval_llm.py

@@ -16,6 +16,8 @@ def init_model(args):
             hidden_size=args.hidden_size,
             num_hidden_layers=args.num_hidden_layers,
             use_moe=bool(args.use_moe),
+            num_experts=args.num_experts,
+            num_experts_per_tok=args.num_experts_per_tok,
             inference_rope_scaling=args.inference_rope_scaling
         ))
         moe_suffix = '_moe' if args.use_moe else ''
@@ -40,6 +42,8 @@ def main():
     parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--num_experts', default=4, type=int, help="MoE专家数量（仅use_moe=1时生效，需与训练时一致）")
+    parser.add_argument('--num_experts_per_tok', default=1, type=int, help="每个token路由到的专家数（仅use_moe=1时生效，需与训练时一致）")
     parser.add_argument('--inference_rope_scaling', default=False, action='store_true', help="启用RoPE位置编码外推（4倍，仅解决位置编码问题）")
     parser.add_argument('--max_new_tokens', default=8192, type=int, help="最大生成长度（注意：并非模型实际长文本能力）")
     parser.add_argument('--temperature', default=0.85, type=float, help="生成温度，控制随机性（0-1，越大越随机）")

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 = torch.cat([attention_mask.repeat_interleave(num_generations, dim=0), attention_mask.new_ones(output_ids.size(0), completion_ids.size(1))], dim=1)
+            full_mask = (output_ids != self.tokenizer.pad_token_id).long()
             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

@@ -106,33 +106,34 @@ 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 = tokenizer(context, add_special_tokens=False)["input_ids"]
-        input_ids = torch.tensor([prompt_ids + response_ids], device=device)
+            prompt_ids = context_ids
         rollout_result = rollout_engine.rollout(
-            prompt_ids=input_ids,
-            attention_mask=torch.ones_like(input_ids),
+            prompt_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
             num_generations=1,
             max_new_tokens=max_new_tokens,
             temperature=0.8,
         )
-        valid_len = int(rollout_result.completion_mask[0].sum().item())
-        new_ids = rollout_result.completion_ids[0, :valid_len].tolist()
-        new_logps = rollout_result.per_token_logps[0, :valid_len].tolist()
-        if len(new_ids) != len(new_logps):
-            raise RuntimeError(f"rollout token/logprob length mismatch: {len(new_ids)} vs {len(new_logps)}")
+        new_ids = rollout_result.completion_ids[0].tolist()
+        new_logps = rollout_result.per_token_logps[0].tolist()
+        if len(new_ids) != len(new_logps): Logger(f"rollout token/logprob length mismatch: {len(new_ids)} vs {len(new_logps)}")
+        pairs = [(t, lp) for t, lp in zip(new_ids, new_logps) if t != tokenizer.pad_token_id and t != tokenizer.eos_token_id]
+        new_ids = [t for t, _ in pairs]
+        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([int(t != tokenizer.eos_token_id) for t in new_ids])
+        response_mask.extend([1] * len(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
-        assistant_message = {"role": "assistant", "content": new_text}
-        messages.append(assistant_message)
+        messages.append({"role": "assistant", "content": new_text})
         for call in calls:
             name, raw = call.get("name", ""), call.get("arguments", {})
             if isinstance(raw, str):
@@ -142,17 +143,10 @@ 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))
@@ -274,7 +268,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 = (torch.arange(max_len, device=args.device).unsqueeze(0) < seq_lens.unsqueeze(1)).long()
+        full_mask = (input_ids != tokenizer.pad_token_id).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_full_sft.py

@@ -99,6 +99,8 @@ if __name__ == "__main__":
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--max_seq_len', default=768, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--num_experts', default=4, type=int, help="MoE专家数量（仅use_moe=1时生效）")
+    parser.add_argument('--num_experts_per_tok', default=1, type=int, help="每个token路由到的专家数（仅use_moe=1时生效）")
     parser.add_argument('--seed', default=42, type=int, help="随机种子（DDP下每个rank为seed+rank，每轮为seed+epoch）")
     parser.add_argument("--data_path", type=str, default="../dataset/sft_t2t_mini.jsonl", help="训练数据路径")
     parser.add_argument('--from_weight', default='pretrain', type=str, help="基于哪个权重训练，为none则不基于任何权重训练")
@@ -115,7 +117,8 @@ if __name__ == "__main__":
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
-    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe))
+    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe),
+                               num_experts=args.num_experts, num_experts_per_tok=args.num_experts_per_tok)
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     
     # ========== 3. 设置混合精度 ==========

trainer/train_grpo.py

@@ -91,7 +91,6 @@ 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]
 

trainer/train_ppo.py

@@ -120,7 +120,6 @@ 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)

trainer/train_pretrain.py

@@ -98,6 +98,8 @@ if __name__ == "__main__":
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--max_seq_len', default=340, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--num_experts', default=4, type=int, help="MoE专家数量（仅use_moe=1时生效）")
+    parser.add_argument('--num_experts_per_tok', default=1, type=int, help="每个token路由到的专家数（仅use_moe=1时生效）")
     parser.add_argument('--seed', default=42, type=int, help="随机种子（DDP下每个rank为seed+rank，每轮为seed+epoch）")
     parser.add_argument("--data_path", type=str, default="../dataset/pretrain_t2t_mini.jsonl", help="预训练数据路径")
     parser.add_argument('--from_weight', default='none', type=str, help="基于哪个权重训练，为none则从头开始")
@@ -114,7 +116,8 @@ if __name__ == "__main__":
     
     # ========== 2. 配置目录、模型参数、检查ckp ==========
     os.makedirs(args.save_dir, exist_ok=True)
-    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe))
+    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe),
+                               num_experts=args.num_experts, num_experts_per_tok=args.num_experts_per_tok)
     ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
     
     # ========== 3. 设置混合精度 ==========