Add an argument to set bucket_cap_mb for PyTorch DDP (#14756)

* [trainer] Set bucket_cap_mb for DDP from arguments

* Put find_unused_parameters into kwargs

src/transformers/trainer.py

@@ -999,19 +999,23 @@ class Trainer:
         elif is_sagemaker_dp_enabled():
             model = DDP(model, device_ids=[dist.get_local_rank()], broadcast_buffers=False)
         elif self.args.local_rank != -1:
+            kwargs = {}
             if self.args.ddp_find_unused_parameters is not None:
-                find_unused_parameters = self.args.ddp_find_unused_parameters
+                kwargs["find_unused_parameters"] = self.args.ddp_find_unused_parameters
             elif isinstance(model, PreTrainedModel):
                 # find_unused_parameters breaks checkpointing as per
                 # https://github.com/huggingface/transformers/pull/4659#issuecomment-643356021
-                find_unused_parameters = not model.is_gradient_checkpointing
+                kwargs["find_unused_parameters"] = not model.is_gradient_checkpointing
             else:
-                find_unused_parameters = True
+                kwargs["find_unused_parameters"] = True
+
+            if self.args.ddp_bucket_cap_mb is not None:
+                kwargs["bucket_cap_mb"] = self.args.ddp_bucket_cap_mb
             model = nn.parallel.DistributedDataParallel(
                 model,
                 device_ids=[self.args.local_rank] if self.args._n_gpu != 0 else None,
                 output_device=self.args.local_rank if self.args._n_gpu != 0 else None,
-                find_unused_parameters=find_unused_parameters,
+                **kwargs,
             )
 
         return model

src/transformers/training_args.py

@@ -348,6 +348,9 @@ class TrainingArguments:
             When using distributed training, the value of the flag :obj:`find_unused_parameters` passed to
             :obj:`DistributedDataParallel`. Will default to :obj:`False` if gradient checkpointing is used, :obj:`True`
             otherwise.
+        ddp_bucket_cap_mb (:obj:`int`, `optional`):
+            When using distributed training, the value of the flag :obj:`bucket_cap_mb` passed to
+            :obj:`DistributedDataParallel`.
         dataloader_pin_memory (:obj:`bool`, `optional`, defaults to :obj:`True`):
             Whether you want to pin memory in data loaders or not. Will default to :obj:`True`.
         skip_memory_metrics (:obj:`bool`, `optional`, defaults to :obj:`True`):
@@ -665,6 +668,13 @@ class TrainingArguments:
             "`DistributedDataParallel`."
         },
     )
+    ddp_bucket_cap_mb: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "When using distributed training, the value of the flag `bucket_cap_mb` passed to "
+            "`DistributedDataParallel`."
+        },
+    )
     dataloader_pin_memory: bool = field(
         default=True, metadata={"help": "Whether or not to pin memory for DataLoader."}
     )