[Reducer] Remove replica size == 1 checks (#58603)

rohan-varma · facebook-github-bot · commit b3015584104e · 2021-05-20T22:34:23.000-07:00
Summary: Pull Request resolved: #58603 No longer need these checks ghstack-source-id: 129498227 Test Plan: CI Reviewed By: zhaojuanmao Differential Revision: D28549893 fbshipit-source-id: a89bf8c3fc3aba311a70fd37e5a6aa5dc14b41b9
diff --git a/torch/lib/c10d/logger.cpp b/torch/lib/c10d/logger.cpp
@@ -260,12 +260,11 @@ void Logger::set_runtime_stats_and_log() {
 
   if (reducer_->replicas_[0][0].is_cuda()) {
 #ifdef USE_CUDA
-    // Cuda time stats are only collected for single process single
-    // device and single device module.
-    if (reducer_->replicas_.size() > 1 || reducer_->is_multi_device_module_) {
+    // Cuda time stats are only collected for single device modules.
+    if (reducer_->is_multi_device_module_) {
       TORCH_WARN_ONCE(
-          "Cuda time stats are not collected for single process "
-          "multiple device program or multi-device modules.");
+        "Cuda time stats are not collected for multi-device modules."
+      );
       return;
     }
     // Check events on the replicas_[0][0].device().
diff --git a/torch/lib/c10d/reducer.cpp b/torch/lib/c10d/reducer.cpp
@@ -1725,9 +1725,8 @@ bool Reducer::should_collect_runtime_stats() {
 void Reducer::record_forward_compute_start_time() {
   if (replicas_[0][0].is_cuda()) {
 #ifdef USE_CUDA
-    // Record event only for single process single device
-    // and single device module.
-    if (replicas_.size() == 1 && !is_multi_device_module_) {
+    // Record event only for single device module.
+    if (!is_multi_device_module_) {
       // Create and record event on the replicas_[0][0].device().
       at::DeviceGuard g(replicas_[0][0].device());
       gpu_timer_.forward_start.record();
@@ -1741,9 +1740,8 @@ void Reducer::record_forward_compute_start_time() {
 void Reducer::record_backward_compute_start_time() {
   if (replicas_[0][0].is_cuda()) {
 #ifdef USE_CUDA
-    // Record event only for single process single device
-    // and single device module.
-    if (replicas_.size() == 1 && !is_multi_device_module_) {
+    // Record event only for single device module.
+    if (!is_multi_device_module_) {
       // Create and record event on the replicas_[0][0].device().
       at::DeviceGuard g(replicas_[0][0].device());
       gpu_timer_.backward_compute_start.record();
@@ -1755,9 +1753,8 @@ void Reducer::record_backward_compute_start_time() {
 void Reducer::record_backward_compute_end_time() {
   if (replicas_[0][0].is_cuda()) {
 #ifdef USE_CUDA
-    // Record event only for single process single device
-    // and single device module.
-    if (replicas_.size() == 1 && !is_multi_device_module_) {
+    // Record event only for single device module.
+    if (!is_multi_device_module_) {
       at::DeviceGuard g(replicas_[0][0].device());
       gpu_timer_.backward_compute_end.record();
     }
@@ -1770,9 +1767,8 @@ void Reducer::record_backward_compute_end_time() {
 void Reducer::record_backward_comm_start_time() {
   if (replicas_[0][0].is_cuda()) {
 #ifdef USE_CUDA
-    // Record event only for single process single device
-    // and single device module.
-    if (replicas_.size() == 1 && !is_multi_device_module_) {
+    // Record event only for single device module
+    if (!is_multi_device_module_) {
       at::DeviceGuard g(replicas_[0][0].device());
       gpu_timer_.backward_comm_start.record();
     }
@@ -1785,9 +1781,8 @@ void Reducer::record_backward_comm_start_time() {
 void Reducer::record_backward_comm_end_time() {
   if (replicas_[0][0].is_cuda()) {
 #ifdef USE_CUDA
-    // Record event only for single process single device
-    // and single device module.
-    if (replicas_.size() == 1 && !is_multi_device_module_) {
+    // Record event only for single device module.
+    if (!is_multi_device_module_) {
       at::DeviceGuard g(replicas_[0][0].device());
       gpu_timer_.backward_comm_end.record();
     }
