Remove USE_CUDA from c10d reducer/logger (#59562)

lw · facebook-github-bot · commit cbcae46fa5df · 2021-06-11T05:06:39.000-07:00
Summary: Pull Request resolved: #59562 Needed to merge c10d into libtorch(_cuda). ghstack-source-id: 131169542 Test Plan: CI Reviewed By: agolynski Differential Revision: D28931378 fbshipit-source-id: 71376b862ff6ef7dbfa7331ec8d269bd3fcc7e0d
diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl
@@ -669,6 +669,7 @@ libtorch_python_distributed_core_sources = [
     "torch/lib/c10d/comm.cpp",
     "torch/lib/c10d/default_comm_hooks.cpp",
     "torch/lib/c10d/reducer.cpp",
+    "torch/lib/c10d/reducer_cuda.cpp",
     "torch/lib/c10d/logger.cpp",
     "torch/csrc/distributed/c10d/python_comm_hook.cpp",
     "torch/csrc/distributed/c10d/init.cpp",
diff --git a/torch/lib/c10d/logger.cpp b/torch/lib/c10d/logger.cpp
@@ -9,12 +9,6 @@ namespace c10d {
 // stats.
 const int LoggingIterations[] = {10, 20, 100, 1000};
 
-namespace {
-
-const int kMilliSecondToNanosSecond = 1000000;
-
-} // anonymous namespace
-
 std::ostream& operator<<(std::ostream& output, const Logger& logger) {
   auto& ddp_logging_data = (*logger.ddp_logging_data_);
 
@@ -176,48 +170,21 @@ void Logger::set_construction_data_and_log(
   at::LogPyTorchDDPUsage(*ddp_logging_data_);
 }
 
-void Logger::calculate_avg_cpu_time(
-    int64_t& avg_time,
-    int64_t& time_duration,
-    int64_t cpu_start_time,
-    int64_t cpu_end_time) {
-  // If cpu_end_time is not recorded in this iteration,
-  // avg_time will return invalid value.
-  // For some cases like DDP runs on non-sync mode, backward compute
-  // end time can not be recorded in this iteration and thus can not
-  // calculate the valid avg_time.
-  // In this case, skip calculating the avg_time and return.
-  TORCH_CHECK(num_iterations_stats_recorded_ > 0);
-  if (cpu_end_time < cpu_start_time) {
-    return;
-  }
-  time_duration = cpu_end_time - cpu_start_time;
-  avg_time = (time_duration + avg_time * (num_iterations_stats_recorded_ - 1)) /
-      num_iterations_stats_recorded_;
-}
-
-#ifdef USE_CUDA
-void Logger::calculate_avg_gpu_time(
+void Logger::calculate_avg_time(
     int64_t& avg_time,
     int64_t& time_duration,
-    at::cuda::CUDAEvent& gpu_start,
-    at::cuda::CUDAEvent& gpu_end) {
+    Timer& timer,
+    Timer::Event start_event,
+    Timer::Event end_event) {
   TORCH_CHECK(num_iterations_stats_recorded_ > 0);
-  float milliseconds = gpu_start.elapsed_time(gpu_end);
-  // If gpu_end is not recorded in this iteration,
-  // milliseconds will have invalid value.
-  // For some cases like DDP runs on non-sync mode,
-  // gpu_end can not be recorded in this iteration and thus can not
-  // calculate the valid avg_time.
-  // In this case, skip calculating the avg_time and return.
-  if (milliseconds < 0) {
+  c10::optional<int64_t> maybe_time_duration = timer.measureDifference(start_event, end_event);
+  if (!maybe_time_duration.has_value()) {
     return;
   }
-  time_duration = int64_t(milliseconds * kMilliSecondToNanosSecond);
+  time_duration = maybe_time_duration.value();
   avg_time = (time_duration + avg_time * (num_iterations_stats_recorded_ - 1)) /
       num_iterations_stats_recorded_;
 }
-#endif
 
 void Logger::reset_performance_stats() {
   ddp_logging_data_->ints_map["forward_compute_time"] = 0;
@@ -260,85 +227,39 @@ void Logger::set_runtime_stats_and_log() {
 
   reset_performance_stats();
 
-  if (reducer_->replicas_[0][0].is_cuda()) {
-#ifdef USE_CUDA
-    // 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 multi-device modules."
-      );
-      return;
-    }
-    // Check events on the replicas_[0][0].device().
-    at::DeviceGuard g(reducer_->replicas_[0][0].device());
-    // It is possible users did not call backward or run codes in
-    // no-sync mode, in this case, some cudaEvents like "backward_compute_end"
-    // or "backward_comm_start" or "backward_comm_end" will not be recorded.
-    // cudaEvent is created when it is first time to be recorded.
-    // If it is never recorded/created, skip synchronize and calculation.
-    // Otherwise it will throw cuda errors.
-    if (!reducer_->gpu_timer_.forward_start.isCreated() ||
-        !reducer_->gpu_timer_.backward_compute_start.isCreated() ||
-        !reducer_->gpu_timer_.backward_compute_end.isCreated() ||
-        !reducer_->gpu_timer_.backward_comm_start.isCreated() ||
-        !reducer_->gpu_timer_.backward_comm_end.isCreated()) {
-      return;
-    }
-
-    // set_runtime_stats_and_log is called at the beginning of forward call,
-    // when it is cheap to synchronize the cuda events of previous iteration,
-    // as mostly all cuda operations are finished in previous iteration.
-    reducer_->gpu_timer_.forward_start.synchronize();
-    reducer_->gpu_timer_.backward_compute_start.synchronize();
-    reducer_->gpu_timer_.backward_compute_end.synchronize();
-    reducer_->gpu_timer_.backward_comm_start.synchronize();
-    reducer_->gpu_timer_.backward_comm_end.synchronize();
-    calculate_avg_gpu_time(
-        ddp_logging_data_->ints_map["avg_forward_compute_time"],
-        ddp_logging_data_->ints_map["forward_compute_time"],
-        reducer_->gpu_timer_.forward_start,
-        reducer_->gpu_timer_.backward_compute_start);
-    calculate_avg_gpu_time(
-        ddp_logging_data_->ints_map["avg_backward_compute_time"],
-        ddp_logging_data_->ints_map["backward_compute_time"],
-        reducer_->gpu_timer_.backward_compute_start,
-        reducer_->gpu_timer_.backward_compute_end);
-    calculate_avg_gpu_time(
-        ddp_logging_data_->ints_map["avg_backward_comm_time"],
-        ddp_logging_data_->ints_map["backward_comm_time"],
-        reducer_->gpu_timer_.backward_comm_start,
-        reducer_->gpu_timer_.backward_comm_end);
-    calculate_avg_gpu_time(
-        ddp_logging_data_->ints_map["avg_backward_compute_comm_overlap_time"],
-        ddp_logging_data_->ints_map["backward_compute_comm_overlap_time"],
-        reducer_->gpu_timer_.backward_comm_start,
-        reducer_->gpu_timer_.backward_compute_end);
-#endif
-  } else {
-    calculate_avg_cpu_time(
-        ddp_logging_data_->ints_map["avg_forward_compute_time"],
-        ddp_logging_data_->ints_map["forward_compute_time"],
-        reducer_->cpu_timer_.forward_start_time,
-        reducer_->cpu_timer_.backward_compute_start_time);
-
-    calculate_avg_cpu_time(
-        ddp_logging_data_->ints_map["avg_backward_compute_time"],
-        ddp_logging_data_->ints_map["backward_compute_time"],
-        reducer_->cpu_timer_.backward_compute_start_time,
-        reducer_->cpu_timer_.backward_compute_end_time);
-
-    calculate_avg_cpu_time(
-        ddp_logging_data_->ints_map["avg_backward_comm_time"],
-        ddp_logging_data_->ints_map["backward_comm_time"],
-        reducer_->cpu_timer_.backward_comm_start_time,
-        reducer_->cpu_timer_.backward_comm_end_time);
-
-    calculate_avg_cpu_time(
-        ddp_logging_data_->ints_map["avg_backward_compute_comm_overlap_time"],
-        ddp_logging_data_->ints_map["backward_compute_comm_overlap_time"],
-        reducer_->cpu_timer_.backward_comm_start_time,
-        reducer_->cpu_timer_.backward_compute_end_time);
+  // Cuda time stats are only collected for single device modules.
+  if (reducer_->replicas_[0][0].is_cuda() && reducer_->is_multi_device_module_) {
+    TORCH_WARN_ONCE(
+      "Cuda time stats are not collected for multi-device modules."
+    );
+    return;
   }
+  TORCH_INTERNAL_ASSERT(reducer_->timer_);
+  calculate_avg_time(
+      ddp_logging_data_->ints_map["avg_forward_compute_time"],
+      ddp_logging_data_->ints_map["forward_compute_time"],
+      *reducer_->timer_,
+      Timer::Event::kForwardStart,
+      Timer::Event::kBackwardComputeStart);
+  calculate_avg_time(
+      ddp_logging_data_->ints_map["avg_backward_compute_time"],
+      ddp_logging_data_->ints_map["backward_compute_time"],
+      *reducer_->timer_,
+      Timer::Event::kBackwardComputeStart,
+      Timer::Event::kBackwardComputeEnd);
+  calculate_avg_time(
+      ddp_logging_data_->ints_map["avg_backward_comm_time"],
+      ddp_logging_data_->ints_map["backward_comm_time"],
+      *reducer_->timer_,
+      Timer::Event::kBackwardCommStart,
+      Timer::Event::kBackwardCommEnd);
+  calculate_avg_time(
+      ddp_logging_data_->ints_map["avg_backward_compute_comm_overlap_time"],
+      ddp_logging_data_->ints_map["backward_compute_comm_overlap_time"],
+      *reducer_->timer_,
+      Timer::Event::kBackwardCommStart,
+      Timer::Event::kBackwardComputeEnd);
+
   // Log runtime stats to stderr if TORCH_DISTRIBUTED_DEBUG=DETAIL is enabled.
   if (parseDistDebugLevel() == DistributedDebugLevel::DETAIL) {
     LOG(INFO) << *this;
diff --git a/torch/lib/c10d/logger.hpp b/torch/lib/c10d/logger.hpp
@@ -41,18 +41,12 @@ class Logger {
 
   // Calculate avg stats using cpu timer and gpu timer
   // that has been recorded in reducer.
-  void calculate_avg_cpu_time(
+  void calculate_avg_time(
       int64_t& avg_time,
       int64_t& time_duration,
-      int64_t cpu_start_time,
-      int64_t cpu_end_time);
-#ifdef USE_CUDA
-  void calculate_avg_gpu_time(
-      int64_t& avg_time,
-      int64_t& time_duration,
-      at::cuda::CUDAEvent& gpu_start,
-      at::cuda::CUDAEvent& gpu_end);
-#endif
+      Timer& timer,
+      Timer::Event start_event,
+      Timer::Event end_event);
   // Set stats that can be collected only during
   // training loop. It is called at the beginning of forward call
   // to record the run time stats of sampled iterations that previouly ran.
diff --git a/torch/lib/c10d/reducer.cpp b/torch/lib/c10d/reducer.cpp
@@ -38,6 +38,67 @@ constexpr int kUnsetDivFactor = -1;
 
 } // namespace
 
+C10_DEFINE_TYPED_REGISTRY(TimerRegistry, c10::DeviceType, Timer, std::unique_ptr, c10::Device);
+
+namespace {
+
+class CpuTimer : public Timer {
+ private:
+  // The timestamp of forward call start time in each iteration.
+  int64_t forward_start_time = -1;
+  // The timestamp of backward computation start and end time in each
+  // iteration.
+  int64_t backward_compute_start_time = -1;
+  int64_t backward_compute_end_time = -1;
+  // The timestamp of first communication call start time in each iteration.
+  int64_t backward_comm_start_time = -1;
+  // The timestamp of last communication call end time in each iteration.
+  int64_t backward_comm_end_time = -1;
+
+  int64_t& getTime(Event event) {
+    switch (event) {
+      case Event::kForwardStart:
+        return forward_start_time;
+      case Event::kBackwardComputeStart:
+        return backward_compute_start_time;
+      case Event::kBackwardComputeEnd:
+        return backward_compute_end_time;
+      case Event::kBackwardCommStart:
+        return backward_comm_start_time;
+      case Event::kBackwardCommEnd:
+        return backward_comm_end_time;
+      default:
+        TORCH_INTERNAL_ASSERT(false);
+    }
+  }
+
+ public:
+  explicit CpuTimer(c10::Device /* unused */) {}
+
+  void record(Event event) override {
+    getTime(event) = current_time_in_nanos();
+  }
+
+  c10::optional<int64_t> measureDifference(Event start, Event end) override {
+    int64_t start_time = getTime(start);
+    int64_t end_time = getTime(end);
+    // If cpu_end_time is not recorded in this iteration,
+    // avg_time will return invalid value.
+    // For some cases like DDP runs on non-sync mode, backward compute
+    // end time can not be recorded in this iteration and thus can not
+    // calculate the valid avg_time.
+    // In this case, skip calculating the avg_time and return.
+    if (end_time < start_time) {
+      return c10::nullopt;
+    }
+    return end_time - start_time;
+  }
+};
+
+C10_REGISTER_TYPED_CLASS(TimerRegistry, c10::kCPU, CpuTimer);
+
+} // namespace
+
 Reducer::Reducer(
     std::vector<std::vector<at::Tensor>> replicas,
     std::vector<std::vector<size_t>> bucket_indices,
@@ -88,6 +149,12 @@ Reducer::Reducer(
     }
   }
 
+  // For CUDA, record events only for single device module.
+  c10::Device device = replicas_[0][0].device();
+  if (!(device.is_cuda() && is_multi_device_module_)) {
+    timer_ = TimerRegistry()->Create(device.type(), device);
+  }
+
   // If `expect_sparse_gradients` is not specified, initialize it such that
   // we do not expect sparse gradients for any parameter.
   if (expect_sparse_gradients_.empty()) {
@@ -724,7 +791,7 @@ void Reducer::mark_variable_ready(size_t variable_index) {
   checkAndRaiseMarkedTwiceError(variable_index);
   perIterationReadyParams_.insert(variable_index);
   backward_stats_[0][variable_index] =
-      current_time_in_nanos() - cpu_timer_.backward_compute_start_time;
+      current_time_in_nanos() - backward_compute_start_time_;
 
   // Any time we mark a variable ready (be it in line due to unused parameters,
   // or via an autograd hook), we require a call to the finalize function. If
@@ -1180,7 +1247,7 @@ void Reducer::prepare_for_backward(
     const std::vector<torch::autograd::Variable>& outputs) {
   std::lock_guard<std::mutex> lock(mutex_);
 
-  cpu_timer_.backward_compute_start_time = current_time_in_nanos();
+  backward_compute_start_time_ = current_time_in_nanos();
   if (should_collect_runtime_stats()) {
     record_backward_compute_start_time();
   }
@@ -1723,72 +1790,32 @@ 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 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();
-    }
-#endif
-  } else {
-    cpu_timer_.forward_start_time = current_time_in_nanos();
+  if (timer_) {
+    timer_->record(Timer::Event::kForwardStart);
   }
 }
 
 void Reducer::record_backward_compute_start_time() {
-  if (replicas_[0][0].is_cuda()) {
-#ifdef USE_CUDA
-    // 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();
-    }
-#endif
+  if (timer_) {
+    timer_->record(Timer::Event::kBackwardComputeStart);
   }
 }
 
 void Reducer::record_backward_compute_end_time() {
-  if (replicas_[0][0].is_cuda()) {
-#ifdef USE_CUDA
-    // 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();
-    }
-#endif
-  } else {
-    cpu_timer_.backward_compute_end_time = current_time_in_nanos();
+  if (timer_) {
+    timer_->record(Timer::Event::kBackwardComputeEnd);
   }
 }
 
 void Reducer::record_backward_comm_start_time() {
-  if (replicas_[0][0].is_cuda()) {
-#ifdef USE_CUDA
-    // 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();
-    }
-#endif
-  } else {
-    cpu_timer_.backward_comm_start_time = current_time_in_nanos();
+  if (timer_) {
+    timer_->record(Timer::Event::kBackwardCommStart);
   }
 }
 
 void Reducer::record_backward_comm_end_time() {
-  if (replicas_[0][0].is_cuda()) {
-#ifdef USE_CUDA
-    // 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();
-    }
-#endif
-  } else {
-    cpu_timer_.backward_comm_end_time = current_time_in_nanos();
+  if (timer_) {
+    timer_->record(Timer::Event::kBackwardCommEnd);
   }
 }
 
diff --git a/torch/lib/c10d/reducer.hpp b/torch/lib/c10d/reducer.hpp
diff --git a/torch/lib/c10d/reducer_cuda.cpp b/torch/lib/c10d/reducer_cuda.cpp
