From 73147fd0ad76a88adc83286c9019464a898fc696 Mon Sep 17 00:00:00 2001 From: Priya Kasimbeg Date: Sat, 10 Feb 2024 00:21:28 +0000 Subject: [PATCH] turn on preconditioning --- .../paper_baselines/shampoo/jax/distributed_shampoo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/reference_algorithms/paper_baselines/shampoo/jax/distributed_shampoo.py b/reference_algorithms/paper_baselines/shampoo/jax/distributed_shampoo.py index 725529cae..1336f3c80 100644 --- a/reference_algorithms/paper_baselines/shampoo/jax/distributed_shampoo.py +++ b/reference_algorithms/paper_baselines/shampoo/jax/distributed_shampoo.py @@ -1318,7 +1318,7 @@ def distributed_shampoo( ### inverse_failure_threshold=0.1, moving_average_for_momentum=True, - skip_preconditioning_dim_size_gt=0, + skip_preconditioning_dim_size_gt=4096, clip_by_scaled_gradient_norm=None, precision=lax.Precision.HIGHEST, tensordot_precision=None,