diff --git a/emerging_optimizers/orthogonalized_optimizers/muon_utils.py b/emerging_optimizers/orthogonalized_optimizers/muon_utils.py index bc003c9e..09e5bd3f 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,6 +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), rounded to 4 decimals. + # Used for extra iterations before the unguarded final step. + (1.8564, -1.2132, 0.3568), (1.8750, -1.2500, 0.3750), ], "cans": [ @@ -89,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, @@ -103,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) @@ -123,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) @@ -236,8 +256,7 @@ 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" + 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 diff --git a/tests/test_distributed_muon_utils_cpu.py b/tests/test_distributed_muon_utils_cpu.py index b100b9d9..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 = torch.nn.functional.normalize(x, dim=(-2, -1), eps=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 3051e71f..004a623d 100644 --- a/tests/test_muon_utils.py +++ b/tests/test_muon_utils.py @@ -261,19 +261,30 @@ 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"]) - coeff.append(coeff[-1]) + 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_pe9, out_ref, atol=2e-6, rtol=1e-7) + torch.testing.assert_close(out_pe, out_ref, atol=2e-6, rtol=1e-7) + + 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( + (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) @parameterized.parameters( (511, 513), @@ -428,6 +439,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"):