Open-sourced update on 03/25/2024 (facebookresearch#97)

Summary:
Pull Request resolved: facebookresearch#97

1. Add `ShampooPreconditionerConfig.inverse_exponent_override` to merge `inv_root_override`, `ShampooPreconditionerConfig.ignored_dims`, and `ShampooPreconditionerConfig.amortized_computation_config.exponent_multiplier` in Shampoo.
2. Add `EigenvalueCorrectedShampooPreconditionerConfig.inverse_exponent_override` to replace `inv_root_override`, and push `PreconditionerConfig.ignored_dims` to `EigenvalueCorrectedShampooPreconditionerConfig. ignored_basis_change_dims` in eigenvalue-corrected Shampoo.

Reviewed By: anana10c

Differential Revision: D71823575

fbshipit-source-id: 14743fdddebb78a89005af01a61112fdbca9a3ef

Author: tsunghsienlee

distributed_shampoo/distributed_shampoo.py

@@ -8,7 +8,7 @@
 """
 
 import logging
-from collections.abc import Callable, Iterator, Sequence
+from collections.abc import Callable, Iterator
 from copy import deepcopy
 from functools import partial
 from typing import Any
@@ -36,7 +36,6 @@
     GraftingConfig,
     HSDPShampooConfig,
     HybridShardShampooConfig,
-    INV_ROOT_OVERRIDE,
     LR,
     MASKED_BLOCKED_GRADS,
     MASKED_BLOCKED_PARAMS,
@@ -265,12 +264,6 @@ class DistributedShampoo(torch.optim.Optimizer):
             (Default: 1)
         start_preconditioning_step (int): Iteration to start computing inverse preconditioner. If -1, uses
             the same value as precondition_frequency. (Default: -1)
-        inv_root_override (int, Sequence[int]): Inverse root to use in Shampoo. If a list [l1, l2, ..., lp], then we will
-            use -1 / l1 for 1-D tensor (vectors), -1 / l2 for 2-D tensors (matrices), and so on. If the order of the
-            tensor exceeds the order of the tensor, reverts to the default value. If 0 is used, uses the default inverse
-            root -1 / (2 * o), where o is the order of the tensor. If preconditioner_config is an instance of
-            EigenvalueCorrectedShampooPreconditionerConfig, the default is -1 / 2.
-            (Default: 0)
         use_nesterov (bool): Flag for using Nesterov momentum. (default: False)
         use_bias_correction (bool): Flag for using bias correction. (Default: True)
         use_decoupled_weight_decay (bool): Flag for using AdamW-style decoupled weight decay. (Default: True)
@@ -303,7 +296,6 @@ def __init__(
         max_preconditioner_dim: int = 1024,
         precondition_frequency: int = 1,
         start_preconditioning_step: int = -1,
-        inv_root_override: int | Sequence[int] = 0,
         use_nesterov: bool = False,
         use_bias_correction: bool = True,
         use_decoupled_weight_decay: bool = True,
@@ -357,16 +349,6 @@ def __init__(
             raise ValueError(
                 f"Invalid start preconditioning step: {start_preconditioning_step}. Must be >= -1."
             )
-        if isinstance(inv_root_override, Sequence):
-            if not all(e >= 0 for e in inv_root_override):
-                raise ValueError(
-                    f"Invalid exponent override list: {inv_root_override}. All values must be >= 0."
-                )
-        else:
-            if not inv_root_override >= 0:
-                raise ValueError(
-                    f"Invalid exponent override: {inv_root_override}. Must be >= 0."
-                )
 
         # Provide warning/error for start_preconditioning_step.
         if start_preconditioning_step == -1:
@@ -387,10 +369,17 @@ def __init__(
                 "Continuing without using momentum or Nesterov acceleration..."
             )
 
-        # Check potential conflict between preconditioner_config.ignored_dims and inv_root_override.
-        if preconditioner_config.ignored_dims != [] and inv_root_override != 0:
+        # No use of preconditioner_config.amortized_computation_config.exponent_multiplier.
+        if (
+            getattr(
+                preconditioner_config.amortized_computation_config,
+                "exponent_multiplier",
+                1.0,
+            )
+            != 1.0
+        ):
             raise ValueError(
-                f"{preconditioner_config.ignored_dims=} is not supported when {inv_root_override=} is not set to 0. Please set {inv_root_override=} to 0 if you set {preconditioner_config.ignored_dims=}."
+                "preconditioner_config.amortized_computation_config.exponent_multiplier is not supported. Please use PreconditionerConfig.inverse_exponent_override instead."
             )
 
         super().__init__(
@@ -406,7 +395,6 @@ def __init__(
                 MAX_PRECONDITIONER_DIM: max_preconditioner_dim,
                 PRECONDITION_FREQUENCY: precondition_frequency,
                 START_PRECONDITIONING_STEP: start_preconditioning_step,
-                INV_ROOT_OVERRIDE: inv_root_override,
                 USE_NESTEROV: use_nesterov,
                 USE_BIAS_CORRECTION: use_bias_correction,
                 USE_DECOUPLED_WEIGHT_DECAY: use_decoupled_weight_decay,
@@ -516,7 +504,6 @@ def _instantiate_shampoo_preconditioner_list(self) -> None:
                 preconditioner_config=group[PRECONDITIONER_CONFIG],
                 beta2=group[BETAS][1],
                 epsilon=group[EPSILON],
-                inv_root_override=group[INV_ROOT_OVERRIDE],
                 use_bias_correction=group[USE_BIAS_CORRECTION],
                 factor_matrix_dtype=group[PRECONDITIONER_DTYPE],
             )

distributed_shampoo/examples/ddp_cifar10_example.py

@@ -108,7 +108,6 @@
         max_preconditioner_dim=args.max_preconditioner_dim,
         precondition_frequency=args.precondition_frequency,
         start_preconditioning_step=args.start_preconditioning_step,
-        inv_root_override=args.inv_root_override,
         exponent_multiplier=args.exponent_multiplier,
         use_nesterov=args.use_nesterov,
         use_bias_correction=args.use_bias_correction,

distributed_shampoo/examples/default_cifar10_example.py

@@ -122,7 +122,6 @@ def train_default_model(
         max_preconditioner_dim=args.max_preconditioner_dim,
         precondition_frequency=args.precondition_frequency,
         start_preconditioning_step=args.start_preconditioning_step,
-        inv_root_override=args.inv_root_override,
         exponent_multiplier=args.exponent_multiplier,
         use_nesterov=args.use_nesterov,
         use_bias_correction=args.use_bias_correction,

distributed_shampoo/examples/fsdp_cifar10_example.py

@@ -102,7 +102,6 @@
         max_preconditioner_dim=args.max_preconditioner_dim,
         precondition_frequency=args.precondition_frequency,
         start_preconditioning_step=args.start_preconditioning_step,
-        inv_root_override=args.inv_root_override,
         exponent_multiplier=args.exponent_multiplier,
         use_nesterov=args.use_nesterov,
         use_bias_correction=args.use_bias_correction,

distributed_shampoo/examples/fully_shard_cifar10_example.py

@@ -124,7 +124,6 @@ def create_model_and_optimizer_and_loss_fn(args, device):
         max_preconditioner_dim=args.max_preconditioner_dim,
         precondition_frequency=args.precondition_frequency,
         start_preconditioning_step=args.start_preconditioning_step,
-        inv_root_override=args.inv_root_override,
         exponent_multiplier=args.exponent_multiplier,
         use_nesterov=args.use_nesterov,
         use_bias_correction=args.use_bias_correction,

distributed_shampoo/examples/hsdp_cifar10_example.py

@@ -115,7 +115,6 @@
         max_preconditioner_dim=args.max_preconditioner_dim,
         precondition_frequency=args.precondition_frequency,
         start_preconditioning_step=args.start_preconditioning_step,
-        inv_root_override=args.inv_root_override,
         exponent_multiplier=args.exponent_multiplier,
         use_nesterov=args.use_nesterov,
         use_bias_correction=args.use_bias_correction,

distributed_shampoo/examples/trainer_utils.py

@@ -371,7 +371,6 @@ def instantiate_optimizer(
     max_preconditioner_dim: int,
     precondition_frequency: int,
     start_preconditioning_step: int,
-    inv_root_override: int,
     exponent_multiplier: float,
     use_nesterov: bool,
     use_bias_correction: bool,
@@ -423,7 +422,6 @@ def instantiate_optimizer(
             max_preconditioner_dim=max_preconditioner_dim,
             precondition_frequency=precondition_frequency,
             start_preconditioning_step=start_preconditioning_step,
-            inv_root_override=inv_root_override,
             use_nesterov=use_nesterov,
             use_bias_correction=use_bias_correction,
             use_decoupled_weight_decay=use_decoupled_weight_decay,

distributed_shampoo/gpu_tests/shampoo_eigenvalue_correction_test.py

@@ -63,7 +63,9 @@ def test_adagrad_eigenvalue_correction_on_quadratic(self) -> None:
                 (math.inf, DefaultSOAPConfig),
                 (
                     1,
-                    EigenvalueCorrectedShampooPreconditionerConfig(ignored_dims=[0, 1]),
+                    EigenvalueCorrectedShampooPreconditionerConfig(
+                        ignored_basis_change_dims={0: [0], 1: [0], 2: [0, 1]}
+                    ),
                 ),
             ),
         ):
@@ -113,7 +115,9 @@ def test_adam_eigenvalue_correction_on_quadratic(self) -> None:
                 (math.inf, DefaultSOAPConfig),
                 (
                     1,
-                    EigenvalueCorrectedShampooPreconditionerConfig(ignored_dims=[0, 1]),
+                    EigenvalueCorrectedShampooPreconditionerConfig(
+                        ignored_basis_change_dims={0: [0], 1: [0], 2: [0, 1]}
+                    ),
                 ),
             ),
         ):
@@ -165,7 +169,9 @@ def test_adamw_eigenvalue_correction_on_quadratic(self) -> None:
                 (math.inf, DefaultSOAPConfig),
                 (
                     1,
-                    EigenvalueCorrectedShampooPreconditionerConfig(ignored_dims=[0, 1]),
+                    EigenvalueCorrectedShampooPreconditionerConfig(
+                        ignored_basis_change_dims={0: [0], 1: [0], 2: [0, 1]}
+                    ),
                 ),
             ),
         ):
@@ -217,7 +223,9 @@ def test_rmsprop_eigenvalue_correction_on_quadratic(self) -> None:
                 (math.inf, DefaultSOAPConfig),
                 (
                     1,
-                    EigenvalueCorrectedShampooPreconditionerConfig(ignored_dims=[0, 1]),
+                    EigenvalueCorrectedShampooPreconditionerConfig(
+                        ignored_basis_change_dims={0: [0], 1: [0], 2: [0, 1]}
+                    ),
                 ),
             ),
         ):