Merge pull request #267 from kozistr/fix/prodigy-optimizer

[Fix] device mismatch problems

Author: kozistr

docs/changelogs/v3.1.1.md

@@ -10,4 +10,5 @@
 
 ### Bug
 
-* Fix to handle the optimizers that only take the `model` instead of the parameters in `create_optimizer()`. (#263)
+* Handle the optimizers that only take the `model` instead of the parameters in `create_optimizer()`. (#263)
+* Move the variable to the same device with the parameter. (#266, #267)

pytorch_optimizer/optimizer/adamg.py

@@ -27,7 +27,7 @@ def __init__(
         lr: float = 1e-3,
         betas: BETAS = (0.95, 0.999, 0.95),
         p: float = 0.5,
-        q: float = 0.25,
+        q: float = 0.24,
         weight_decay: float = 0.0,
         weight_decouple: bool = False,
         fixed_decay: bool = False,

pytorch_optimizer/optimizer/adan.py

@@ -83,7 +83,7 @@ def get_global_gradient_norm(self) -> Union[torch.Tensor, float]:
         if self.defaults['max_grad_norm'] == 0.0:
             return 1.0
 
-        global_grad_norm = get_global_gradient_norm(self.param_groups, self.param_groups[0]['params'][0].device)
+        global_grad_norm = get_global_gradient_norm(self.param_groups)
         global_grad_norm.sqrt_().add_(self.defaults['eps'])
 
         return torch.clamp(self.defaults['max_grad_norm'] / global_grad_norm, max=1.0)

pytorch_optimizer/optimizer/alig.py

@@ -52,7 +52,7 @@ def reset(self):
     @torch.no_grad()
     def compute_step_size(self, loss: float) -> float:
         r"""Compute step_size."""
-        global_grad_norm = get_global_gradient_norm(self.param_groups, torch.device('cpu'))
+        global_grad_norm = get_global_gradient_norm(self.param_groups)
         global_grad_norm.add_(1e-6)
 
         return loss / global_grad_norm.item()

pytorch_optimizer/optimizer/dadapt.py

@@ -473,7 +473,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
         numerator_weighted = group['numerator_weighted']
 
         if group['step'] == 0:
-            group['g0_norm'] = get_global_gradient_norm(self.param_groups, device).sqrt_().item()
+            group['g0_norm'] = get_global_gradient_norm(self.param_groups).sqrt_().item()
         g0_norm = group['g0_norm']
 
         if g0_norm == 0:

pytorch_optimizer/optimizer/lamb.py

@@ -103,7 +103,7 @@ def get_global_gradient_norm(self) -> Union[torch.Tensor, float]:
         if self.defaults['max_grad_norm'] == 0.0:
             return 1.0
 
-        global_grad_norm = get_global_gradient_norm(self.param_groups, self.param_groups[0]['params'][0].device)
+        global_grad_norm = get_global_gradient_norm(self.param_groups)
         global_grad_norm.sqrt_().add_(self.defaults['eps'])
 
         return torch.clamp(self.defaults['max_grad_norm'] / global_grad_norm, max=1.0)

pytorch_optimizer/optimizer/prodigy.py

@@ -110,6 +110,8 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
         if 'd_numerator' not in group:
             group['d_numerator'] = torch.tensor([0.0], device=device)
+        elif group['d_numerator'].device != device:
+            group['d_numerator'] = group['d_numerator'].to(device)
 
         d_numerator = group['d_numerator']
         d_numerator.mul_(beta3)

pytorch_optimizer/optimizer/utils.py

@@ -272,9 +272,9 @@ def l2_projection(parameters: PARAMETERS, max_norm: float = 1e2):
 
 
 @torch.no_grad()
-def get_global_gradient_norm(param_groups: List[Dict], device: torch.device) -> torch.Tensor:
+def get_global_gradient_norm(param_groups: List[Dict]) -> torch.Tensor:
     r"""Get global gradient norm."""
-    global_grad_norm = torch.zeros(1, dtype=torch.float32, device=device)
+    global_grad_norm = torch.zeros(1, dtype=torch.float32, device=param_groups[0]['params'][0].device)
 
     for group in param_groups:
         for p in group['params']:

tests/constants.py

@@ -428,8 +428,8 @@
     (Prodigy, {'lr': 1e0, 'beta3': 0.999, 'weight_decay': 1e-3, 'safeguard_warmup': True}, 15),
     (PAdam, {'lr': 1e0, 'weight_decay': 1e-3}, 5),
     (Tiger, {'lr': 1e0, 'weight_decay': 1e-3}, 5),
-    (CAME, {'lr': 7.5e-1, 'weight_decay': 1e-3}, 75),
-    (CAME, {'lr': 7.5e-1, 'weight_decay': 1e-3, 'ams_bound': True}, 75),
+    (CAME, {'lr': 7.5e-1, 'weight_decay': 1e-3}, 70),
+    (CAME, {'lr': 7.5e-1, 'weight_decay': 1e-3, 'ams_bound': True}, 70),
     (Aida, {'lr': 1e0, 'weight_decay': 1e-3, 'ams_bound': True}, 5),
     (
         GaLore,