Set gradient as bucket view

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

Add a config variable: DDP_GRADIENT_AS_BUCKET_VIEW. Pass it to DDP. This variable reduces the memory consumption of the model.

Reviewed By: tglik

Differential Revision: D44273339

fbshipit-source-id: 272e2ffbea89532a55df0ebdb3bd49f0df7d78a5

d2go/runner/config_defaults.py

@@ -72,6 +72,8 @@ def _add_detectron2go_runner_default_cfg(_C: CN) -> None:
     _C.MODEL.DDP_FIND_UNUSED_PARAMETERS = False
     # Set FP16 gradient compression for DistributedDataParallel.
     _C.MODEL.DDP_FP16_GRAD_COMPRESS = False
+    # Specify the gradients as views
+    _C.MODEL.DDP_GRADIENT_AS_BUCKET_VIEW = False
 
     # Set default optimizer
     _C.SOLVER.OPTIMIZER = "sgd"

tools/train_net.py

@@ -93,6 +93,7 @@ def main(
                 else [comm.get_local_rank()],
                 broadcast_buffers=False,
                 find_unused_parameters=cfg.MODEL.DDP_FIND_UNUSED_PARAMETERS,
+                gradient_as_bucket_view=cfg.MODEL.DDP_GRADIENT_AS_BUCKET_VIEW,
             )
 
         logger.info("Starting train..")