github
/
minimind
mirror of https://github.com/jingyaogong/minimind.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Compare commits

base: github:5f26767e41704d3503824bc829a09f5f2ebe7955
Branches Tags
github:master
github:website
github:v2
github:minimind-v1
...
compare: github:87d708e4d288d7c2341c6d78c4f3796f01390b18
Branches Tags
github:master
github:website
github:v2
github:minimind-v1

8 Commits
5f26767e41 ... 87d708e4d2

Author SHA1 Message Date
Yuchen Lv 87d708e4d2
Merge 0c91dea4a2 into cc312c1cc6 2026-09-19 09:03:44 +08: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
TOPSAIL\78190 0c91dea4a2 [refactor] remove redundant rollout update_policy after torch.compile 2026-08-24 12:43:50 +08:00
5 changed files with 10 additions and 5 deletions
Show Stats Expand all files Collapse all files

3
eval_llm.py
View File

@ -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') 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__": if __name__ == "__main__":
try:
main() main()
except (EOFError, KeyboardInterrupt):
print()

7
model/model_minimind.py
View File

@ -158,7 +158,12 @@ class MOEFeedForward(nn.Module):
x_flat = x.view(-1, hidden_dim) x_flat = x.view(-1, hidden_dim)
scores = F.softmax(self.gate(x_flat), dim=-1) 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) 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) y = torch.zeros_like(x_flat)
for i, expert in enumerate(self.experts): for i, expert in enumerate(self.experts):
mask = (topk_idx == i) mask = (topk_idx == i)

1
trainer/train_agent.py
View File

@ -471,7 +471,6 @@ if __name__ == "__main__":
if args.use_compile == 1: if args.use_compile == 1:
model = torch.compile(model) model = torch.compile(model)
Logger('torch.compile enabled') Logger('torch.compile enabled')
rollout_engine.update_policy(model)
if dist.is_initialized(): if dist.is_initialized():
# 同 train_ppo：RoPE buffer 各 rank 一致，每步广播纯属浪费 # 同 train_ppo：RoPE buffer 各 rank 一致，每步广播纯属浪费
model = DistributedDataParallel(model, device_ids=[local_rank], broadcast_buffers=False) model = DistributedDataParallel(model, device_ids=[local_rank], broadcast_buffers=False)

1
trainer/train_grpo.py
View File

@ -312,7 +312,6 @@ if __name__ == "__main__":
if args.use_compile == 1: if args.use_compile == 1:
model = torch.compile(model) model = torch.compile(model)
Logger('torch.compile enabled') Logger('torch.compile enabled')
rollout_engine.update_policy(model)
if dist.is_initialized(): if dist.is_initialized():
# 同 train_ppo：RoPE buffer 各 rank 一致，每步广播纯属浪费 # 同 train_ppo：RoPE buffer 各 rank 一致，每步广播纯属浪费
model = DistributedDataParallel(model, device_ids=[local_rank], broadcast_buffers=False) model = DistributedDataParallel(model, device_ids=[local_rank], broadcast_buffers=False)

1
trainer/train_ppo.py
View File

@ -427,7 +427,6 @@ if __name__ == "__main__":
if args.use_compile == 1: if args.use_compile == 1:
actor_model = torch.compile(actor_model) actor_model = torch.compile(actor_model)
Logger('torch.compile enabled') Logger('torch.compile enabled')
rollout_engine.update_policy(actor_model)
if dist.is_initialized(): if dist.is_initialized():
# freqs_cos/freqs_sin 各 rank 由 config 确定性算出，默认每步广播一次纯属浪费 # freqs_cos/freqs_sin 各 rank 由 config 确定性算出，默认每步广播一次纯属浪费
actor_model = DistributedDataParallel(actor_model, device_ids=[local_rank], broadcast_buffers=False) actor_model = DistributedDataParallel(actor_model, device_ids=[local_rank], broadcast_buffers=False)