patch for smddp import (#18244)

* add import * format

patch for smddp import (#18244)
* add import * format
ee67e7ad · Carolyn Wang · GitHub · 68097dcc · ee67e7ad
Unverified Commit ee67e7ad authored Jul 26, 2022 by Carolyn Wang Committed by GitHub Jul 26, 2022
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 0 deletions

src/transformers/training_args.py src/transformers/training_args.py +2 -0

No files found.
--- a/src/transformers/training_args.py
+++ b/src/transformers/training_args.py
@@ -1333,6 +1333,8 @@ class TrainingArguments:
            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)