Accumulate opt state dict on do_rank 0 (#11481)

src/transformers/trainer.py

@@ -1420,6 +1420,7 @@ class Trainer:
                 xm.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt"))
                 reissue_pt_warnings(caught_warnings)
         elif is_sagemaker_mp_enabled():
+            if smp.dp_rank() == 0:
                 # Consolidate the state dict on all processed of dp_rank 0
                 opt_state_dict = self.optimizer.state_dict()
                 # Save it and the scheduler on the main process