upgrade to latest torch shampoo comit

PrimeIntellect-ai · Jan 6, 2025 · ea23984 · ea23984
1 parent a64ca4b
commit ea23984
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/src/zeroband/optimizers/__init__.py b/src/zeroband/optimizers/__init__.py
@@ -3,7 +3,7 @@
 import torch
 from zeroband.optimizers.muon import Muon, AdamConfig, MuonConfig
 from distributed_shampoo import (
-    EighEigenvalueCorrectionConfig,
+    DefaultSOAPConfig,
     DistributedShampoo,
     FullyShardShampooConfig,
     ShampooPT2CompileConfig,
@@ -52,9 +52,7 @@ def get_optimizer(params: list[torch.nn.Parameter], config: OptimizersConfig) ->
             max_preconditioner_dim=config.max_preconditioner_dim,
             precondition_frequency=config.precondition_frequency,
             use_decoupled_weight_decay=True,
-            # This can also be set to `QREigenvalueCorrectionConfig` which is less expensive
-            # and might therefore allow for a smaller `precondition_frequency`.
-            preconditioner_computation_config=EighEigenvalueCorrectionConfig(),
+            preconditioner_config=DefaultSOAPConfig,
             distributed_config=FullyShardShampooConfig(),
             shampoo_pt2_compile_config=ShampooPT2CompileConfig(enable_shampoo_pt2_dynamic_shape=False),
         )

diff --git a/uv.lock b/uv.lock