From 9d36f14782fd7e5d0de57d95c3a2600d96e5dbfc Mon Sep 17 00:00:00 2001 From: LCStayingdullCircuit <2216215909@qq.com> Date: Mon, 13 Jul 2026 14:23:45 +0800 Subject: [PATCH 1/3] =?UTF-8?q?fix:=20muon=E6=94=AF=E6=8C=81lr=5Fratio?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- python/paddle/optimizer/muon.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/python/paddle/optimizer/muon.py b/python/paddle/optimizer/muon.py index 1251abac5f506b..72065f20ab64a3 100644 --- a/python/paddle/optimizer/muon.py +++ b/python/paddle/optimizer/muon.py @@ -606,6 +606,9 @@ def ortho_fn(m): self._master_weights[param.name] if find_master else None ) + lr_ratio = 1.0 if self._lr_ratio is None else self._lr_ratio(param) + effective_lr = lr * lr_ratio + with_decay = True if ( self._apply_decay_param_fun is not None @@ -614,11 +617,11 @@ def ortho_fn(m): with_decay = False if with_decay and weight_decay > 0: if find_master: - master_weight.scale_(1.0 - lr * weight_decay) + master_weight.scale_(1.0 - effective_lr * weight_decay) else: - param.scale_(1.0 - lr * weight_decay) + param.scale_(1.0 - effective_lr * weight_decay) - final_step = orthogonal_update * lr + final_step = orthogonal_update * effective_lr if find_master: master_weight.subtract_(final_step) From c4bfcaf4e4ca0733ece65816aa2930b330b0d271 Mon Sep 17 00:00:00 2001 From: LCStayingdullCircuit <2216215909@qq.com> Date: Mon, 13 Jul 2026 16:10:54 +0800 Subject: [PATCH 2/3] =?UTF-8?q?fix:=20muon=E6=94=AF=E6=8C=81lr=5Fratio?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- test/ai_edited_test/test_ai_optimizer_muon.py | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/test/ai_edited_test/test_ai_optimizer_muon.py b/test/ai_edited_test/test_ai_optimizer_muon.py index 3ff6595e8e831e..5932f3a3f858b1 100644 --- a/test/ai_edited_test/test_ai_optimizer_muon.py +++ b/test/ai_edited_test/test_ai_optimizer_muon.py @@ -274,5 +274,38 @@ def test_zero_steps(self): self.assertEqual(result.shape, [4, 4]) +class TestMuonLrRatio(unittest.TestCase): + """测试 Muon 优化器的 lr_ratio + Test lr_ratio in Muon optimizer""" + + def test_freeze_parameter(self): + """测试 lr_ratio=0 时冻结指定的 Muon 参数 + Test that lr_ratio=0 freezes the selected Muon parameter""" + paddle.seed(2026) + frozen = paddle.create_parameter(shape=[4, 4], dtype='float32') + trainable = paddle.create_parameter(shape=[4, 4], dtype='float32') + + optimizer = Muon( + learning_rate=0.02, + parameters=[frozen, trainable], + lr_ratio=lambda param: 0.0 if param.name == frozen.name else 1.0, + weight_decay=0.01, + muon_param_info_map={ + frozen.name: MuonParamInfo(use_muon=True), + trainable.name: MuonParamInfo(use_muon=True), + }, + ns_matmul_dtype=paddle.float32, + ) + + frozen_before = frozen.numpy().copy() + trainable_before = trainable.numpy().copy() + loss = frozen.sum() + trainable.sum() + loss.backward() + optimizer.step() + + np.testing.assert_array_equal(frozen.numpy(), frozen_before) + self.assertFalse(np.array_equal(trainable.numpy(), trainable_before)) + + if __name__ == '__main__': unittest.main() From 85570e578270f30f7f983a4619a72ad2bad0fc33 Mon Sep 17 00:00:00 2001 From: adam_xiaoyao Date: Fri, 17 Jul 2026 10:59:06 +0800 Subject: [PATCH 3/3] test: add lr_ratio regression test for Muon optimizer --- test/ai_edited_test/test_ai_optimizer_muon.py | 40 +++++++++++++++++-- 1 file changed, 36 insertions(+), 4 deletions(-) diff --git a/test/ai_edited_test/test_ai_optimizer_muon.py b/test/ai_edited_test/test_ai_optimizer_muon.py index 5932f3a3f858b1..950b880532c174 100644 --- a/test/ai_edited_test/test_ai_optimizer_muon.py +++ b/test/ai_edited_test/test_ai_optimizer_muon.py @@ -275,12 +275,10 @@ def test_zero_steps(self): class TestMuonLrRatio(unittest.TestCase): - """测试 Muon 优化器的 lr_ratio - Test lr_ratio in Muon optimizer""" + """Test lr_ratio in Muon optimizer.""" def test_freeze_parameter(self): - """测试 lr_ratio=0 时冻结指定的 Muon 参数 - Test that lr_ratio=0 freezes the selected Muon parameter""" + """Test that lr_ratio=0 freezes the selected Muon parameter.""" paddle.seed(2026) frozen = paddle.create_parameter(shape=[4, 4], dtype='float32') trainable = paddle.create_parameter(shape=[4, 4], dtype='float32') @@ -306,6 +304,40 @@ def test_freeze_parameter(self): np.testing.assert_array_equal(frozen.numpy(), frozen_before) self.assertFalse(np.array_equal(trainable.numpy(), trainable_before)) + def test_muon_lr_ratio_scales_muon_update(self): + """lr_ratio=0.0 freezes the param; lr_ratio=0.5 update is half of lr_ratio=1.0.""" + paddle.disable_static() + weight_np = np.array([[0.2, -0.4], [0.6, 0.8]], dtype="float32") + grad_np = np.array([[0.1, 0.3], [-0.2, 0.4]], dtype="float32") + + def run(ratio): + p = paddle.create_parameter(shape=[2, 2], dtype="float32") + p.set_value(weight_np) + p.grad = paddle.to_tensor(grad_np) + opt = Muon( + parameters=[p], + learning_rate=0.01, + weight_decay=0.01, + ns_steps=1, + ns_matmul_dtype=paddle.float32, + muon_param_info_map={p.name: MuonParamInfo(use_muon=True)}, + lr_ratio=lambda _: ratio, + ) + opt.step() + return p.numpy() + + full = run(1.0) + half = run(0.5) + zero = run(0.0) + + np.testing.assert_allclose(zero, weight_np, rtol=1e-6, atol=1e-6) + np.testing.assert_allclose( + weight_np - half, + 0.5 * (weight_np - full), + rtol=1e-5, + atol=1e-6, + ) + if __name__ == '__main__': unittest.main()