patch for smddp import (huggingface#18244)

* add import * format
oneraghavan · Sep 26, 2022 · 858683a · 858683a
1 parent 174d422
commit 858683a
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/transformers/training_args.py b/src/transformers/training_args.py
@@ -1333,6 +1333,8 @@ def _setup_devices(self) -> "torch.device":
             device = torch.device("cuda", local_rank)
             self._n_gpu = 1
         elif is_sagemaker_dp_enabled():
+            import smdistributed.dataparallel.torch.torch_smddp  # noqa: F401
+
             dist.init_process_group(backend="smddp")
             self.local_rank = int(os.getenv("SMDATAPARALLEL_LOCAL_RANK"))
             device = torch.device("cuda", self.local_rank)