From 8fa38fb4803b8d415d3ec1b9c2adf5717adeb11c Mon Sep 17 00:00:00 2001 From: mkhona Date: Wed, 23 Sep 2026 15:33:50 -0700 Subject: [PATCH 1/5] Scale Muon normalization norm by 1.01 Signed-off-by: mkhona --- .../orthogonalized_optimizers/muon_utils.py | 11 ++++++----- tests/test_distributed_muon_utils_cpu.py | 2 +- tests/test_muon_utils.py | 4 ++-- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/emerging_optimizers/orthogonalized_optimizers/muon_utils.py b/emerging_optimizers/orthogonalized_optimizers/muon_utils.py index 2187b3e4..269d908e 100644 --- a/emerging_optimizers/orthogonalized_optimizers/muon_utils.py +++ b/emerging_optimizers/orthogonalized_optimizers/muon_utils.py @@ -141,9 +141,9 @@ def distributed_normalize_p2( x_sq_sum = (x_sq * x_sq).sum() torch.distributed.all_reduce(x_sq_sum, op=torch.distributed.ReduceOp.SUM, group=group) norm = torch.sqrt(x_sq_sum).to(x.dtype) - if not normalize_in_double: - norm.clamp_min_(eps) - return x / norm + if normalize_in_double: + return x / (1.01 * norm) + return x / (1.01 * norm + eps) def newton_schulz( @@ -213,13 +213,14 @@ def newton_schulz( X = distributed_normalize_p2(x, eps, tp_group, normalize_in_double) else: if not normalize_in_double: - X = torch.nn.functional.normalize(x, p=2, dim=(-2, -1), eps=eps) # type: ignore[arg-type] + norm = torch.linalg.vector_norm(x, dim=(-2, -1), keepdim=True) + X = x / (1.01 * norm + eps) else: # eps is ignored when normalize in double so that zero division can happen if norm is exact 0. # However, if norm is 0 in double precision, it means the entire input is 0, which usually # suggests something wrong has happened in training. So we don't guard it here. norm = torch.linalg.vector_norm(x, dim=(-2, -1), keepdim=True, dtype=torch.float64).to(x.dtype) - X = x / norm + X = x / (1.01 * norm) if coefficient_type in _COEFFICIENT_SETS: coefficient_sets = _COEFFICIENT_SETS[coefficient_type] diff --git a/tests/test_distributed_muon_utils_cpu.py b/tests/test_distributed_muon_utils_cpu.py index b100b9d9..4b3ae3a2 100644 --- a/tests/test_distributed_muon_utils_cpu.py +++ b/tests/test_distributed_muon_utils_cpu.py @@ -115,7 +115,7 @@ def test_distributed_normalize_close_to_non_distributed(self, shape): local_x = x.chunk(world_size, dim=1)[rank] dist_out = muon_utils.distributed_normalize_p2(local_x, eps=1e-7, group=torch.distributed.group.WORLD) - ref_out = torch.nn.functional.normalize(x, dim=(-2, -1), eps=1e-7) + ref_out = x / (1.01 * torch.linalg.vector_norm(x, dim=(-2, -1), keepdim=True) + 1e-7) torch.testing.assert_close(ref_out.chunk(world_size, dim=1)[rank], dist_out) diff --git a/tests/test_muon_utils.py b/tests/test_muon_utils.py index 3051e71f..de261362 100644 --- a/tests/test_muon_utils.py +++ b/tests/test_muon_utils.py @@ -54,7 +54,7 @@ def newton_schulz_ref(x: torch.Tensor, coefficient_sets: list[tuple[float, float x = x.mT # Ensure spectral norm is at most 1 - X = x / x.norm(dim=(-2, -1), keepdim=True).clamp_min_(1e-7) + X = x / (1.01 * x.norm(dim=(-2, -1), keepdim=True) + 1e-15) # Perform the NS iterations for i in range(steps): @@ -128,7 +128,7 @@ def test_preserve_values_with_underflowed_norm_in_fp64(self): norm_ref = torch.linalg.vector_norm(x, dtype=torch.double) assert norm_ref != 0 out = muon_utils.newton_schulz(x, steps=0, normalize_in_double=True) - torch.testing.assert_close(x / norm_ref, out, atol=0, rtol=1e-6) + torch.testing.assert_close(x / (1.01 * norm_ref), out, atol=0, rtol=1e-6) @parameterized.parameters( (2, 256, 256), From 0c070a6b329f3404ed2f59df0298826c24a4713f Mon Sep 17 00:00:00 2001 From: mkhona Date: Wed, 23 Sep 2026 15:38:18 -0700 Subject: [PATCH 2/5] Guard repeated Polar Express coefficients until final step Signed-off-by: mkhona --- .../orthogonalized_optimizers/muon_utils.py | 11 ++++++++++- tests/test_muon_utils.py | 4 +++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/emerging_optimizers/orthogonalized_optimizers/muon_utils.py b/emerging_optimizers/orthogonalized_optimizers/muon_utils.py index 269d908e..a17943a6 100644 --- a/emerging_optimizers/orthogonalized_optimizers/muon_utils.py +++ b/emerging_optimizers/orthogonalized_optimizers/muon_utils.py @@ -239,7 +239,16 @@ def newton_schulz( repeat_last_types = ("polar_express", "cans", "deepseekv4") iter_mode: CoeffIterMode = "repeat_last" if coefficient_type in repeat_last_types else "cycle" - coeff_iter = get_coefficient_iterator(steps, coefficient_sets, mode=iter_mode) + if coefficient_type == "polar_express" and steps > len(coefficient_sets): + a, b, c = coefficient_sets[-1] + guarded_last = (a / 1.01, b / 1.01**3, c / 1.01**5) + coeff_iter = chain( + coefficient_sets[:-1], + repeat(guarded_last, steps - len(coefficient_sets)), + (coefficient_sets[-1],), + ) + else: + coeff_iter = get_coefficient_iterator(steps, coefficient_sets, mode=iter_mode) ns_step_fn = newton_schulz_step if X.ndim == 2 else batched_newton_schulz_step # Perform the NS iterations diff --git a/tests/test_muon_utils.py b/tests/test_muon_utils.py index de261362..396f2f6f 100644 --- a/tests/test_muon_utils.py +++ b/tests/test_muon_utils.py @@ -271,7 +271,9 @@ def test_polar_express_9steps_close_to_reference(self, dim1, dim2): out_pe9 = muon_utils.newton_schulz(x, steps=9, coefficient_type="polar_express") coeff = deepcopy(muon_utils._COEFFICIENT_SETS["polar_express"]) - coeff.append(coeff[-1]) + a, b, c = coeff[-1] + coeff[-1] = (a / 1.01, b / 1.01**3, c / 1.01**5) + coeff.append((a, b, c)) out_ref = newton_schulz_ref(x, coefficient_sets=coeff) torch.testing.assert_close(out_pe9, out_ref, atol=2e-6, rtol=1e-7) From 790c6b38e5b3b3755875853e6cf7879c6ff7872a Mon Sep 17 00:00:00 2001 From: mkhona Date: Mon, 12 Oct 2026 10:20:14 -0700 Subject: [PATCH 3/5] Precompute guarded Polar Express coefficients Signed-off-by: mkhona --- .../orthogonalized_optimizers/muon_utils.py | 26 +++++++++-------- tests/test_distributed_muon_utils_cpu.py | 20 ++++++++----- tests/test_muon_utils.py | 28 +++++++++---------- 3 files changed, 41 insertions(+), 33 deletions(-) diff --git a/emerging_optimizers/orthogonalized_optimizers/muon_utils.py b/emerging_optimizers/orthogonalized_optimizers/muon_utils.py index a17943a6..269e9982 100644 --- a/emerging_optimizers/orthogonalized_optimizers/muon_utils.py +++ b/emerging_optimizers/orthogonalized_optimizers/muon_utils.py @@ -54,6 +54,9 @@ (1.8726, -1.2307, 0.3585), (1.8564, -1.2132, 0.3568), (1.8750, -1.2500, 0.3750), + # Guarded copy of the final tuple: (a / 1.01, b / 1.01**3, c / 1.01**5). + # Used for extra iterations before the unguarded final step. + (1.85643564, -1.21323768, 0.35679963), ], "cans": [ # CANS from: http://arxiv.org/abs/2506.10935 @@ -141,9 +144,9 @@ def distributed_normalize_p2( x_sq_sum = (x_sq * x_sq).sum() torch.distributed.all_reduce(x_sq_sum, op=torch.distributed.ReduceOp.SUM, group=group) norm = torch.sqrt(x_sq_sum).to(x.dtype) - if normalize_in_double: - return x / (1.01 * norm) - return x / (1.01 * norm + eps) + if not normalize_in_double: + norm.clamp_min_(eps) + return x / norm def newton_schulz( @@ -213,14 +216,13 @@ def newton_schulz( X = distributed_normalize_p2(x, eps, tp_group, normalize_in_double) else: if not normalize_in_double: - norm = torch.linalg.vector_norm(x, dim=(-2, -1), keepdim=True) - X = x / (1.01 * norm + eps) + X = torch.nn.functional.normalize(x, p=2, dim=(-2, -1), eps=eps) # type: ignore[arg-type] else: # eps is ignored when normalize in double so that zero division can happen if norm is exact 0. # However, if norm is 0 in double precision, it means the entire input is 0, which usually # suggests something wrong has happened in training. So we don't guard it here. norm = torch.linalg.vector_norm(x, dim=(-2, -1), keepdim=True, dtype=torch.float64).to(x.dtype) - X = x / (1.01 * norm) + X = x / norm if coefficient_type in _COEFFICIENT_SETS: coefficient_sets = _COEFFICIENT_SETS[coefficient_type] @@ -239,13 +241,13 @@ def newton_schulz( repeat_last_types = ("polar_express", "cans", "deepseekv4") iter_mode: CoeffIterMode = "repeat_last" if coefficient_type in repeat_last_types else "cycle" - if coefficient_type == "polar_express" and steps > len(coefficient_sets): - a, b, c = coefficient_sets[-1] - guarded_last = (a / 1.01, b / 1.01**3, c / 1.01**5) + coeff_iter: Iterator[tuple[float, float, float]] + if coefficient_type == "polar_express" and steps >= len(coefficient_sets): + # Repeat the predefined guarded tuple, then finish with the unguarded tuple. coeff_iter = chain( - coefficient_sets[:-1], - repeat(guarded_last, steps - len(coefficient_sets)), - (coefficient_sets[-1],), + coefficient_sets[:-2], + repeat(coefficient_sets[-1], steps - len(coefficient_sets) + 1), + (coefficient_sets[-2],), ) else: coeff_iter = get_coefficient_iterator(steps, coefficient_sets, mode=iter_mode) diff --git a/tests/test_distributed_muon_utils_cpu.py b/tests/test_distributed_muon_utils_cpu.py index 4b3ae3a2..2c88f8e8 100644 --- a/tests/test_distributed_muon_utils_cpu.py +++ b/tests/test_distributed_muon_utils_cpu.py @@ -102,20 +102,26 @@ def test_with_partial_tp(self, shape, tp_size): class DistributedNewtonSchulzCpuTest(parameterized.TestCase): - @parameterized.parameters( - {"shape": (21, 16)}, - {"shape": (16, 32)}, + @parameterized.product( + shape=[(21, 16), (16, 32)], + normalize_in_double=[False, True], + scale=[1.0, 1e-30], ) - def test_distributed_normalize_close_to_non_distributed(self, shape): - x = torch.randint(-5, 5, shape, device="cpu", dtype=torch.float32) + def test_distributed_normalize_close_to_non_distributed(self, shape, normalize_in_double, scale): + x = torch.randint(-5, 5, shape, device="cpu", dtype=torch.float32) * scale torch.distributed.all_reduce(x, op=torch.distributed.ReduceOp.SUM) world_size = torch.distributed.get_world_size() rank = torch.distributed.get_rank() local_x = x.chunk(world_size, dim=1)[rank] - dist_out = muon_utils.distributed_normalize_p2(local_x, eps=1e-7, group=torch.distributed.group.WORLD) - ref_out = x / (1.01 * torch.linalg.vector_norm(x, dim=(-2, -1), keepdim=True) + 1e-7) + dist_out = muon_utils.distributed_normalize_p2( + local_x, eps=1e-7, group=torch.distributed.group.WORLD, normalize_in_double=normalize_in_double + ) + if normalize_in_double: + ref_out = x / torch.linalg.vector_norm(x, dim=(-2, -1), keepdim=True, dtype=torch.float64).to(x.dtype) + else: + ref_out = torch.nn.functional.normalize(x, dim=(-2, -1), eps=1e-7) torch.testing.assert_close(ref_out.chunk(world_size, dim=1)[rank], dist_out) diff --git a/tests/test_muon_utils.py b/tests/test_muon_utils.py index 396f2f6f..a75e2cef 100644 --- a/tests/test_muon_utils.py +++ b/tests/test_muon_utils.py @@ -54,7 +54,7 @@ def newton_schulz_ref(x: torch.Tensor, coefficient_sets: list[tuple[float, float x = x.mT # Ensure spectral norm is at most 1 - X = x / (1.01 * x.norm(dim=(-2, -1), keepdim=True) + 1e-15) + X = x / x.norm(dim=(-2, -1), keepdim=True).clamp_min_(1e-7) # Perform the NS iterations for i in range(steps): @@ -128,7 +128,7 @@ def test_preserve_values_with_underflowed_norm_in_fp64(self): norm_ref = torch.linalg.vector_norm(x, dtype=torch.double) assert norm_ref != 0 out = muon_utils.newton_schulz(x, steps=0, normalize_in_double=True) - torch.testing.assert_close(x / (1.01 * norm_ref), out, atol=0, rtol=1e-6) + torch.testing.assert_close(x / norm_ref, out, atol=0, rtol=1e-6) @parameterized.parameters( (2, 256, 256), @@ -261,21 +261,21 @@ def test_polar_express_16steps_almost_orthogonal(self, size): out = muon_utils.newton_schulz(x, steps=16, coefficient_type="polar_express") assert_close_to_orthogonal(out, diag_atol=1e-5, off_diag_atol=1e-5) - @parameterized.parameters( - (511, 513), - (511, 257), - (257, 513), + @parameterized.product( + shape=[(511, 513), (511, 257), (257, 513), (2, 31, 17)], + steps=[0, 1, 7, 8, 9, 10, 16], ) - def test_polar_express_9steps_close_to_reference(self, dim1, dim2): - x = torch.randn(dim1, dim2, device=self.device, dtype=torch.float32) - out_pe9 = muon_utils.newton_schulz(x, steps=9, coefficient_type="polar_express") + def test_polar_express_close_to_reference(self, shape, steps): + x = torch.randn(*shape, device=self.device, dtype=torch.float32) + out_pe = muon_utils.newton_schulz(x, steps=steps, coefficient_type="polar_express") - coeff = deepcopy(muon_utils._COEFFICIENT_SETS["polar_express"]) - a, b, c = coeff[-1] - coeff[-1] = (a / 1.01, b / 1.01**3, c / 1.01**5) + coeff = deepcopy(muon_utils._COEFFICIENT_SETS["polar_express"][:7]) + a, b, c = (1.8750, -1.2500, 0.3750) + if steps > 8: + coeff.extend([(a / 1.01, b / 1.01**3, c / 1.01**5)] * (steps - 8)) coeff.append((a, b, c)) - out_ref = newton_schulz_ref(x, coefficient_sets=coeff) - torch.testing.assert_close(out_pe9, out_ref, atol=2e-6, rtol=1e-7) + out_ref = newton_schulz_ref(x, coefficient_sets=coeff[:steps]) + torch.testing.assert_close(out_pe, out_ref, atol=2e-6, rtol=1e-7) @parameterized.parameters( (511, 513), From 8d5898dc5d9715ec4d8b2c59755183359a1a2b55 Mon Sep 17 00:00:00 2001 From: mkhona Date: Mon, 12 Oct 2026 14:32:26 -0700 Subject: [PATCH 4/5] Reserve final coefficients in the shared iterator Signed-off-by: mkhona --- .../orthogonalized_optimizers/muon_utils.py | 37 +++++++++++-------- tests/test_muon_utils.py | 36 +++++++++++++++--- 2 files changed, 52 insertions(+), 21 deletions(-) diff --git a/emerging_optimizers/orthogonalized_optimizers/muon_utils.py b/emerging_optimizers/orthogonalized_optimizers/muon_utils.py index 269e9982..ef2492b7 100644 --- a/emerging_optimizers/orthogonalized_optimizers/muon_utils.py +++ b/emerging_optimizers/orthogonalized_optimizers/muon_utils.py @@ -23,7 +23,7 @@ __all__ = ["newton_schulz", "newton_schulz_tp", "NSCoeffT", "get_coefficient_iterator", "CoeffIterMode"] -CoeffIterMode = Literal["cycle", "repeat_last"] +CoeffIterMode = Literal["cycle", "repeat_last", "repeat_penultimate"] NSCoeffT = Literal["simple", "quintic", "polar_express", "cans", "aol", "deepseekv4", "custom", "cubic5"] @@ -53,10 +53,10 @@ (2.2779, -1.6198, 0.3985), (1.8726, -1.2307, 0.3585), (1.8564, -1.2132, 0.3568), - (1.8750, -1.2500, 0.3750), # Guarded copy of the final tuple: (a / 1.01, b / 1.01**3, c / 1.01**5). # Used for extra iterations before the unguarded final step. (1.85643564, -1.21323768, 0.35679963), + (1.8750, -1.2500, 0.3750), ], "cans": [ # CANS from: http://arxiv.org/abs/2506.10935 @@ -92,6 +92,12 @@ ], } +_COEFFICIENT_ITER_MODES: dict[NSCoeffT, CoeffIterMode] = { + "polar_express": "repeat_penultimate", + "cans": "repeat_last", + "deepseekv4": "repeat_last", +} + def get_coefficient_iterator( steps: int, @@ -106,12 +112,15 @@ def get_coefficient_iterator( mode: Iteration mode: - "cycle": After the last element, restart from the beginning. - "repeat_last": After the last element, keep yielding the last tuple. + - "repeat_penultimate": Consume the prefix, repeat the penultimate tuple for + extra steps, and finish with the last tuple. Short schedules truncate the prefix. Yields: Tuples (a, b, c) from coefficient_sets according to the specified mode. Raises: ValueError: If coefficient_sets is empty. + ValueError: If "repeat_penultimate" is used with fewer than two coefficient tuples. ValueError: If an invalid mode is provided. """ logging.debug("Iterating through %s steps with %s mode.", steps, mode) @@ -126,8 +135,16 @@ def get_coefficient_iterator( elif mode == "repeat_last": # Chain the original list with an infinite repeat of the last item base = chain(coefficient_sets, repeat(coefficient_sets[-1])) + elif mode == "repeat_penultimate": + if len(coefficient_sets) < 2: + raise ValueError("repeat_penultimate requires at least two coefficient tuples.") + base = chain( + coefficient_sets[:-2], + repeat(coefficient_sets[-2], max(0, steps - len(coefficient_sets) + 1)), + coefficient_sets[-1:], + ) else: - raise ValueError(f"Invalid mode: {mode}. Expected 'cycle' or 'repeat_last'.") + raise ValueError(f"Invalid mode: {mode}. Expected 'cycle', 'repeat_last', or 'repeat_penultimate'.") return islice(base, steps) @@ -239,18 +256,8 @@ def newton_schulz( "Use steps=5, or pass explicit custom coefficients." ) - repeat_last_types = ("polar_express", "cans", "deepseekv4") - iter_mode: CoeffIterMode = "repeat_last" if coefficient_type in repeat_last_types else "cycle" - coeff_iter: Iterator[tuple[float, float, float]] - if coefficient_type == "polar_express" and steps >= len(coefficient_sets): - # Repeat the predefined guarded tuple, then finish with the unguarded tuple. - coeff_iter = chain( - coefficient_sets[:-2], - repeat(coefficient_sets[-1], steps - len(coefficient_sets) + 1), - (coefficient_sets[-2],), - ) - else: - coeff_iter = get_coefficient_iterator(steps, coefficient_sets, mode=iter_mode) + iter_mode = _COEFFICIENT_ITER_MODES.get(coefficient_type, "cycle") + coeff_iter = get_coefficient_iterator(steps, coefficient_sets, mode=iter_mode) ns_step_fn = newton_schulz_step if X.ndim == 2 else batched_newton_schulz_step # Perform the NS iterations diff --git a/tests/test_muon_utils.py b/tests/test_muon_utils.py index a75e2cef..323f305a 100644 --- a/tests/test_muon_utils.py +++ b/tests/test_muon_utils.py @@ -269,14 +269,21 @@ def test_polar_express_close_to_reference(self, shape, steps): x = torch.randn(*shape, device=self.device, dtype=torch.float32) out_pe = muon_utils.newton_schulz(x, steps=steps, coefficient_type="polar_express") - coeff = deepcopy(muon_utils._COEFFICIENT_SETS["polar_express"][:7]) - a, b, c = (1.8750, -1.2500, 0.3750) - if steps > 8: - coeff.extend([(a / 1.01, b / 1.01**3, c / 1.01**5)] * (steps - 8)) - coeff.append((a, b, c)) - out_ref = newton_schulz_ref(x, coefficient_sets=coeff[:steps]) + coeff = list( + muon_utils.get_coefficient_iterator( + steps, muon_utils._COEFFICIENT_SETS["polar_express"], mode="repeat_penultimate" + ) + ) + out_ref = newton_schulz_ref(x, coefficient_sets=coeff) torch.testing.assert_close(out_pe, out_ref, atol=2e-6, rtol=1e-7) + def test_polar_express_guarded_coefficient_matches_safety_factor(self): + coeff = muon_utils._COEFFICIENT_SETS["polar_express"] + a, b, c = coeff[-1] + expected = torch.tensor((a / 1.01, b / 1.01**3, c / 1.01**5), dtype=torch.float32) + actual = torch.tensor(coeff[-2], dtype=torch.float32) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + @parameterized.parameters( (511, 513), (511, 257), @@ -430,6 +437,23 @@ def test_get_coefficient_iterator_empty_raises_value_error(self) -> None: with self.assertRaisesRegex(ValueError, "must be non-empty"): list(muon_utils.get_coefficient_iterator(5, [])) + @parameterized.named_parameters( + ("zero_steps", 0, []), + ("truncated_prefix", 1, [0]), + ("no_repetition", 2, [0, 2]), + ("full_schedule", 3, [0, 1, 2]), + ("one_extra_step", 4, [0, 1, 1, 2]), + ("several_extra_steps", 7, [0, 1, 1, 1, 1, 1, 2]), + ) + def test_get_coefficient_iterator_repeat_penultimate(self, steps, expected_indices): + coeff = [(1.0, 2.0, 3.0), (4.0, 5.0, 6.0), (7.0, 8.0, 9.0)] + actual = list(muon_utils.get_coefficient_iterator(steps, coeff, mode="repeat_penultimate")) + self.assertEqual(actual, [coeff[index] for index in expected_indices]) + + def test_get_coefficient_iterator_repeat_penultimate_requires_two_coefficients(self): + with self.assertRaisesRegex(ValueError, "requires at least two"): + muon_utils.get_coefficient_iterator(5, [(1.0, 2.0, 3.0)], mode="repeat_penultimate") + def test_get_coefficient_iterator_invalid_mode_raises_value_error(self) -> None: """Test that get_coefficient_iterator raises ValueError for invalid mode.""" with self.assertRaisesRegex(ValueError, "Invalid mode.*invalid"): From 0c484d0b70c2b65d46c64312ce4e9f1a18c64d87 Mon Sep 17 00:00:00 2001 From: mkhona Date: Mon, 12 Oct 2026 15:20:25 -0700 Subject: [PATCH 5/5] Round guarded Polar Express coefficients to four decimals Signed-off-by: mkhona --- emerging_optimizers/orthogonalized_optimizers/muon_utils.py | 4 ++-- tests/test_muon_utils.py | 6 ++++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/emerging_optimizers/orthogonalized_optimizers/muon_utils.py b/emerging_optimizers/orthogonalized_optimizers/muon_utils.py index f825642d..09e5bd3f 100644 --- a/emerging_optimizers/orthogonalized_optimizers/muon_utils.py +++ b/emerging_optimizers/orthogonalized_optimizers/muon_utils.py @@ -53,9 +53,9 @@ (2.2779, -1.6198, 0.3985), (1.8726, -1.2307, 0.3585), (1.8564, -1.2132, 0.3568), - # Guarded copy of the final tuple: (a / 1.01, b / 1.01**3, c / 1.01**5). + # Guarded copy of the final tuple: (a / 1.01, b / 1.01**3, c / 1.01**5), rounded to 4 decimals. # Used for extra iterations before the unguarded final step. - (1.85643564, -1.21323768, 0.35679963), + (1.8564, -1.2132, 0.3568), (1.8750, -1.2500, 0.3750), ], "cans": [ diff --git a/tests/test_muon_utils.py b/tests/test_muon_utils.py index 323f305a..004a623d 100644 --- a/tests/test_muon_utils.py +++ b/tests/test_muon_utils.py @@ -277,10 +277,12 @@ def test_polar_express_close_to_reference(self, shape, steps): out_ref = newton_schulz_ref(x, coefficient_sets=coeff) torch.testing.assert_close(out_pe, out_ref, atol=2e-6, rtol=1e-7) - def test_polar_express_guarded_coefficient_matches_safety_factor(self): + def test_polar_express_guarded_coefficient_matches_rounded_safety_factor(self): coeff = muon_utils._COEFFICIENT_SETS["polar_express"] a, b, c = coeff[-1] - expected = torch.tensor((a / 1.01, b / 1.01**3, c / 1.01**5), dtype=torch.float32) + expected = torch.tensor( + (round(a / 1.01, 4), round(b / 1.01**3, 4), round(c / 1.01**5, 4)), dtype=torch.float32 + ) actual = torch.tensor(coeff[-2], dtype=torch.float32) torch.testing.assert_close(actual, expected, atol=0, rtol=0)