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:3681271736a9b5769dc483bfbd7c134431a523b6
Branches Tags
github:master
github:website
github:v2
github:minimind-v1
...
compare: github:92a71d2a37b67e60470aa6614e22f93765ec0f91
Branches Tags
github:master
github:website
github:v2
github:minimind-v1

8 Commits
3681271736 ... 92a71d2a37

Author SHA1 Message Date
Corey52HZ 92a71d2a37
Merge 655d4705eb 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
Lingxu 655d4705eb feat: support Apple Silicon (MPS) for training and web demo 
`device_type` currently only distinguishes "cuda" from "cpu", so running
with `--device mps` on Apple Silicon falls into the "cpu" branch, disables
autocast entirely and forces fp32 training.

Changes (9 files, mechanical):
- `device_type` now resolves to cuda / mps / cpu
- `torch.cuda.amp.autocast(...)` -> `torch.amp.autocast(device_type, ...)`
- `torch.cuda.amp.GradScaler(...)` -> `torch.amp.GradScaler(device_type, ...)`,
  enabled only for CUDA + fp16 (MPS has no gradient scaler)
- `web_demo.py` device selection now prefers mps when available

This also silences the `torch.cuda.amp.*` deprecation warnings on torch 2.x.
CUDA behaviour is unchanged.

Measured on Apple M5 Pro (24GB), hidden_size=512 / 8 layers, bs=16, seq=512:
- before (fp32, no autocast): 11,784 tok/s
- after  (bf16 autocast):     15,492 tok/s   (+31%)

Verified end-to-end on MPS: pretrain, full_sft, lora, dpo, ppo, grpo,
distillation and agent RL all run to completion, and the resulting weights
export correctly to transformers / GGUF.

Note for macOS users: `--num_workers 0` is currently required, as DataLoader
worker processes deadlock on macOS + Python 3.12 due to fork safety.
 2026-09-09 21:42:03 +08:00
11 changed files with 32 additions and 24 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
scripts/web_demo.py
View File

@ -71,7 +71,7 @@ st.markdown("""
</style> </style>
""", unsafe_allow_html=True) """, unsafe_allow_html=True)
device = "cuda" if torch.cuda.is_available() else "cpu" device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
# 多语言文本 # 多语言文本
LANG_TEXTS = { LANG_TEXTS = {

4
trainer/train_agent.py
View File

@ -422,9 +422,9 @@ if __name__ == "__main__":
max_seq_len=args.max_seq_len + args.max_gen_len, use_moe=bool(args.use_moe)) max_seq_len=args.max_seq_len + args.max_gen_len, use_moe=bool(args.use_moe))
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume == 1 else None ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume == 1 else None
device_type = "cuda" if "cuda" in args.device else "cpu" device_type = "cuda" if "cuda" in args.device else ("mps" if "mps" in args.device else "cpu")
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) autocast_ctx = nullcontext() if device_type == "cpu" else torch.amp.autocast(device_type, dtype=dtype)
wandb = None wandb = None
if args.use_wandb and is_main_process(): if args.use_wandb and is_main_process():

6
trainer/train_distillation.py
View File

@ -188,9 +188,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config_student, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None ckp_data = lm_checkpoint(lm_config_student, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ========== # ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu" device_type = "cuda" if "cuda" in args.device else ("mps" if "mps" in args.device else "cpu")
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) autocast_ctx = nullcontext() if device_type == "cpu" else torch.amp.autocast(device_type, dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None
@ -210,7 +210,7 @@ if __name__ == "__main__":
Logger(f'教师模型总参数量：{sum(p.numel() for p in teacher_model.parameters()) / 1e6:.3f} M') Logger(f'教师模型总参数量：{sum(p.numel() for p in teacher_model.parameters()) / 1e6:.3f} M')
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.amp.GradScaler(device_type, enabled=(args.dtype == 'float16' and device_type == 'cuda'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# ========== 6. 从ckp恢复状态 ========== # ========== 6. 从ckp恢复状态 ==========

6
trainer/train_dpo.py
View File

@ -166,9 +166,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ========== # ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu" device_type = "cuda" if "cuda" in args.device else ("mps" if "mps" in args.device else "cpu")
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) autocast_ctx = nullcontext() if device_type == "cpu" else torch.amp.autocast(device_type, dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None
@ -190,7 +190,7 @@ if __name__ == "__main__":
train_ds = DPODataset(args.data_path, tokenizer, max_length=args.max_seq_len) train_ds = DPODataset(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.amp.GradScaler(device_type, enabled=(args.dtype == 'float16' and device_type == 'cuda'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# ========== 6. 从ckp恢复状态 ========== # ========== 6. 从ckp恢复状态 ==========

6
trainer/train_full_sft.py
View File

@ -119,9 +119,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ========== # ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu" device_type = "cuda" if "cuda" in args.device else ("mps" if "mps" in args.device else "cpu")
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) autocast_ctx = nullcontext() if device_type == "cpu" else torch.amp.autocast(device_type, dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None
@ -136,7 +136,7 @@ if __name__ == "__main__":
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device) model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
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.amp.GradScaler(device_type, enabled=(args.dtype == 'float16' and device_type == 'cuda'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# ========== 6. 从ckp恢复状态 ========== # ========== 6. 从ckp恢复状态 ==========

4
trainer/train_grpo.py
View File

@ -257,9 +257,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ========== # ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu" device_type = "cuda" if "cuda" in args.device else ("mps" if "mps" in args.device else "cpu")
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) autocast_ctx = nullcontext() if device_type == "cpu" else torch.amp.autocast(device_type, dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None

6
trainer/train_lora.py
View File

@ -112,9 +112,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config, weight=args.lora_name, save_dir='../checkpoints') if args.from_resume==1 else None ckp_data = lm_checkpoint(lm_config, weight=args.lora_name, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ========== # ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu" device_type = "cuda" if "cuda" in args.device else ("mps" if "mps" in args.device else "cpu")
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) autocast_ctx = nullcontext() if device_type == "cpu" else torch.amp.autocast(device_type, dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None
@ -148,7 +148,7 @@ if __name__ == "__main__":
# ========== 6. 定义数据和优化器 ========== # ========== 6. 定义数据和优化器 ==========
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.amp.GradScaler(device_type, enabled=(args.dtype == 'float16' and device_type == 'cuda'))
optimizer = optim.AdamW(lora_params, lr=args.learning_rate) optimizer = optim.AdamW(lora_params, lr=args.learning_rate)
# ========== 7. 从ckp恢复状态 ========== # ========== 7. 从ckp恢复状态 ==========

4
trainer/train_ppo.py
View File

@ -364,9 +364,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ========== # ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu" device_type = "cuda" if "cuda" in args.device else ("mps" if "mps" in args.device else "cpu")
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) autocast_ctx = nullcontext() if device_type == "cpu" else torch.amp.autocast(device_type, dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None

6
trainer/train_pretrain.py
View File

@ -118,9 +118,9 @@ if __name__ == "__main__":
ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
# ========== 3. 设置混合精度 ========== # ========== 3. 设置混合精度 ==========
device_type = "cuda" if "cuda" in args.device else "cpu" device_type = "cuda" if "cuda" in args.device else ("mps" if "mps" in args.device else "cpu")
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) autocast_ctx = nullcontext() if device_type == "cpu" else torch.amp.autocast(device_type, dtype=dtype)
# ========== 4. 配wandb ========== # ========== 4. 配wandb ==========
wandb = None wandb = None
@ -135,7 +135,7 @@ if __name__ == "__main__":
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device) model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
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.amp.GradScaler(device_type, enabled=(args.dtype == 'float16' and device_type == 'cuda'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# ========== 6. 从ckp恢复状态 ========== # ========== 6. 从ckp恢复状态 ==========