Merge pull request #376 from kozistr/feature/optimizers

[Feature] Implement Fira, RACS and Alice optimizers

Author: kozistr

Makefile

@@ -1,4 +1,4 @@
-.PHONY: init format test check requirements docs
+.PHONY: init format test check requirements visualize docs
 
 init:
 	python -m pip install -q -U poetry isort black ruff pytest pytest-cov
@@ -8,16 +8,19 @@ format:
 	isort --profile black -l 119 pytorch_optimizer examples tests hubconf.py
 	black -S -l 119 pytorch_optimizer examples tests hubconf.py
 
-test:
-	python -m pytest -p no:pastebin -p no:nose -p no:doctest -sv -vv --cov=pytorch_optimizer --cov-report=xml ./tests
-
 check:
 	black -S -l 119 --check pytorch_optimizer examples tests hubconf.py
 	ruff check pytorch_optimizer examples tests hubconf.py
 
+test:
+	python -m pytest -p no:pastebin -p no:nose -p no:doctest -sv -vv --cov=pytorch_optimizer --cov-report=xml ./tests
+
 requirements:
 	poetry export -f requirements.txt --output requirements.txt --without-hashes
 	poetry export -f requirements.txt --output requirements-dev.txt --without-hashes --with dev
 
+visualize:
+	python -m examples.visualize_optimizers
+
 docs:
 	mkdocs serve

README.md

@@ -0,0 +1,12 @@
+## Change Log
+
+### Feature
+
+* Implement `Fira` optimizer. (#376)
+    * [Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint?](https://arxiv.org/abs/2410.01623) 
+* Implement `RACS` and `Alice optimizer. (#376)
+    * [Towards Efficient Optimizer Design for LLM via Structured Fisher Approximation with a Low-Rank Extension](https://arxiv.org/abs/2502.07752)
+
+### Fix
+
+* Fix shape mismatch issues in the Galore projection for `reverse_std`, `right` and `full` projection types. (#376)

docs/changelogs/v3.5.2.md

@@ -188,6 +188,10 @@
     :docstring:
     :members:
 
+::: pytorch_optimizer.Fira
+    :docstring:
+    :members:
+
 ::: pytorch_optimizer.FOCUS
     :docstring:
     :members:

docs/index.md

@@ -174,6 +174,10 @@
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rastrigin_FAdam.png)
 
+### Fira
+
+![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rastrigin_Fira.png)
+
 ### FOCUS
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rastrigin_FOCUS.png)
@@ -556,6 +560,10 @@
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rosenbrock_FAdam.png)
 
+### Fira
+
+![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rosenbrock_Fira.png)
+
 ### FOCUS
 
 ![image](https://raw.githubusercontent.com/kozistr/pytorch_optimizer/main/docs/visualizations/rosenbrock_FOCUS.png)

docs/optimizer.md

@@ -10,18 +10,18 @@ homepage = "https://github.com/kozistr/pytorch_optimizer"
 repository = "https://github.com/kozistr/pytorch_optimizer"
 documentation = "https://pytorch-optimizers.readthedocs.io/en/latest"
 keywords = [
-    "pytorch", "deep-learning", "optimizer", "lr scheduler", "A2Grad", "ASGD", "AccSGD", "AdaBelief", "AdaBound",
-    "AdaDelta", "AdaFactor", "AdaGC", "AdaMax", "AdamG", "AdaMod", "AdaNorm", "AdaPNM", "AdaSmooth", "AdEMAMix",
-    "Simplified-AdEMAMix", "ADOPT", "AdaHessian", "Adai", "Adalite", "AdaLomo", "AdamMini", "AdamP", "AdamS", "Adan",
-    "AggMo", "Aida", "AliG", "Amos", "Apollo", "APOLLO", "AvaGrad", "bSAM", "CAME", "DAdaptAdaGrad", "DAdaptAdam",
-    "DAdaptAdan", "DAdaptSGD", "DAdaptLion", "DeMo", "DiffGrad", "EXAdam", "FAdam", "FOCUS", "Fromage", "FTRL",
-    "GaLore", "Grams", "Gravity", "GrokFast", "GSAM", "Kate", "Lamb", "LaProp", "LARS", "Lion", "LOMO", "Lookahead",
-    "MADGRAD", "MARS", "MSVAG", "Muno", "Nero", "NovoGrad", "OrthoGrad", "PAdam", "PCGrad", "PID", "PNM", "Prodigy",
-    "PSGD", "QHAdam", "QHM", "RAdam", "Ranger", "Ranger21", "RotoGrad", "SAM", "GCSAM", "LookSAM", "ScheduleFreeSGD",
-    "ScheduleFreeAdamW", "ScheduleFreeRAdam", "SCION", "SGDP", "Shampoo", "ScalableShampoo", "SGDW", "SignSGD", "SM3",
-    "SOAP", "SopihaH", "SPAM", "StableSPAM", "SRMM", "StableAdamW", "SWATS", "TAM", "Tiger", "TRAC", "WSAM", "Yogi",
-    "BCE", "BCEFocal", "Focal", "FocalCosine", "SoftF1", "Dice", "LDAM", "Jaccard", "Bi-Tempered", "Tversky",
-    "FocalTversky", "LovaszHinge", "bitsandbytes", "WSD", "QGaLore",
+    "pytorch", "deep-learning", "optimizer", "lr scheduler", "A2Grad", "Alice", "ASGD", "AccSGD", "AdaBelief",
+    "AdaBound", "AdaDelta", "AdaFactor", "AdaGC", "AdaMax", "AdamG", "AdaMod", "AdaNorm", "AdaPNM", "AdaSmooth",
+    "AdEMAMix", "Simplified-AdEMAMix", "ADOPT", "AdaHessian", "Adai", "Adalite", "AdaLomo", "AdamMini", "AdamP",
+    "AdamS", "Adan", "AggMo", "Aida", "AliG", "Amos", "Apollo", "APOLLO", "AvaGrad", "bSAM", "CAME", "DAdaptAdaGrad",
+    "DAdaptAdam", "DAdaptAdan", "DAdaptSGD", "DAdaptLion", "DeMo", "DiffGrad", "EXAdam", "FAdam", "Fira", "FOCUS",
+    "Fromage", "FTRL", "GaLore", "Grams", "Gravity", "GrokFast", "GSAM", "Kate", "Lamb", "LaProp", "LARS", "Lion",
+    "LOMO", "Lookahead", "MADGRAD", "MARS", "MSVAG", "Muno", "Nero", "NovoGrad", "OrthoGrad", "PAdam", "PCGrad", "PID",
+    "PNM", "Prodigy", "PSGD", "QHAdam", "QHM", "RACS", "RAdam", "Ranger", "Ranger21", "RotoGrad", "SAM", "GCSAM",
+    "LookSAM", "ScheduleFreeSGD", "ScheduleFreeAdamW", "ScheduleFreeRAdam", "SCION", "SGDP", "Shampoo",
+    "ScalableShampoo", "SGDW", "SignSGD", "SM3", "SOAP", "SopihaH", "SPAM", "StableSPAM", "SRMM", "StableAdamW",
+    "SWATS", "TAM", "Tiger", "TRAC", "WSAM", "Yogi", "BCE", "BCEFocal", "Focal", "FocalCosine", "SoftF1", "Dice",
+    "LDAM", "Jaccard", "Bi-Tempered", "Tversky", "FocalTversky", "LovaszHinge", "bitsandbytes", "WSD", "QGaLore",
 ]
 classifiers = [
     "License :: OSI Approved :: Apache Software License",

docs/visualization.md

@@ -57,6 +57,7 @@
     PID,
     PNM,
     QHM,
+    RACS,
     SAM,
     SCION,
     SGDP,
@@ -96,6 +97,7 @@
     AdEMAMix,
     AggMo,
     Aida,
+    Alice,
     AliG,
     Amos,
     ApolloDQN,
@@ -110,6 +112,7 @@
     DynamicLossScaler,
     EXAdam,
     FAdam,
+    Fira,
     Fromage,
     GaLore,
     Grams,