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

[fix] dtype & lr

pull/825/head
jingyaogong 2025-12-09 13:01:38 +08:00
parent b6b0c8f675
commit 37300329c5
2 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
model/model_minimind.py
View File

@ -313,7 +313,7 @@ class MOEFeedForward(nn.Module):
flat_topk_idx = topk_idx.view(-1) flat_topk_idx = topk_idx.view(-1)
if self.training: if self.training:
x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0) x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)
y = torch.empty_like(x, dtype=torch.float16) y = torch.empty_like(x, dtype=x.dtype)
for i, expert in enumerate(self.experts): for i, expert in enumerate(self.experts):
y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(y.dtype) # 确保类型一致 y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(y.dtype) # 确保类型一致
y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1) y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)

3
trainer/trainer_utils.py
View File

@ -23,8 +23,7 @@ def Logger(content):
def get_lr(current_step, total_steps, lr): def get_lr(current_step, total_steps, lr):
min_lr = lr / 10 return lr*(0.1 + 0.45*(1 + math.cos(math.pi * current_step / total_steps)))
return min_lr + 0.5 * (lr - min_lr) * (1 + math.cos(math.pi * current_step / total_steps))
def init_distributed_mode(): def init_distributed_mode():