Merge 56d4cddfd8 into cc312c1cc6

[fix] unify Ctrl+C handling
Merge pull request #857 from wbbeyourself/fix-ctrl-c-exception
2026-09-18 10:20:18 -07:00 · 2026-09-18 16:33:09 +08:00 · 2026-09-18 15:54:43 +08:00 · 2026-09-18 15:28:50 +08:00 · 2026-09-18 15:03:36 +08:00 · 2026-09-17 21:31:41 +05:30
3 changed files with 20 additions and 8 deletions
--- a/dataset/lm_dataset.py
+++ b/dataset/lm_dataset.py
@ -28,10 +28,13 @@ def pre_processing_chat(conversations, add_system_ratio=0.2):
            return [{'role': 'system', 'content': random.choice(SYSTEM_PROMPTS)}] + conversations
    return conversations

-def post_processing_chat(prompt_content, empty_think_ratio=0.2):
+def post_processing_chat(prompt_content, empty_think_ratio=0.2, remove_empty_think=None):
    # 以80%概率移除空思考标签
-    if '<think>\n\n</think>\n\n' in prompt_content and random.random() > empty_think_ratio:
-        prompt_content = prompt_content.replace('<think>\n\n</think>\n\n', '')
+    if '<think>\n\n</think>\n\n' in prompt_content:
+        if remove_empty_think is None:
+            remove_empty_think = random.random() > empty_think_ratio
+        if remove_empty_think:
+            prompt_content = prompt_content.replace('<think>\n\n</think>\n\n', '')
    return prompt_content

 class PretrainDataset(Dataset):
@ -139,12 +142,13 @@ class DPODataset(Dataset):
        chosen_prompt = self.tokenizer.apply_chat_template(
            chosen, tokenize=False, add_generation_prompt=False
        )
-        chosen_prompt = post_processing_chat(chosen_prompt)

        rejected_prompt = self.tokenizer.apply_chat_template(
            rejected, tokenize=False, add_generation_prompt=False
        )
-        rejected_prompt = post_processing_chat(rejected_prompt)
+        remove_empty_think = random.random() > 0.2
+        chosen_prompt = post_processing_chat(chosen_prompt, remove_empty_think=remove_empty_think)
+        rejected_prompt = post_processing_chat(rejected_prompt, remove_empty_think=remove_empty_think)
        chosen_encoding = self.tokenizer(
            chosen_prompt, truncation=True, max_length=self.max_length, padding='max_length'
        )
@ -253,4 +257,4 @@ class AgentRLDataset(Dataset):


 if __name__ == "__main__":
-    pass
+    pass
--- a/eval_llm.py
+++ b/eval_llm.py
@ -91,4 +91,7 @@ def main():
        print(f'\n[Speed]: {gen_tokens / (time.time() - st):.2f} tokens/s\n\n') if args.show_speed else print('\n\n')

 if __name__ == "__main__":
-    main()
+    try:
+        main()
+    except (EOFError, KeyboardInterrupt):
+        print()
--- a/model/model_minimind.py
+++ b/model/model_minimind.py
@ -158,7 +158,12 @@ class MOEFeedForward(nn.Module):
        x_flat = x.view(-1, hidden_dim)
        scores = F.softmax(self.gate(x_flat), dim=-1)
        topk_weight, topk_idx = torch.topk(scores, k=self.config.num_experts_per_tok, dim=-1, sorted=False)
-        if self.config.norm_topk_prob: topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
+        if self.config.norm_topk_prob:
+            if self.config.num_experts_per_tok > 1: 
+                topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
+            else:
+                top1 = torch.topk(F.softmax(self.gate(x_flat.detach()), dim=-1), k=1, dim=-1, sorted=False)[0]
+                topk_weight = top1 - top1.detach() + 1.0 # k=1: 1.0 forward on purpose, gradient via straight-through
        y = torch.zeros_like(x_flat)
        for i, expert in enumerate(self.experts):
            mask = (topk_idx == i)
Author	SHA1	Message	Date
Daoyuan Li	ccfaee37da	Merge `56d4cddfd8` into `cc312c1cc6`	2026-09-18 10:20:18 -07:00
jingyaogong	cc312c1cc6	[fix] unify Ctrl+C handling	2026-09-18 16:33:09 +08:00
jingyaogong	e2ee878519	Merge pull request #857 from wbbeyourself/fix-ctrl-c-exception fix: eval_llm.py Ctrl+C 退出异常修复	2026-09-18 15:54:43 +08:00
jingyaogong	b63a9565e3	[fix] restore top-1 router gradient #858	2026-09-18 15:28:50 +08:00
jingyaogong	1aae25911b	Merge pull request #859 from basil-k-aji-dev/fix/moe-router-gradient-topk1 [fix] keep the MoE router trainable at num_experts_per_tok=1	2026-09-18 15:03:36 +08:00
basil-k-aji-dev	1c876435ef	[fix] keep the MoE router trainable at num_experts_per_tok=1 Normalising a single top-k weight by its own sum gives exactly 1.0 for every token, so the gate leaves the graph entirely: topk_idx is a discrete index and the weight is now a constant, which leaves the LM loss no differentiable path to the router. It then trains on aux_loss alone, whose optimum is a uniform distribution, and the experts never specialise. Measured on MOEFeedForward with router_aux_loss_coef=0, so only the LM path can contribute a gradient: num_experts_per_tok=1, norm_topk_prob=True 7.02e-10 <- current default num_experts_per_tok=1, norm_topk_prob=False 1.02e-02 num_experts_per_tok=2, norm_topk_prob=True 1.67e-02 Only normalise when there is more than one weight to normalise. Note this changes the forward pass for existing num_experts_per_tok=1 MoE checkpoints: the expert output was scaled by a constant 1.0 and is now scaled by the gate probability. Closes #858.	2026-09-17 21:31:41 +05:30
wbbeyourself	9fa620e2e5	fix: eval_llm.py Ctrl+C 退出异常修复	2026-09-16 17:47:55 +08:00
Daoyuan Li	56d4cddfd8	[fix] keep DPO preference pair think formatting aligned	2026-08-29 18:29:00 -07:00