Fix Polar Express Coefficient Iterator for >8 step polar express (#303)

Signed-off-by: mkhona <mkhona@nvidia.com>

Author: mkhona-nvidia

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

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)
 

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"):