Add option to use fused adamw optimizer

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

adamw recently added an option to use a fused optimizer. It may give better performance than foreach argument. However, we cannot enable it by default, since it requires all parameters to be in CUDA and maybe some other restrictions. So, enable it on a per project basis.

On DALLE2, it results about 23ms faster.

Reviewed By: newstzpz

Differential Revision: D43027327

fbshipit-source-id: 82c6855116094e86386ad2edeea3a74f9e555174

d2go/optimizer/build.py

@@ -280,7 +280,8 @@ def adamw(cfg, model: torch.nn.Module) -> torch.optim.Optimizer:
         lr=cfg.SOLVER.BASE_LR,
         betas=cfg.SOLVER.BETAS,
         eps=cfg.SOLVER.EPS,
-        foreach=True,
+        foreach=True if cfg.SOLVER.FUSED is False else False,
+        fused=True if cfg.SOLVER.FUSED else False,
     )
 
 

d2go/runner/config_defaults.py

@@ -80,6 +80,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
+    _C.SOLVER.FUSED = False
 
     # RECOMPUTE_BOXES for LSJ Training
     _C.INPUT.RECOMPUTE_BOXES = False