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

2026-07-08 19:42:20 +08:00
parent 8035be9b1f
commit c8567a6f65
1 changed files with 7 additions and 1 deletions
@@ -43,7 +43,13 @@ class MuonMix(optim.Optimizer):
        for name, param in model.named_parameters():
            if not param.requires_grad:
                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)
            else:
                other_params.append(param)