update: GrokFast optimizer

Author: kozistr

pytorch_optimizer/__init__.py

@@ -60,6 +60,7 @@
 from pytorch_optimizer.optimizer.galore import GaLore, GaLoreProjector
 from pytorch_optimizer.optimizer.gc import centralize_gradient
 from pytorch_optimizer.optimizer.gravity import Gravity
+from pytorch_optimizer.optimizer.grokfast import GrokFastAdamW, gradfilter_ema, gradfilter_ma
 from pytorch_optimizer.optimizer.lamb import Lamb
 from pytorch_optimizer.optimizer.lars import LARS
 from pytorch_optimizer.optimizer.lion import Lion
@@ -192,6 +193,7 @@
     ScheduleFreeSGD,
     ScheduleFreeAdamW,
     FAdam,
+    GrokFastAdamW,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 

tests/constants.py

@@ -49,6 +49,7 @@
     Fromage,
     GaLore,
     Gravity,
+    GrokFastAdamW,
     Lamb,
     Lion,
     Nero,
@@ -129,6 +130,7 @@
     'bsam',
     'schedulefreeadamw',
     'fadam',
+    'grokfastadamw',
 ]
 
 VALID_LR_SCHEDULER_NAMES: List[str] = [
@@ -448,6 +450,7 @@
     (ScheduleFreeSGD, {'lr': 1e0, 'weight_decay': 1e-3}, 5),
     (ScheduleFreeAdamW, {'lr': 1e0, 'weight_decay': 1e-3}, 5),
     (FAdam, {'lr': 1e0, 'weight_decay': 1e-3}, 5),
+    (GrokFastAdamW, {'lr': 1e0, 'weight_decay': 1e-3}, 10),
 ]
 ADANORM_SUPPORTED_OPTIMIZERS: List[Tuple[Any, Dict[str, Union[float, bool, int]], int]] = [
     (AdaBelief, {'lr': 5e-1, 'weight_decay': 1e-3, 'adanorm': True}, 10),

tests/test_load_modules.py

@@ -38,7 +38,7 @@ def test_load_lr_scheduler_invalid(invalid_lr_scheduler_names):
 
 
 def test_get_supported_optimizers():
-    assert len(get_supported_optimizers()) == 67
+    assert len(get_supported_optimizers()) == 68
 
 
 def test_get_supported_lr_schedulers():