Adding shampoo training

run_mlm_flax.py

@@ -60,6 +60,7 @@ from transformers import (
 )
 from transformers.file_utils import get_full_repo_name
 
+from distributed_shampoo import distributed_shampoo, GraftingType
 
 MODEL_CONFIG_CLASSES = list(FLAX_MODEL_FOR_MASKED_LM_MAPPING.keys())
 MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
@@ -88,6 +89,10 @@ class TrainingArguments:
         default=8, metadata={"help": "Batch size per GPU/TPU core/CPU for evaluation."}
     )
     learning_rate: float = field(default=5e-5, metadata={"help": "The initial learning rate for AdamW."})
+    distributed_shampoo: bool = field(
+        default=False,
+        metadata={"help": "Use Distributed Shampoo optimizer instead of AdamW."},
+    )
     weight_decay: float = field(default=0.0, metadata={"help": "Weight decay for AdamW if we apply some."})
     adam_beta1: float = field(default=0.9, metadata={"help": "Beta1 for AdamW optimizer"})
     adam_beta2: float = field(default=0.999, metadata={"help": "Beta2 for AdamW optimizer"})
@@ -629,6 +634,32 @@ def main():
         optimizer = optax.adafactor(
             learning_rate=linear_decay_lr_schedule_fn,
         )
+    elif training_args.distributed_shampoo:
+        # parameters from https://github.com/tensorflow/lingvo/blob/03ee9d7cd50764b0424c7c863733c91fc0b053ec/lingvo/jax/optimizers.py#L729
+        # Notes:
+        # - mask for weight decay is not implemented but we don't use it anyway
+        optimizer = distributed_shampoo(
+            linear_decay_lr_schedule_fn,
+            block_size=1024,  # recommended default for large LM is 1536
+            beta1=training_args.adam_beta1,  # 0.9,
+            beta2=training_args.adam_beta2,  # 0.999,
+            diagonal_epsilon=training_args.adam_epsilon,  # 1e-10,
+            matrix_epsilon=1e-8,
+            weight_decay=training_args.weight_decay,  # 0.0,
+            start_preconditioning_step=1001,
+            preconditioning_compute_steps=10,
+            statistics_compute_steps=1,
+            best_effort_shape_interpretation=True,
+            graft_type=GraftingType.RMSPROP_NORMALIZED,
+            nesterov=False,
+            exponent_override=0,
+            batch_axis_name="batch",
+            inverse_failure_threshold=0.1,
+            moving_average_for_momentum=True,
+            skip_preconditioning_dim_size_gt=4096,
+            clip_by_scaled_gradient_norm=None,
+            precision=jax.lax.Precision.HIGHEST,
+        )
     else:
         optimizer = optax.adamw(
             learning_rate=linear_decay_lr_schedule_fn,