wip shampoo optim support

2024-09-18 08:10:52 -07:00
parent 7b9f669a3a
commit eb3eab3450
3 changed files with 91 additions and 2 deletions
--- a/requirements.txt
+++ b/requirements.txt
@@ -35,6 +35,7 @@ python-dotenv==1.0.1
 autoawq>=0.2.5
 triton>=2.3.0
 liger-kernel==0.2.1
+distributed_shampoo @ git+https://github.com/facebookresearch/optimizers.git@main

 mamba-ssm==1.2.0.post1