ViperEkura
/
AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: exclude embedding, lm_head, bias, and norm params from Muon optimizer, use AdamW

This commit is contained in:
ViperEkura 2026-07-08 19:41:16 +08:00
parent 8035be9b1f
commit c8567a6f65
1 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
scripts/tools/train.py
View File

@ -43,7 +43,13 @@ class MuonMix(optim.Optimizer):
for name, param in model.named_parameters(): for name, param in model.named_parameters():
if not param.requires_grad: if not param.requires_grad:
continue continue
if param.dim() >= 2 and "norm" not in name and "bias" not in name: if (
param.dim() >= 2
and "norm" not in name
and "bias" not in name
and "embed" not in name
and "lm_head" not in name
):
matrix_params.append(param) matrix_params.append(param)
else: else:
other_params.append(param) other_params.append(param)