Repository: mosaicml/composer

Issue:
Composer prevents loss from converging when using DecoupledAdamW
** Environment **
- OS: Ubuntu 22.04 LTS
- GPU:
```
  *-display                 
       description: VGA compatible controller
       product: GP102 [GeForce GTX 1080 Ti]
       vendor: NVIDIA Corporation
       physical id: 0
       bus info: pci@0000:01:00.0
       version: a1
       width: 64 bits
       clock: 33MHz
       capabilities: pm msi pciexpress vga_controller bus_master cap_list rom
       configuration: driver=nvidia latency=0
       resources: irq:58 memory:fa000000-faffffff memory:c0000000-cfffffff memory:d0000000-d1ffffff ioport:e000(size=128) memory:c0000-dffff
  *-graphics
       product: EFI VGA
       physical id: 2
       logical name: /dev/fb0
       capabilities: fb
       configuration: depth=32 resolution=1024,768
```
- Cuda: 11.7
- Composer: 0.8.2


I'm trying to implement a paper in Pytorch, using composer for my training loop.
See here: https://github.com/vedantroy/improved-ddpm-pytorch/commit/24909c064953bf54bf3a50c64923fdfef714f3cc

Composer is breaking stuff in at least 2 different ways:
1. If I replace this line with composer's `DecoupledAdamW` then the loss goes from 1.0 => 0.6 but not much further down
   https://github.com/vedantroy/improved-ddpm-pytorch/blob/24909c064953bf54bf3a50c64923fdfef714f3cc/openai_iddpm.py#L103
2. ~If I use the `OpenAIIDDPM` model instead of the `manual_train` function, the loss does not go down from around 1.0 at all
  https://github.com/vedantroy/improved-ddpm-pytorch/blob/24909c064953bf54bf3a50c64923fdfef714f3cc/openai_iddpm.py#L72~
 - I think only the 1st thing was a bug b/c once I switched to torch's `AdamW` and removed the warmup, the learning rate progressed as normal

I've uploaded the data here, so you can reproduce it: https://drive.google.com/drive/folders/1zK-3SkGOETGFajIBSXbEPqPSRdwqajUq?usp=sharing

It's just the CelebHQ dataset in 64x64 resolution, stored as parquet. Any image dataset should work and you can probably swap out my torch data based dataloader for something more standard.


Discussion:
Hey, thank you for raising this issue!

I think this is a bug on our side where we copied over torch's `AdamW` defaults, in particular the `weight_decay=1e-2`.

Since our decoupled optimizers [implement the original paper](https://arxiv.org/abs/1711.05101) and completely separate `lr` and `weight_decay`, the default of `1e-2` is much too high, and would decay your weights by 1/100th every batch! This is likely what was causing the poor convergence. I think our internal benchmarks that used `DecoupledAdamW` always set `weight_decay` explicitly which is why we missed this bug :( 

I will update the default value to `weight_decay=1e-5`, which should be close if not identical to the behavior of torch's `AdamW` which has defaults `lr=1e-3` and `weight_decay=1e-2`.



Resolving patch:
diff --git a/composer/optim/decoupled_weight_decay.py b/composer/optim/decoupled_weight_decay.py
--- a/composer/optim/decoupled_weight_decay.py
+++ b/composer/optim/decoupled_weight_decay.py
@@ -11,7 +11,7 @@
 
 import logging
 import math
-from typing import Iterable, List, Tuple
+from typing import Iterable, List, Tuple, Union
 
 import torch
 from torch.optim import SGD, AdamW
@@ -25,26 +25,50 @@
 class DecoupledSGDW(SGD):
     """SGD optimizer with the weight decay term decoupled from the learning rate.
 
+    NOTE: Since `weight_decay` is no longer scaled by `lr`, you will likely want to use much smaller values
+    for `weight_decay` than you would if using `torch.optim.SGD`. In this optimizer, the value `weight_decay` translates exactly to:
+    'On every optimizer update, every weight element will be multiplied by `(1.0 - weight_decay_t)`'.
+    The term `weight_decay_t` will follow the same schedule as `lr_t` but crucially will not be scaled by `lr`.
+
     Argument defaults are copied from :class:`torch.optim.SGD`.
 
-    The standard `SGD <https://pytorch.org/docs/stable/generated/torch.optim.SGD.html?highlight=sgd#torch.optim.SGD>`_
+    Why use this optimizer? The standard `SGD <https://pytorch.org/docs/stable/generated/torch.optim.SGD.html?highlight=sgd#torch.optim.SGD>`_
     optimizer couples the weight decay term with the gradient calculation. This ties the optimal value
     of :attr:`weight_decay` to :attr:`lr` and can also hurt generalization in practice. For more details
     on why decoupling might be desirable, see `Decoupled Weight Decay Regularization <https://arxiv.org/abs/1711.05101>`_.
 
     Args:
-        params (list): List of parameters to optimize or dicts defining parameter groups.
-        lr (float, optional): Learning rate.
+        params (iterable): Iterable of parameters to optimize or dicts defining parameter groups.
+        lr (float): Learning rate.
         momentum (int, optional): Momentum factor. Default: ``0``.
         dampening (int, optional): Dampening factor applied to the momentum. Default: ``0``.
         weight_decay (int, optional): Decoupled weight decay factor. Default: ``0``.
         nesterov (bool, optional): Enables Nesterov momentum updates. Default: ``False``.
     """
 
+    def __init__(self,
+                 params: Union[Iterable[torch.Tensor], Iterable[dict]],
+                 lr: float = required,
+                 momentum: float = 0,
+                 dampening: float = 0,
+                 weight_decay: float = 0,
+                 nesterov: bool = False):
+        if weight_decay >= 1e-3:
+            log.warning(
+                f'You are using a high value of `weight_decay={weight_decay}` for the `DecoupledSGDW` optimizer. Are you sure you want to do this? '
+                f'Your model\'s weights will be multiplied by {1.0 - weight_decay} on every step!')
+        super().__init__(params=params,
+                         lr=lr,
+                         momentum=momentum,
+                         dampening=dampening,
+                         weight_decay=weight_decay,
+                         nesterov=nesterov)
+        for group in self.param_groups:
+            group['initial_lr'] = group['lr']
+
     @staticmethod
-    def sgdw(params: Iterable[torch.nn.parameter.Parameter], d_p_list: List[torch.Tensor],
-             momentum_buffer_list: List[torch.Tensor], *, weight_decay: float, momentum: float, lr: float,
-             initial_lr: float, dampening: float, nesterov: bool):
+    def sgdw(params: List[torch.Tensor], d_p_list: List[torch.Tensor], momentum_buffer_list: List[torch.Tensor], *,
+             weight_decay: float, momentum: float, lr: float, initial_lr: float, dampening: float, nesterov: bool):
         r"""Functional API that performs SGDW algorithm computation.
 
         Args:
@@ -77,22 +101,11 @@ def sgdw(params: Iterable[torch.nn.parameter.Parameter], d_p_list: List[torch.Te
                     d_p = buf
 
             if weight_decay != 0:
-                decay_factor = lr / initial_lr
+                decay_factor = (lr / initial_lr) if initial_lr else 1.0
                 param.mul_(1 - decay_factor * weight_decay)
 
             param.add_(d_p, alpha=-lr)
 
-    def __init__(self,
-                 params: List[torch.Tensor],
-                 lr: float = required,
-                 momentum: float = 0,
-                 dampening: float = 0,
-                 weight_decay: float = 0,
-                 nesterov: bool = False):
-        super().__init__(params, lr, momentum, dampening, weight_decay, nesterov)
-        for group in self.param_groups:
-            group['initial_lr'] = group['lr']
-
     @torch.no_grad()
     def step(self, closure=None):
         """Performs a single optimization step.
@@ -149,23 +162,45 @@ def step(self, closure=None):
 class DecoupledAdamW(AdamW):
     """Adam optimizer with the weight decay term decoupled from the learning rate.
 
-    Argument defaults are copied from :class:`torch.optim.AdamW`.
+    NOTE: Since `weight_decay` is no longer scaled by `lr`, you will likely want to use much smaller values
+    for `weight_decay` than you would if using `torch.optim.Adam` or `torch.optim.AdamW`. In this optimizer, the value `weight_decay` translates exactly to:
+    'On every optimizer update, every weight element will be multiplied by `(1.0 - weight_decay_t)`'.
+    The term `weight_decay_t` will follow the same schedule as `lr_t` but crucially will not be scaled by `lr`.
+
+    Argument defaults are similar to :class:`torch.optim.AdamW` but we make two changes:
+    * The default for ``weight_decay`` is changed from ``1e-2`` -> ``1e-5`` because in `DecoupledAdamW`, the weight decay is decoupled and no longer scaled by the `lr=1e-3`.
+    * The default for ``betas`` is changed from ``(0.9, 0.999)`` to ``(0.9, 0.95)`` to reflect community best-practices for the beta2 hyperparameter.
 
-    The standard `AdamW <https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html#torch.optim.AdamW>`_
+    Why use this optimizer? The standard `AdamW <https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html#torch.optim.AdamW>`_
     optimizer explicitly couples the weight decay term with the learning rate. This ties the
     optimal value of :attr:`weight_decay` to :attr:`lr` and can also hurt generalization in practice. For more details on
     why decoupling might be desirable, see `Decoupled Weight Decay Regularization <https://arxiv.org/abs/1711.05101>`_.
 
     Args:
-        params (list): List of parameters to update.
+        params (iterable): Iterable of parameters to optimize or dicts defining parameter groups.
         lr (float, optional): Learning rate. Default: ``1e-3``.
         betas (tuple, optional): Coefficients used for computing running averages of gradient and its square
-                                 Default: ``(0.9, 0.999)``.
+                                 Default: ``(0.9, 0.95)``.
         eps (float, optional): Term added to the denominator to improve numerical stability. Default: ``1e-8``.
-        weight_decay (float, optional): Decoupled weight decay factor. Default: ``1e-2``.
+        weight_decay (float, optional): Decoupled weight decay factor. Default: ``1e-5``.
         amsgrad (bool, optional): Enables the amsgrad variant of Adam. Default: ``False``.
     """
 
+    def __init__(self,
+                 params: Union[Iterable[torch.Tensor], Iterable[dict]],
+                 lr: float = 1e-3,
+                 betas: Tuple[float, float] = (0.9, 0.95),
+                 eps: float = 1e-8,
+                 weight_decay: float = 1e-5,
+                 amsgrad: bool = False):
+        if weight_decay >= 1e-3:
+            log.warning(
+                f'You are using a high value of `weight_decay={weight_decay}` for the `DecoupledAdamW` optimizer. Are you sure you want to do this? '
+                f'Your model\'s weights will be multiplied by {1.0 - weight_decay} on every step!')
+        super().__init__(params=params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, amsgrad=amsgrad)
+        for group in self.param_groups:
+            group['initial_lr'] = group['lr']
+
     @staticmethod
     def adamw(params: List[torch.Tensor], grads: List[torch.Tensor], exp_avgs: List[torch.Tensor],
               exp_avg_sqs: List[torch.Tensor], max_exp_avg_sqs: List[torch.Tensor], state_steps: List[int], *,
@@ -174,12 +209,12 @@ def adamw(params: List[torch.Tensor], grads: List[torch.Tensor], exp_avgs: List[
         r"""Functional API that performs AdamW algorithm computation with decoupled weight decay.
 
         Args:
-            params (List[torch.Tensor]): List of parameters to update.
-            grads (List[torch.Tensor]): List of parameter gradients.
-            exp_avgs (List[torch.Tensor]): List of average gradients.
-            exp_avg_sqs (List[torch.Tensor]): List of average squared gradients.
-            max_exp_avg_sqs (List[torch.Tensor]): List of max average squared gradients for amsgrad updates.
-            state_steps (Iterable[int]): List of steps taken for all parameters.
+            params (list): List of parameters to update.
+            grads (list): List of parameter gradients.
+            exp_avgs (list): List of average gradients.
+            exp_avg_sqs (list): List of average squared gradients.
+            max_exp_avg_sqs (list): List of max average squared gradients for amsgrad updates.
+            state_steps (list): List of steps taken for all parameters.
             amsgrad (bool): Enables amsgrad variant of Adam.
             beta1 (float): Coefficient for computing the moving average of gradient values.
             beta2 (float): Coefficient for computing the moving average of squared gradient values.
@@ -196,7 +231,7 @@ def adamw(params: List[torch.Tensor], grads: List[torch.Tensor], exp_avgs: List[
 
             # Perform stepweight decay
             if weight_decay != 0:
-                decay_factor = lr / initial_lr
+                decay_factor = (lr / initial_lr) if initial_lr else 1.0
                 param.mul_(1 - decay_factor * weight_decay)
 
             bias_correction1 = 1 - beta1**step
@@ -217,17 +252,6 @@ def adamw(params: List[torch.Tensor], grads: List[torch.Tensor], exp_avgs: List[
 
             param.addcdiv_(exp_avg, denom, value=-step_size)
 
-    def __init__(self,
-                 params: List[torch.Tensor],
-                 lr: float = 1e-3,
-                 betas: Tuple[float, float] = (0.9, 0.999),
-                 eps: float = 1e-8,
-                 weight_decay: float = 1e-2,
-                 amsgrad: bool = False):
-        super().__init__(params, lr, betas, eps, weight_decay, amsgrad)
-        for group in self.param_groups:
-            group['initial_lr'] = group['lr']
-
     @torch.no_grad()
     def step(self, closure=None):
         """Performs a single optimization step.


Test patch: