feat: update shampoo

tools/train/scalable_shampoo/README.md

@@ -4,4 +4,4 @@ Files copied from [google-research/scalable_shampoo/optax](https://github.com/go
 
 Imports have been modified to be relative.
 
-This will be replaced with `optax-shampoo` package eventually.
+This will eventually be replaced with `optax-shampoo` package.

tools/train/scalable_shampoo/distributed_shampoo.py

@@ -25,13 +25,12 @@
 # Authors: Rohan Anil (rohananil at google dot com)
 #    &     Vineet Gupta (vineet at google dot com)
 #
-
 """Distributed Shampoo Implementation."""
 
 import enum
 import functools
 import itertools
-from typing import Any, List, NamedTuple
+from typing import Any, List, NamedTuple, Tuple
 
 import chex
 import jax
@@ -43,6 +42,7 @@ from flax import struct
 from jax import lax
 
 from .quantization_utils import QuantizedValue
+from .symmetric_matrices import symmetric_matrices
 
 # Dtype for inverse-pth root routine
 # Switch to f64 if you have hardware that supports it. Enable the jax flag
@@ -141,7 +141,10 @@ class GraftingType(enum.IntEnum):
 
 
 def power_iteration(
-    matrix, num_iters=100, error_tolerance=1e-6, precision=lax.Precision.HIGHEST
+    matrix,
+    num_iters=100,
+    error_tolerance=1e-6,
+    precision=lax.Precision.HIGHEST,
 ):
     r"""Power iteration algorithm.
 
@@ -156,10 +159,10 @@ def power_iteration(
       matrix: the symmetric PSD matrix.
       num_iters: Number of iterations.
       error_tolerance: Iterative exit condition.
-      precision: precision XLA related flag, the available options are:
-        a) lax.Precision.DEFAULT (better step time, but not precise)
-        b) lax.Precision.HIGH (increased precision, slower)
-        c) lax.Precision.HIGHEST (best possible precision, slowest)
+      precision: precision XLA related flag, the available options are: a)
+        lax.Precision.DEFAULT (better step time, but not precise) b)
+        lax.Precision.HIGH (increased precision, slower) c) lax.Precision.HIGHEST
+        (best possible precision, slowest)
 
     Returns:
       eigen vector, eigen value
@@ -196,7 +199,11 @@ def power_iteration(
     return v_out, s_out
 
 
-def mat_power(mat_m, p, precision=lax.Precision.HIGHEST):
+def mat_power(
+    mat_m,
+    p,
+    precision=lax.Precision.HIGHEST,
+):
     """A simple matrix power method. M^p where p can be TracedValue."""
     power = jnp.eye(mat_m.shape[0], dtype=_MAT_INV_PTH_ROOT_DTYPE)
 
@@ -245,15 +252,19 @@ def matrix_inverse_pth_root(
       num_iters: Maximum number of iterations.
       ridge_epsilon: Ridge epsilon added to make the matrix positive definite.
       error_tolerance: Error indicator, useful for early termination.
-      precision: precision XLA related flag, the available options are:
-        a) lax.Precision.DEFAULT (better step time, but not precise)
-        b) lax.Precision.HIGH (increased precision, slower)
-        c) lax.Precision.HIGHEST (best possible precision, slowest)
+      precision: precision XLA related flag, the available options are: a)
+        lax.Precision.DEFAULT (better step time, but not precise) b)
+        lax.Precision.HIGH (increased precision, slower) c) lax.Precision.HIGHEST
+        (best possible precision, slowest)
 
     Returns:
       matrix^(-1/p)
     """
 
+    # If the input is not square, materialize it from the concatenated form.
+    if matrix.shape[0] != matrix.shape[1]:
+        matrix = symmetric_matrices.materialize_matrix_from_concat(matrix)
+
     assert matrix.shape[0] == matrix.shape[1]
 
     # We use _MAT_INV_PTH_ROOT_DTYPE for the matrix inverse pth root.
@@ -336,8 +347,8 @@ def merge_small_dims(shape_to_merge, max_dim):
     return resulting_shape
 
 
-def pad_matrix(mat, max_size):
-    """Pad a matrix to a max_size.
+def pad_square_matrix(mat, max_size):
+    """Pad a square matrix up to max_size.
 
     Args:
       mat: a matrix to pad.
@@ -346,19 +357,132 @@ def pad_matrix(mat, max_size):
     Returns:
       Given M returns [[M, 0], [0, I]]
     """
-    size = mat.shape[0]
-    assert size <= max_size
-    if size == max_size:
+    rows, cols = mat.shape
+    if rows != cols:
+        raise ValueError(
+            "Must have rows == cols, instead got " f"rows={rows}, cols={cols}"
+        )
+    if cols > max_size:
+        raise ValueError(
+            "Must have cols <= max_size. Instead got "
+            f"cols={cols}, max_size={max_size}."
+        )
+    if rows == max_size:
         return mat
-    pad_size = max_size - size
-    zs1 = jnp.zeros([size, pad_size], dtype=mat.dtype)
-    zs2 = jnp.zeros([pad_size, size], dtype=mat.dtype)
+    pad_size = max_size - rows
+
+    zs1 = jnp.zeros([rows, pad_size], dtype=mat.dtype)
+    zs2 = jnp.zeros([pad_size, rows], dtype=mat.dtype)
     eye = jnp.eye(pad_size, dtype=mat.dtype)
     mat = jnp.concatenate([mat, zs1], 1)
     mat = jnp.concatenate([mat, jnp.concatenate([zs2, eye], 1)], 0)
     return mat
 
 
+def make_sliced_padding(
+    symmetric_block_size,
+    num_blocks,
+    starting_block,
+    dtype,
+):
+    """Returns padding for symmetric block matrix.
+
+    Specifically, the padding is given concatenated rectangular matrices
+    representing the lower-triangular rows below the starting block. For example,
+    if we want to pad the symmetric matrix
+
+    M = [[A, B^T]
+         [B, C]],
+
+    the desired output (in terms of the full matrix) with num_blocks = 4 is
+
+    M_padded = [[A, B^T, 0, 0]
+                [B, C,   0, 0]
+                [0, 0,   I, 0]
+                 0, 0,   0, I].
+
+    We would represent M as the block matrix mat = [A, B, C]. In this form, the
+    additional padding to provide has form [0, 0, I, 0, 0, 0, I] (only the lower
+    triangular parts in the third and fourth rows).
+
+    Args:
+      symmetric_block_size: The size of each block.
+      num_blocks: The total number of blocks.
+      starting_block: The block where to start the padding.
+      dtype: The type to use for the blocks.
+    """
+    if starting_block == num_blocks:
+        return jnp.zeros(shape=(symmetric_block_size, 0), dtype=dtype)
+
+    blocks = []
+    for i in range(starting_block, num_blocks):
+        blocks.append(
+            jnp.zeros(
+                shape=(symmetric_block_size, symmetric_block_size * i), dtype=dtype
+            )
+        )
+        blocks.append(jnp.eye(symmetric_block_size, dtype=dtype))
+    return jnp.concatenate(blocks, axis=-1)
+
+
+def pad_block_symmetric_matrix(
+    mat,
+    symmetric_block_size,
+    max_num_blocks,
+):
+    """Returns the padded blocked symmetric matrix.
+
+    The size of the padded matrix will be:
+      [symmetric_block_size, symmetric_block_size * max_num_blocks]
+
+    The input matrix can either:
+      - Be square with size less or equal to symmetric_block_size. In this case,
+        mat will first be padded to a square matrix of size symmetric_block_size,
+        and then be padded again up to the full size of the blocked matrix.
+      - Be a rectangle with number of rows equal to block size.
+        In this case, number of columns must be a multiple of number of rows, and
+        the ratio must correspond to a block representation of a symmetric matrix.
+        That is, the ratio must have form x * (x + 1) / 2. Here, x represents the
+        number of block rows represented by the matrix.
+
+    Args:
+      mat: The input block matrix.
+      symmetric_block_size: The size of blocks.
+      max_num_blocks: The largest number of blocks to pad to.
+    """
+    rows, cols = mat.shape
+    if rows > symmetric_block_size:
+        raise ValueError(
+            "Must have rows <= symmetric_block_size. Instead got "
+            f"rows={rows}, symmetric_block_size={symmetric_block_size}."
+        )
+    if rows > cols:
+        raise ValueError(
+            "Must have rows <= cols, instead got " f"rows={rows}, cols={cols}."
+        )
+    if cols > symmetric_block_size * max_num_blocks:
+        raise ValueError(
+            "Must have cols <= symmetric_block_size * max_num_blocks "
+            f"Instead got cols={cols}, "
+            f"symmetric_block_size={symmetric_block_size}, "
+            f"max_num_blocks={max_num_blocks}."
+        )
+    if rows < symmetric_block_size:
+        mat = pad_square_matrix(mat, max_size=symmetric_block_size)
+    # Update rows and cols after possibly padding in pad_square_matrix.
+    rows, cols = mat.shape
+    assert rows == symmetric_block_size
+    assert cols % rows == 0
+    filled_blocks = cols // rows
+    padding_blocks = make_sliced_padding(
+        symmetric_block_size=symmetric_block_size,
+        num_blocks=symmetric_matrices.num_blocks_from_total_blocks(max_num_blocks),
+        starting_block=symmetric_matrices.num_blocks_from_total_blocks(filled_blocks),
+        dtype=mat.dtype,
+    )
+    return jnp.concatenate([mat, padding_blocks], axis=-1)
+
+
 def pad_vector(vec, max_size):
     """Pad a vector to a max_size.
 
@@ -694,18 +818,17 @@ def distributed_shampoo(
       num_devices_for_pjit: Number of devices to parallelize over when using pjit.
       shard_optimizer_states: Shard optimizer states to save memory in model
         parallel training.
-      best_effort_memory_usage_reduction: Best effort memory usage reduction.
-        diagonal_statistics -> jnp.bfloat16
-        momentum buffers (2x) -> jnp.int8
+      best_effort_memory_usage_reduction: Best effort memory usage reduction. -
+        diagonal_statistics -> jnp.bfloat16 - momentum buffers (2x) -> jnp.int8 -
         statistics, preconditioners -> jnp.int16 + diagonals
       inverse_failure_threshold: numerics are hard and inverses fail sometimes; we
         determine that using this threshold.
       moving_average_for_momentum: Whether to use moving average for momentum
         instead of exponential moving average.
       skip_preconditioning_dim_size_gt: Skip if preconditioning dim size is
-          greater than this value.
-      clip_by_scaled_gradient_norm: Clip by scaled gradient norm (only useful
-        when using RMSProp Grafting).
+        greater than this value.
+      clip_by_scaled_gradient_norm: Clip by scaled gradient norm (only useful when
+        using RMSProp Grafting).
       precision: precision XLA related flag, the available options are: a)
         lax.Precision.DEFAULT (better step time, but not precise) b)
         lax.Precision.HIGH (increased precision, slower) c) lax.Precision.HIGHEST
@@ -1167,7 +1290,7 @@ def distributed_shampoo(
         new_padded_statistics = []
         for stat in new_stats_flat:
             new_padded_statistics.extend(
-                [pad_matrix(stat, max_size) for stat in stat.statistics]
+                [pad_square_matrix(stat, max_size) for stat in stat.statistics]
             )
 
         # Create global stats
@@ -1388,7 +1511,7 @@ def distributed_shampoo(
         num_devices = lax.psum(1, batch_axis_name)
         num_statistics = len(statistics)
         # Pad statistics and exponents to next multiple of num_devices.
-        packed_statistics = [pad_matrix(stat, max_size) for stat in statistics]
+        packed_statistics = [pad_square_matrix(stat, max_size) for stat in statistics]
         to_pad = -num_statistics % num_devices
         packed_statistics.extend(
             [jnp.eye(max_size, dtype=packed_statistics[0].dtype) for _ in range(to_pad)]
@@ -1540,7 +1663,7 @@ def distributed_shampoo(
         # diagonals [d] f32
         # bucket_sizes [d] f32
         packed_quantized_statistics = [
-            pad_matrix(stat.quantized, max_size) for stat in statistics
+            pad_square_matrix(stat.quantized, max_size) for stat in statistics
         ]
         packed_quantized_diagonals = [
             pad_vector(stat.diagonal, max_size) for stat in statistics
@@ -1772,7 +1895,7 @@ def distributed_shampoo(
         """
         num_statistics = len(statistics)
         to_pad = -num_statistics % num_devices_for_pjit
-        padded_statistics = [pad_matrix(stat, max_size) for stat in statistics]
+        padded_statistics = [pad_square_matrix(stat, max_size) for stat in statistics]
         padded_statistics.extend(
             [jnp.eye(max_size, dtype=padded_statistics[0].dtype) for _ in range(to_pad)]
         )

tools/train/scalable_shampoo/symmetric_matrices/symmetric_matrices.py

@@ -16,7 +16,7 @@
 """JAX Ops for symmetric matrices used by the Shampoo optimizer."""
 
 import functools
-from typing import Any, List, Sequence, Union
+from typing import Any, List, Optional, Sequence, Union
 
 import jax
 import jax.numpy as jnp
@@ -192,7 +192,7 @@ def materialize_matrix(symmetric_matrix):
 @functools.partial(jax.jit, static_argnames=("num_blocks"))
 def materialize_matrix_from_concat(
     block_rows_concat,
-    num_blocks,
+    num_blocks=None,
 ):
     """Returns a materialized symmetric matrix from concatenated slices.
 
@@ -200,7 +200,11 @@ def materialize_matrix_from_concat(
       block_rows_concat: The matrix represented as the concatenated
         lower-triangular blocks.
       num_blocks: The number of block-rows used to represent the symmetric matrix.
+        If not specified, it is inferred from the shape of block_rows_concat.
     """
+    if num_blocks is None:
+        num_blocks = find_num_blocks(block_rows_concat)
+
     block_size = block_rows_concat.shape[-2]
 
     block_rows = [
@@ -251,6 +255,28 @@ def update_sliced_rows(
     )
 
 
+def num_blocks_from_total_blocks(total_blocks):
+    """Returns the number of blocks (i.e.
+
+    block rows) from the total blocks.
+
+    This is the inverse of the function x -> x*(x+1)/2.
+
+    For example, the matrix M = [[A, B^T], [B, C]] may be represented using a
+    total of 3 blocks ([A, B, C]). The number of corresponding block rows is 2.
+
+    Args:
+      total_blocks: The total blocks used to represent the matrix.
+    """
+    num_blocks = np.round((np.sqrt(8 * total_blocks + 1) - 1) / 2).astype(np.int32)
+    if (num_blocks * (num_blocks + 1)) / 2 != total_blocks:
+        raise ValueError(
+            f"total_blocks={total_blocks} does not correspond to "
+            "a symmetric matrix. It must have the form total_blocks = x*(x+1)/2."
+        )
+    return num_blocks
+
+
 def find_num_blocks(block_rows_concat):
     """Returns the number of (row) blocks representing the concatenated matrix.
 
@@ -270,11 +296,147 @@ def find_num_blocks(block_rows_concat):
     # Compute the number of square blocks used to represent the matrix.
     total_blocks = block_rows_concat.shape[-1] / block_rows_concat.shape[-2]
     # Determine the number of block rows by inverting y = x*(x+1)/2.
-    num_blocks = np.round((np.sqrt(8 * total_blocks + 1) - 1) / 2).astype(np.int32)
-    if num_blocks * (num_blocks + 1) / 2 != total_blocks:
+    return num_blocks_from_total_blocks(total_blocks)
+
+
+@functools.partial(jax.jit, static_argnames=("block_size"))
+def slice_symmetric_matrix(
+    mat,
+    block_size,
+):
+    """Returns sliced row blocks.
+
+    Args:
+      mat: A symmetric matrix.
+      block_size: The size of the row slices.
+    """
+    num_rows = mat.shape[-2]
+    num_cols = mat.shape[-1]
+    if num_rows != num_cols:
+        raise ValueError("mat is not square.")
+    if num_rows % block_size != 0:
         raise ValueError(
-            "Could not determine an appropriate number of blocks for "
-            "the concatenated matrix."
+            "block size does not evenly divide rows. "
+            f"num_rows={num_rows}, block_size={block_size}"
         )
-    else:
-        return num_blocks
+    return SlicedSymmetricMatrix(
+        block_rows=[
+            mat[
+                Ellipsis,
+                i * block_size : (i + 1) * block_size,
+                0 : (i + 1) * block_size,
+            ]
+            for i in range(num_rows // block_size)
+        ]
+    )
+
+
+@functools.partial(jax.jit, static_argnames=("block_size"))
+def slice_symmetric_matrix_concat(
+    mat,
+    block_size,
+):
+    """Returns the concatenated sliced row blocks.
+
+    Args:
+      mat: A symmetric matrix.
+      block_size: The size of the row slices.
+    """
+    sliced_symmetric_matrix = slice_symmetric_matrix(mat=mat, block_size=block_size)
+    return jnp.concatenate(sliced_symmetric_matrix.block_rows, axis=-1)
+
+
+def sliced_matrix_diag(mat):
+    """Returns the diagonal of the symmetric matrix.
+
+    Args:
+      mat: The symmetric matrix represented in concatenated block form.
+    """
+    rows, cols = mat.shape
+    total_blocks = cols // rows
+    num_blocks = num_blocks_from_total_blocks(total_blocks)
+    diags = []
+    for i in range(num_blocks):
+        last_index = rows * ((i + 2) * (i + 1)) // 2
+        first_index = last_index - rows
+        diags.append(jnp.diag(mat[Ellipsis, first_index:last_index]))
+    return jnp.concatenate(diags, axis=-1)
+
+
+def diag_as_concat(diag, block_size):
+    """Returns the representation of a diagonal matrix in symmetric block form.
+
+    Args:
+      diag: The 1D array for the diagonals.
+      block_size: The size of blocks to use. Must divide the length of diag.
+    """
+    assert len(diag.shape) == 1  # diag must be 1D.
+    assert len(diag) % block_size == 0
+    num_diag_blocks = len(diag) // block_size
+    blocks = []
+    for i in range(num_diag_blocks):
+        blocks.append(jnp.zeros(shape=(block_size, block_size * i), dtype=diag.dtype))
+        blocks.append(jnp.diag(diag[i * block_size : (i + 1) * block_size]))
+    return jnp.concatenate(blocks, axis=-1)
+
+
+def row_abs_maxes(mat):
+    """Returns the max of the absolute values of the rows of the full matrix.
+
+    For example the symmetric matrix M = [[1, 6], [6, 2]] is represented using
+    mat = [1, 6, 2] with block_size = 1. In this case the function returns the
+    aboslute row maxes of the original symmetric matrix, [6, 6].
+
+    Args:
+      mat: The symmetric matrix represented as the concatenated blocks.
+    """
+    rows, cols = mat.shape
+
+    # Find col and row max for each block.
+    col_maxes = []
+    row_maxes = []
+    for i in range(cols // rows):
+        block = jnp.abs(mat[Ellipsis, i * rows : (i + 1) * rows])
+        col_maxes.append(jnp.max(block, axis=1))
+        row_maxes.append(jnp.max(block, axis=0))
+
+    # global row max from block maxes.
+    num_blocks = num_blocks_from_total_blocks(cols // rows)
+    maxes = []
+    for i in range(num_blocks):
+        maxes.append(
+            jnp.concatenate(
+                row_maxes[(i * (i + 1) // 2) : ((i + 2) * (i + 1) // 2)]
+                + [
+                    col_maxes[((j + 1) * (j + 2)) // 2 - (j - i + 1)]
+                    for j in range(i + 1, num_blocks)
+                ],
+                axis=-1,
+            )
+        )
+
+    return jnp.max(jnp.stack(maxes), axis=0)
+
+
+def times_vector(mat, vec):
+    """Returns the symmetric block-concatenated matrix multiplied by a vector.
+
+    Specifically, each value in the vector is multiplied by a row of the full
+    matrix. That is, the vector is broadcast and multiplied element-wise. Note
+    this would be the transpose of full_mat * vec if full_mat represented the full
+    symmetric matrix.
+
+    Args:
+      mat: The symmetric matrix represented as the concatenated blocks.
+      vec: The vector, having the same dimension as the materialized matrix.
+    """
+    rows, cols = mat.shape
+    num_blocks = num_blocks_from_total_blocks(cols // rows)
+    multiplied = []
+    for i in range(num_blocks):
+        mat_block = mat[
+            Ellipsis, rows * ((i + 1) * i) // 2 : rows * ((i + 1) * (i + 2)) // 2
+        ]
+        vec_block = vec[Ellipsis, rows * i : rows * (i + 1)]
+        multiplied.append(jnp.einsum("...ij,...i->ij", mat_block, vec_block))
+    return jnp.concatenate(multiplied, axis=-1)