wehub-resource-sync
279 lines
10 KiB
Python
279 lines
10 KiB
Python
# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
# [AUTO-GENERATED] Test file for paddle.optimizer.muon
|
|
# 覆盖模块: paddle/optimizer/muon.py
|
|
# 未覆盖行: 135,141,232,236,238,245,307,308,344,345,346,347,348,349,351,356,414,415,417,546,547,548,549,589,602,609,657,662,684,685
|
|
# Covered module: paddle/optimizer/muon.py
|
|
# Uncovered lines: 135,141,232,236,238,245,307,308,344,345,346,347,348,349,351,356,414,415,417,546,547,548,549,589,602,609,657,662,684,685
|
|
|
|
import unittest
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
from paddle.optimizer.muon import (
|
|
_NS_COEFFICIENT_SETS,
|
|
Muon,
|
|
MuonParamInfo,
|
|
_default_should_use_muon,
|
|
)
|
|
|
|
|
|
class TestDefaultShouldUseMuon(unittest.TestCase):
|
|
"""测试 _default_should_use_muon 函数
|
|
Test _default_should_use_muon function"""
|
|
|
|
def test_should_use_muon_2d(self):
|
|
"""测试2D参数应该使用 Muon
|
|
Test 2D parameter should use Muon"""
|
|
result = _default_should_use_muon("weight", (128, 64), [])
|
|
self.assertTrue(result)
|
|
|
|
def test_should_use_muon_3d(self):
|
|
"""测试3D参数应该使用 Muon
|
|
Test 3D parameter should use Muon"""
|
|
result = _default_should_use_muon("weight", (4, 128, 64), [])
|
|
self.assertTrue(result)
|
|
|
|
def test_should_not_use_muon_1d(self):
|
|
"""测试1D参数不应使用 Muon
|
|
Test 1D parameter should not use Muon"""
|
|
result = _default_should_use_muon("bias", (128,), [])
|
|
self.assertFalse(result)
|
|
|
|
def test_should_not_use_muon_4d(self):
|
|
"""测试4D参数不应使用 Muon
|
|
Test 4D parameter should not use Muon"""
|
|
result = _default_should_use_muon("conv_weight", (3, 3, 3, 3), [])
|
|
self.assertFalse(result)
|
|
|
|
def test_should_not_use_muon_exclude_pattern(self):
|
|
"""测试匹配排除模式的参数不应使用 Muon
|
|
Test parameter matching exclude pattern should not use Muon"""
|
|
result = _default_should_use_muon("embed_weight", (128, 64), ['embed'])
|
|
self.assertFalse(result)
|
|
|
|
def test_should_not_use_muon_bias(self):
|
|
"""测试 bias 参数不应使用 Muon
|
|
Test bias parameter should not use Muon"""
|
|
result = _default_should_use_muon("linear.bias", (128, 64), ['bias'])
|
|
self.assertFalse(result)
|
|
|
|
def test_should_use_muon_none_patterns(self):
|
|
"""测试 exclude_patterns 为 None 时抛出 ValueError
|
|
Test ValueError when exclude_patterns is None"""
|
|
with self.assertRaises(ValueError):
|
|
_default_should_use_muon("weight", (128, 64), None)
|
|
|
|
def test_should_use_muon_case_insensitive(self):
|
|
"""测试排除模式匹配不区分大小写
|
|
Test exclude pattern matching is case-insensitive"""
|
|
result = _default_should_use_muon("EMBED_weight", (128, 64), ['embed'])
|
|
self.assertFalse(result)
|
|
|
|
|
|
class TestNSCoefficientSets(unittest.TestCase):
|
|
"""测试 Newton-Schulz 系数集
|
|
Test Newton-Schulz coefficient sets"""
|
|
|
|
def test_simple_coefficients(self):
|
|
"""测试 simple 系数集
|
|
Test simple coefficient set"""
|
|
self.assertIn("simple", _NS_COEFFICIENT_SETS)
|
|
self.assertEqual(len(_NS_COEFFICIENT_SETS["simple"]), 1)
|
|
|
|
def test_quintic_coefficients(self):
|
|
"""测试 quintic 系数集
|
|
Test quintic coefficient set"""
|
|
self.assertIn("quintic", _NS_COEFFICIENT_SETS)
|
|
self.assertEqual(len(_NS_COEFFICIENT_SETS["quintic"]), 5)
|
|
|
|
def test_polar_express_coefficients(self):
|
|
"""测试 polar_express 系数集
|
|
Test polar_express coefficient set"""
|
|
self.assertIn("polar_express", _NS_COEFFICIENT_SETS)
|
|
self.assertEqual(len(_NS_COEFFICIENT_SETS["polar_express"]), 8)
|
|
|
|
def test_aol_coefficients(self):
|
|
"""测试 aol 系数集
|
|
Test aol coefficient set"""
|
|
self.assertIn("aol", _NS_COEFFICIENT_SETS)
|
|
self.assertEqual(len(_NS_COEFFICIENT_SETS["aol"]), 4)
|
|
|
|
def test_deepseekv4_coefficients(self):
|
|
"""测试 deepseekv4 系数集
|
|
Test deepseekv4 coefficient set"""
|
|
self.assertIn("deepseekv4", _NS_COEFFICIENT_SETS)
|
|
self.assertEqual(len(_NS_COEFFICIENT_SETS["deepseekv4"]), 10)
|
|
|
|
|
|
class TestMuonParamInfo(unittest.TestCase):
|
|
"""测试 MuonParamInfo 数据类
|
|
Test MuonParamInfo dataclass"""
|
|
|
|
def test_default_values(self):
|
|
"""测试 MuonParamInfo 默认值
|
|
Test MuonParamInfo default values"""
|
|
info = MuonParamInfo()
|
|
self.assertTrue(info.use_muon)
|
|
self.assertIsNone(info.split_concat_func)
|
|
|
|
def test_custom_values(self):
|
|
"""测试 MuonParamInfo 自定义值
|
|
Test MuonParamInfo with custom values"""
|
|
|
|
def split_fn(matrix, ortho_fn, **kwargs):
|
|
return ortho_fn(matrix)
|
|
|
|
info = MuonParamInfo(use_muon=False, split_concat_func=split_fn)
|
|
self.assertFalse(info.use_muon)
|
|
self.assertEqual(info.split_concat_func, split_fn)
|
|
|
|
|
|
class TestMuonInit(unittest.TestCase):
|
|
"""测试 Muon 优化器初始化
|
|
Test Muon optimizer initialization"""
|
|
|
|
def test_muon_none_parameters(self):
|
|
"""测试 parameters=None 时抛出 ValueError
|
|
Test ValueError when parameters is None"""
|
|
with self.assertRaises(ValueError):
|
|
Muon(parameters=None)
|
|
|
|
def test_muon_dict_parameters(self):
|
|
"""测试字典参数列表时抛出 TypeError
|
|
Test TypeError when parameters is a list of dicts"""
|
|
linear = paddle.nn.Linear(10, 10)
|
|
with self.assertRaises(TypeError):
|
|
Muon(parameters=[{'params': linear.parameters()}])
|
|
|
|
def test_muon_non_list_parameters(self):
|
|
"""测试非列表参数时抛出 TypeError
|
|
Test TypeError when parameters is not a list"""
|
|
with self.assertRaises(TypeError):
|
|
Muon(parameters="invalid")
|
|
|
|
def test_muon_invalid_grad_clip(self):
|
|
"""测试无效 grad_clip 时抛出 TypeError
|
|
Test TypeError with invalid grad_clip"""
|
|
linear = paddle.nn.Linear(10, 10)
|
|
with self.assertRaises(TypeError):
|
|
Muon(
|
|
parameters=linear.parameters(),
|
|
grad_clip="invalid_clip",
|
|
)
|
|
|
|
|
|
class TestMuonScalingFn(unittest.TestCase):
|
|
"""测试 Muon._scaling_fn 静态方法
|
|
Test Muon._scaling_fn static method"""
|
|
|
|
def test_scaling_fn_version1(self):
|
|
"""测试 version=1 的 scaling 函数
|
|
Test scaling function with version=1"""
|
|
tensor = paddle.randn([8, 4])
|
|
result = Muon._scaling_fn(tensor, version=1, extra_scale_factor=0.2)
|
|
# version 1: scale = max(1, dout/din)^0.5 * extra_scale_factor
|
|
expected_scale = max(1, 4 / 8) ** 0.5 * 0.2
|
|
np.testing.assert_allclose(
|
|
result.numpy(),
|
|
tensor.numpy() * expected_scale,
|
|
rtol=1e-5,
|
|
atol=1e-5,
|
|
)
|
|
|
|
def test_scaling_fn_version2(self):
|
|
"""测试 version=2 的 scaling 函数
|
|
Test scaling function with version=2"""
|
|
tensor = paddle.randn([4, 8])
|
|
result = Muon._scaling_fn(tensor, version=2, extra_scale_factor=1.0)
|
|
# version 2: scale = (dout/din)^0.5 * extra_scale_factor
|
|
expected_scale = (8 / 4) ** 0.5 * 1.0
|
|
np.testing.assert_allclose(
|
|
result.numpy(),
|
|
tensor.numpy() * expected_scale,
|
|
rtol=1e-5,
|
|
atol=1e-5,
|
|
)
|
|
|
|
def test_scaling_fn_version3(self):
|
|
"""测试 version=3 的 scaling 函数
|
|
Test scaling function with version=3"""
|
|
tensor = paddle.randn([4, 8])
|
|
result = Muon._scaling_fn(tensor, version=3, extra_scale_factor=0.2)
|
|
# version 3: scale = max(dout, din)^0.5 * extra_scale_factor
|
|
expected_scale = max(8, 4) ** 0.5 * 0.2
|
|
np.testing.assert_allclose(
|
|
result.numpy(),
|
|
tensor.numpy() * expected_scale,
|
|
rtol=1e-5,
|
|
atol=1e-5,
|
|
)
|
|
|
|
|
|
class TestZeropowerNewtonschulz5(unittest.TestCase):
|
|
"""测试 _zeropower_via_newtonschulz5 静态方法
|
|
Test _zeropower_via_newtonschulz5 static method"""
|
|
|
|
def test_basic_orthogonalization(self):
|
|
"""测试基本的正交化
|
|
Test basic orthogonalization"""
|
|
x = paddle.randn([4, 4], dtype='float32')
|
|
result = Muon._zeropower_via_newtonschulz5(
|
|
x, steps=3, ns_matmul_dtype=paddle.float32
|
|
)
|
|
self.assertEqual(result.shape, [4, 4])
|
|
|
|
def test_tall_matrix(self):
|
|
"""测试高矩阵 (rows > cols) 时的转置处理
|
|
Test tall matrix (rows > cols) transpose handling"""
|
|
x = paddle.randn([8, 4], dtype='float32')
|
|
result = Muon._zeropower_via_newtonschulz5(
|
|
x, steps=3, ns_matmul_dtype=paddle.float32
|
|
)
|
|
self.assertEqual(result.shape, [8, 4])
|
|
|
|
def test_wide_matrix(self):
|
|
"""测试宽矩阵 (rows < cols)
|
|
Test wide matrix (rows < cols)"""
|
|
x = paddle.randn([4, 8], dtype='float32')
|
|
result = Muon._zeropower_via_newtonschulz5(
|
|
x, steps=3, ns_matmul_dtype=paddle.float32
|
|
)
|
|
self.assertEqual(result.shape, [4, 8])
|
|
|
|
def test_custom_coeffs(self):
|
|
"""测试使用自定义系数
|
|
Test with custom coefficients"""
|
|
x = paddle.randn([4, 4], dtype='float32')
|
|
custom_coeffs = [(3.4445, -4.7750, 2.0315)]
|
|
result = Muon._zeropower_via_newtonschulz5(
|
|
x, steps=1, ns_coeffs=custom_coeffs, ns_matmul_dtype=paddle.float32
|
|
)
|
|
self.assertEqual(result.shape, [4, 4])
|
|
|
|
def test_zero_steps(self):
|
|
"""测试0步迭代 (仅归一化)
|
|
Test zero-step iteration (normalization only)"""
|
|
x = paddle.randn([4, 4], dtype='float32')
|
|
result = Muon._zeropower_via_newtonschulz5(
|
|
x, steps=0, ns_matmul_dtype=paddle.float32
|
|
)
|
|
self.assertEqual(result.shape, [4, 4])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|