From c8567a6f65165a57548029caff20d07bc11ee969 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Wed, 8 Jul 2026 19:41:16 +0800
Subject: [PATCH] fix: exclude embedding, lm_head, bias, and norm params from
 Muon optimizer, use AdamW

---
 scripts/tools/train.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/scripts/tools/train.py b/scripts/tools/train.py
index f6b82e3..a1a6559 100644
--- a/scripts/tools/train.py
+++ b/scripts/tools/train.py
@@ -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)