From cf0af9a31beb84e8feec77af51f72d063ba905aa Mon Sep 17 00:00:00 2001 From: heya5 Date: Mon, 20 Mar 2023 21:17:34 +0800 Subject: [PATCH] [Trainer] Add optional communication backends for torch.distributed when using GPU (#22247) Update training_args.py --- src/transformers/training_args.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/transformers/training_args.py b/src/transformers/training_args.py index ea91ddfe81..4c13ad0af4 100644 --- a/src/transformers/training_args.py +++ b/src/transformers/training_args.py @@ -1641,7 +1641,10 @@ class TrainingArguments: # Here, we'll use torch.distributed. # Initializes the distributed backend which will take care of synchronizing nodes/GPUs if not torch.distributed.is_initialized(): - torch.distributed.init_process_group(backend="nccl", timeout=self.ddp_timeout_delta) + if self.xpu_backend and self.xpu_backend in ("mpi", "gloo"): + torch.distributed.init_process_group(backend=self.xpu_backend, timeout=self.ddp_timeout_delta) + else: + torch.distributed.init_process_group(backend="nccl", timeout=self.ddp_timeout_delta) device = torch.device("cuda", self.local_rank) self._n_gpu = 1