remove max precondition dim

Signed-off-by: Hao Wu <skyw@nvidia.com>

Author: skyw

emerging_optimizers/soap/soap.py

@@ -61,7 +61,6 @@ class SOAP(optim.Optimizer):
         precondition_warmup_steps: How many steps to warm up the preconditioner (i.e. update every step)
         adam_warmup_steps: How many steps to skip preconditioning in the beginning (i.e. use standard AdamW updates)
         precondition_1d: Whether to precondition 1D gradients (like biases).
-        max_precond_dim: Maximum dimension of the preconditioner matrices. Skips preconditioning if any tensor dimension exceeds.
         trace_normalization: Whether to normalize update by the trace of the kronecker factor matrix
         normalize_preconditioned_grads: Whether to normalize preconditioned gradients per layer
         correct_bias: Whether to use bias correction in Inner Adam and Kronecker factor matrices EMA
@@ -91,7 +90,6 @@ def __init__(
         precondition_warmup_steps: int = 0,
         adam_warmup_steps: int = 1,
         precondition_1d: bool = False,
-        max_precond_dim: int = 8192,
         trace_normalization: bool = False,
         normalize_preconditioned_grads: bool = False,
         correct_bias: bool = True,
@@ -141,7 +139,6 @@ def __init__(
             "precondition_warmup_steps": precondition_warmup_steps,
             "adam_warmup_steps": adam_warmup_steps,
             "precondition_1d": precondition_1d,
-            "max_precond_dim": max_precond_dim,
             "trace_normalization": trace_normalization,
             "normalize_preconditioned_grads": normalize_preconditioned_grads,
             "use_nesterov": use_nesterov,
@@ -194,7 +191,6 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
                     state["GG"] = init_kronecker_factors(
                         grad,
                         precondition_1d=group["precondition_1d"],
-                        max_precond_dim=group["max_precond_dim"],
                     )
 
                     # Update preconditioner matrices with gradient statistics, do not use shampoo_beta for EMA at first step
@@ -204,7 +200,6 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
                             grad=grad,
                             shampoo_beta=0.0,
                             precondition_1d=group["precondition_1d"],
-                            max_precond_dim=group["max_precond_dim"],
                         )
 
                 # Increment step counter
@@ -284,7 +279,6 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
                         grad=grad,
                         shampoo_beta=shampoo_beta,
                         precondition_1d=group["precondition_1d"],
-                        max_precond_dim=group["max_precond_dim"],
                     )
                 torch.cuda.nvtx.range_pop()
 
@@ -330,7 +324,6 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
 def init_kronecker_factors(
     grad: torch.Tensor,
     precondition_1d: bool = False,
-    max_precond_dim: int = 8192,
 ) -> List[torch.Tensor]:
     """Initializes the kronecker factor matrices for the SOAP optimizer.
 
@@ -354,8 +347,6 @@ def init_kronecker_factors(
             The shape of this tensor determines the size of the kronecker factor matrices.
         precondition_1d: Whether to create kronecker factor matrices for 1D tensors
             (like biases). If False, 1D tensors will skip preconditioning.
-        max_precond_dim: Maximum dimension of the preconditioner matrices.
-            Skips preconditioning if any tensor dimension exceeds.
 
     Returns:
         List[torch.Tensor]: List of kronecker factor matrices (L and R in paper).
@@ -387,21 +378,11 @@ def init_kronecker_factors(
         else:
             # Create a square preconditioner matrix for 1D tensors
             size = grad.shape[0]
-            if size > max_precond_dim:
-                # if tensor dimension is larger than max_precond_dim, skip preconditioning this dimension
-                # append empty tensor to kronecker_factor_list so that subsequent check that use numel() to check if preconditioner is initialized will not fail
-                kronecker_factor_list.append(torch.empty(0, device=grad.device))
-            else:
-                kronecker_factor_list.append(torch.zeros(size, size, device=grad.device))
+            kronecker_factor_list.append(torch.zeros(size, size, device=grad.device))
     else:
         # Create a square kronecker factor matrix for each dimension
         for size in grad.shape:
-            if size > max_precond_dim:
-                # append empty tensor to kronecker_factor_list so that subsequent check that use numel() to check if preconditioner is initialized will not fail
-                # skip preconditioning this dimension
-                kronecker_factor_list.append(torch.empty(0, device=grad.device))
-            else:
-                kronecker_factor_list.append(torch.zeros(size, size, device=grad.device))
+            kronecker_factor_list.append(torch.zeros(size, size, device=grad.device))
 
     return kronecker_factor_list
 
@@ -412,7 +393,6 @@ def update_kronecker_factors(
     grad: torch.Tensor,
     shampoo_beta: float,
     precondition_1d: bool = False,
-    max_precond_dim: int = 8192,
 ) -> None:
     """Updates the preconditioner matrices using gradient outer products.
 
@@ -429,8 +409,6 @@ def update_kronecker_factors(
             Controls how much weight to give to new vs old gradient statistics.
         precondition_1d: Whether to apply preconditioning to 1D tensors (like biases).
             If False, 1D tensors will skip preconditioning.
-        max_precond_dim: Maximum dimension of the preconditioner matrices.
-            Skips preconditioning if any tensor dimension exceeds.
 
     Example:
         >>> grad = torch.randn(10, 20)
@@ -446,20 +424,22 @@ def update_kronecker_factors(
             kronecker_factor_list[0].lerp_(outer_product, 1 - shampoo_beta)
         else:
             # For 1D tensors, skip preconditioning
+            logging.error(
+                "1D tensor is passed to update_kronecker_factors, but precondition_1d is not set to True, skipping preconditioning."
+            )
             return
     else:
         # For higher dimensional tensors, compute outer products for each dimension
         for idx, dim_size in enumerate(grad.shape):
-            if dim_size <= max_precond_dim:
-                # Compute outer product by contracting all dimensions except idx
-                contract_dims = [*chain(range(idx), range(idx + 1, grad.dim()))]
-                outer_product = torch.tensordot(
-                    grad,
-                    grad,
-                    dims=[contract_dims] * 2,
-                )
-                # Update the corresponding Kronecker factor
-                kronecker_factor_list[idx].lerp_(outer_product, 1 - shampoo_beta)
+            # Compute outer product by contracting all dimensions except idx
+            contract_dims = [*chain(range(idx), range(idx + 1, grad.dim()))]
+            outer_product = torch.tensordot(
+                grad,
+                grad,
+                dims=[contract_dims] * 2,
+            )
+            # Update the corresponding Kronecker factor
+            kronecker_factor_list[idx].lerp_(outer_product, 1 - shampoo_beta)
 
 
 @torch.no_grad()  # type: ignore[misc]

tests/test_soap_functions.py

@@ -35,28 +35,12 @@ def test_init_preconditioner_multidim_tensor_shapes(self) -> None:
         """Tests init_preconditioner with a multi-dimensional tensor."""
         grad = torch.randn(3, 4, 5)
         state: dict[str, Any] = {}
-        # No merge_dims: each dimension gets its own preconditioner unless dimension > max_precond_dim.
-        state["GG"] = init_kronecker_factors(grad, precondition_1d=False, max_precond_dim=8192)
+        state["GG"] = init_kronecker_factors(grad, precondition_1d=False)
         self.assertEqual(len(state["GG"]), grad.dim())
         self.assertEqual(state["GG"][0].shape, (3, 3))
         self.assertEqual(state["GG"][1].shape, (4, 4))
         self.assertEqual(state["GG"][2].shape, (5, 5))
 
-    def test_init_kronecker_factors_max_precond_dim(self) -> None:
-        """Tests init_kronecker_factors respects max_precond_dim."""
-        max_dim = 8
-        grad = torch.randn(3, max_dim + 2, 5)  # Second dimension exceeds max_dim
-        kronecker_factors = init_kronecker_factors(grad, precondition_1d=False, max_precond_dim=max_dim)
-
-        self.assertEqual(len(kronecker_factors), grad.dim())
-        # Dimension 0 (size 3) <= max_dim
-        self.assertEqual(kronecker_factors[0].shape, (3, 3))
-        # Dimension 1 (size max_dim + 2) > max_dim -> Should be empty
-        self.assertEqual(kronecker_factors[1].shape, (0,))
-        self.assertEqual(kronecker_factors[1].numel(), 0)
-        # Dimension 2 (size 5) <= max_dim
-        self.assertEqual(kronecker_factors[2].shape, (5, 5))
-
     @parameterized.parameters(
         (1,),
         (2,),
@@ -97,14 +81,13 @@ def test_adam_warmup_steps(self, adam_warmup_steps: int) -> None:
             self.assertEqual(state["Q"][1].shape, (3, 3))
 
     def test_update_kronecker_factors(self) -> None:
-        """Tests update_kronecker_factors, including max_precond_dim handling."""
         max_dim = 8
         shampoo_beta = 0.9
         dim0, dim1, dim2 = 3, max_dim + 2, 5
         grad = torch.randn(dim0, dim1, dim2)
 
         # Initialize factors
-        initial_factors = init_kronecker_factors(grad, precondition_1d=False, max_precond_dim=max_dim)
+        initial_factors = init_kronecker_factors(grad, precondition_1d=False)
 
         kronecker_factors = [f.clone() for f in initial_factors]
 
@@ -113,25 +96,15 @@ def test_update_kronecker_factors(self) -> None:
             grad=grad,
             shampoo_beta=shampoo_beta,
             precondition_1d=False,
-            max_precond_dim=max_dim,
         )
 
         self.assertEqual(len(kronecker_factors), grad.dim())
 
-        # Dimension 0 (size 3) <= max_dim: Should be updated
         contract_dims_0 = [1, 2]
         outer_product_0 = torch.tensordot(grad, grad, dims=[contract_dims_0] * 2)
         expected_factor_0 = initial_factors[0] * shampoo_beta + outer_product_0 * (1 - shampoo_beta)
         torch.testing.assert_close(kronecker_factors[0], expected_factor_0, atol=1e-6, rtol=1e-6)
 
-        # Dimension 1 (size 10) > max_dim: Should NOT be updated (still empty)
-        self.assertEqual(kronecker_factors[1].shape, (0,))
-        self.assertEqual(kronecker_factors[1].numel(), 0)
-
-        # Check it's the same object or has same properties as initial empty tensor
-        self.assertTrue(torch.equal(kronecker_factors[1], initial_factors[1]))
-
-        # Dimension 2 (size 5) <= max_dim: Should be updated
         contract_dims_2 = [0, 1]
         outer_product_2 = torch.tensordot(grad, grad, dims=[contract_dims_2] * 2)
         expected_factor_2 = initial_factors[2] * shampoo_beta + outer_product_2 * (1 - shampoo_beta)