Add a config for customize eps of the the AdamW optimizer

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/424

The diffusion model need to customized eps for the AdamW optimizer. For example the prior of Dalle2 trained with adam_eps=1.0e-06 while its decoder trained with adam_eps=1.0e-08

Add the config for AdamW's eps to d2go pipeline with its default vaule same as the official doc: https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html

Reviewed By: newstzpz

Differential Revision: D41363142

fbshipit-source-id: 72f9a4084e229312807aad28b7aba8fec9116013

d2go/optimizer/build.py

@@ -264,7 +264,7 @@ def adam(cfg, model: torch.nn.Module) -> torch.optim.Optimizer:
     params = get_optimizer_param_groups(model, cfg)
 
     return maybe_add_gradient_clipping(cfg, torch.optim.Adam)(
-        params=params, lr=cfg.SOLVER.BASE_LR, betas=cfg.SOLVER.BETAS
+        params=params, lr=cfg.SOLVER.BASE_LR, betas=cfg.SOLVER.BETAS, eps=cfg.SOLVER.EPS
     )
 
 
@@ -276,7 +276,7 @@ def adamw(cfg, model: torch.nn.Module) -> torch.optim.Optimizer:
     params = get_optimizer_param_groups(model, cfg)
 
     return maybe_add_gradient_clipping(cfg, torch.optim.AdamW)(
-        params=params, lr=cfg.SOLVER.BASE_LR, betas=cfg.SOLVER.BETAS
+        params=params, lr=cfg.SOLVER.BASE_LR, betas=cfg.SOLVER.BETAS, eps=cfg.SOLVER.EPS
     )
 
 
@@ -307,7 +307,7 @@ def adamw_mt(cfg, model: torch.nn.Module) -> torch.optim.Optimizer:
     """
     params = get_optimizer_param_groups(model, cfg)
     return maybe_add_gradient_clipping(cfg, torch.optim._multi_tensor.AdamW)(
-        params=params, lr=cfg.SOLVER.BASE_LR
+        params=params, lr=cfg.SOLVER.BASE_LR, eps=cfg.SOLVER.EPS
     )
 
 

d2go/runner/config_defaults.py

@@ -78,6 +78,7 @@ def _add_detectron2go_runner_default_cfg(_C: CN) -> None:
 
     # Betas are used in the AdamW optimizer
     _C.SOLVER.BETAS = (0.9, 0.999)
+    _C.SOLVER.EPS = 1e-08
 
     # RECOMPUTE_BOXES for LSJ Training
     _C.INPUT.RECOMPUTE_BOXES = False