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:08e1bdb9d810189c006ba21a26454433c63fa012
Branches Tags
github:master
github:website
github:v2
github:minimind-v1
...
compare: github:a1080ec37014a3fc77aa10493cc8b5f286a20493
Branches Tags
github:master
github:website
github:v2
github:minimind-v1

9 Commits
08e1bdb9d8 ... a1080ec370

Author SHA1 Message Date
Neil Ge a1080ec370
Merge 9cb07dd1ec 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
neil.ge 9cb07dd1ec perf: enable fused AdamW in pretrain training 
Same fused AdamW optimization as the full SFT script, using the
existing device_type variable so the CPU-only path is unaffected.

Co-authored-by: TRAE CLI <traecli@bytedance.com>
 2026-09-05 19:16:51 +08:00
neil.ge 859c8ee2a3 perf: enable fused AdamW in full SFT training 
Use PyTorch's fused AdamW implementation to reduce optimizer kernel
launch overhead and speed up the training step on CUDA devices.
Guard with device check so the CPU-only fallback path is unaffected.

Co-authored-by: TRAE CLI <traecli@bytedance.com>
 2026-09-05 18:37:04 +08:00
4 changed files with 12 additions and 4 deletions
Show Stats Expand all files Collapse all files

5
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__":
main() try:
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)

2
trainer/train_full_sft.py
View File

@ -137,7 +137,7 @@ if __name__ == "__main__":
train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len) train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate, fused=('cuda' in args.device))
# ========== 6. 从ckp恢复状态 ========== # ========== 6. 从ckp恢复状态 ==========
start_epoch, start_step = 0, 0 start_epoch, start_step = 0, 0

2
trainer/train_pretrain.py
View File

@ -136,7 +136,7 @@ if __name__ == "__main__":
train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len) train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate, fused=(device_type == 'cuda'))
# ========== 6. 从ckp恢复状态 ========== # ========== 6. 从ckp恢复状态 ==========
start_epoch, start_step = 0, 0 start_epoch, start_step = 0, 0