Implement optimization of overlapping wait communication with forward; Support time measurement of staring and waiting communication

Change-Id: I6c5af8d85fbb8a81353142752a37dfce8cd1870d

Author: fzou1

Makefile

@@ -77,11 +77,13 @@ ifeq ($(CAFFE_PER_LAYER_TIMINGS), 1)
 endif
 
 ifeq ($(CAFFE_MLSL_SHUFFLE), 1)
-        COMMON_FLAGS += -DCAFFE_MLSL_SHUFFLE
+	COMMON_FLAGS += -DCAFFE_MLSL_SHUFFLE
 endif
 
+ifneq ($(FW_OVERLAP_OPT), 0)
+	COMMON_FLAGS += -DFW_OVERLAP_OPT
+endif
 endif
-
 #################### MLSL ####################
 
 

cmake/Dependencies.cmake

@@ -115,6 +115,17 @@ if(USE_MLSL)
   include_directories(SYSTEM "${MLSL_ROOT}/intel64/include")
   link_directories(SYSTEM "${MLSL_ROOT}/intel64/lib")
   list(APPEND Caffe_LINKER_LIBS mlsl)
+
+  if(CAFFE_PER_LAYER_TIMINGS)
+    add_definitions("-DCAFFE_PER_LAYER_TIMINGS")
+  endif()
+  if(CAFFE_MLSL_SHUFFLE)
+    add_definitions("-DCAFFE_MLSL_SHUFFLE")
+  endif()
+  if(FW_OVERLAP_OPT OR NOT DEFINED FW_OVERLAP_OPT)
+    message(STATUS "Forward overlapping optimization is enabled!")
+    add_definitions("-DFW_OVERLAP_OPT")
+  endif()
 endif()
 
 # ---[ BLAS

include/caffe/multinode/multi_solver.hpp

@@ -60,6 +60,12 @@ class MultiSolver {
       iter_size(root_solver_->param().iter_size()) {
     root_solver_->set_forward_backward(
       boost::bind(&MultiSolver<Dtype>::ForwardBackward, this));
+#ifdef FW_OVERLAP_OPT    
+    Net<Dtype>& net = *root_solver_->net();
+    const std::vector<shared_ptr<Layer<Dtype>>> & layers{ net.layers() };
+    layer_finished_flags_.resize(layers.size());
+    std::fill(layer_finished_flags_.begin(), layer_finished_flags_.end(), false);
+#endif
   }
 
 
@@ -99,14 +105,23 @@ class MultiSolver {
   boost::shared_ptr<Solver<Dtype>> root_solver() {
     return root_solver_;
   }
-
+#ifdef FW_OVERLAP_OPT
+  void set_layer_finished_flag(int layer_id, bool flag) {
+    layer_finished_flags_[layer_id] = flag;
+  }
+#endif
  private:
   virtual Dtype ForwardBackwardImpl(bool first, bool last);
+  bool IsSkipWaitGradient(int layer_id);
+  void WaitAndUpdateGradient(int layer_id);
 
  protected:
   boost::shared_ptr<Solver<Dtype>> root_solver_;
   int iter_size;
   vector<Callback*> callbacks_;
+#ifdef FW_OVERLAP_OPT
+  vector<bool> layer_finished_flags_;
+#endif
 };
 
 }  // namespace caffe

include/caffe/multinode/multi_sync.hpp

@@ -74,6 +74,10 @@ namespace caffe {
     shared_ptr<Net<Dtype>> net;
     const vector<Blob<Dtype> *> &net_params;
     vector<vector<int>> layer_param_ids;
+#ifdef FW_OVERLAP_OPT
+    vector<vector<bool>> param_ids_finished_flags;
+#endif
+
     // layer_id -> blob_id -> cached blob to restore
     // statistics
     vector<vector<shared_ptr<Blob<Dtype>>>> cached_stats;
@@ -160,6 +164,12 @@ namespace caffe {
                    << " ENABLED"
 #else
                    << " DISABLED"
+#endif
+                   << ", FORWARD OVERLAP OPTIMIZATION IS"
+#ifdef FW_OVERLAP_OPT
+                   << " ENABLED"
+#else
+                   << " DISABLED"
 #endif
                    << ", SINGLE DB SPLITTING IS"
 #ifdef CAFFE_MLSL_SHUFFLE
@@ -201,6 +211,12 @@ namespace caffe {
         return;
       }
 
+#ifdef FW_OVERLAP_OPT
+      std::fill(param_ids_finished_flags[layer_id].begin(),
+          param_ids_finished_flags[layer_id].end(),
+          false);
+#endif
+
       std::vector<int> &param_ids = layer_param_ids[layer_id];
       for (int i = 0; i < param_ids.size(); ++i) {
         if (!layer->ParamNeedReduce(i)) continue;
@@ -215,15 +231,41 @@ namespace caffe {
     void on_delwt_wait(int layer_id) {
       boost::shared_ptr<Layer<Dtype>> &layer = layers[layer_id];
       if (layer->layerOp == nullptr) {
+#ifdef FW_OVERLAP_OPT
+        solver->set_layer_finished_flag(layer_id, true);
+#endif
         return;
       }
 
       std::vector<int> &param_ids = layer_param_ids[layer_id];
 
+#ifdef FW_OVERLAP_OPT
+      int finished_count = 0;
+#endif
+
       for (int i=0; i<param_ids.size(); i++) {
-        if (!layer->ParamNeedReduce(i)) continue;
+        if (!layer->ParamNeedReduce(i)
+#ifdef FW_OVERLAP_OPT
+            || (param_ids_finished_flags[layer_id][i] == true)) {
+          finished_count++;
+#else
+          ) {
+#endif
+          continue;
+        }
+
+#ifdef FW_OVERLAP_OPT
+        bool is_completed = false;
+        Dtype *delwt_buf{(Dtype *) layer->layerOp->GetParameterSet(i)->TestGradientComm(&is_completed)};
+#else
         Dtype *delwt_buf{(Dtype *) layer->layerOp->GetParameterSet(i)->WaitGradientComm()};
+#endif
         if (delwt_buf) {
+#ifdef FW_OVERLAP_OPT
+          assert(is_completed);
+          param_ids_finished_flags[layer_id][i] = true;
+          finished_count++;
+#endif
           if (CAN_USE_PRV(net_params[param_ids[i]])) {
             if (delwt_buf != net_params[param_ids[i]]->prv_diff())
               caffe_copy(net_params[param_ids[i]]->count(),
@@ -235,6 +277,12 @@ namespace caffe {
                 net_params[param_ids[i]]->mutable_cpu_diff());
         }
       }
+
+#ifdef FW_OVERLAP_OPT
+      if (finished_count == param_ids.size()) {
+        solver->set_layer_finished_flag(layer_id, true);
+      }
+#endif
     }
 
     void on_gradients_ready() {

include/caffe/solver.hpp

@@ -165,10 +165,18 @@ class Solver {
   std::vector<double> forward_time_per_layer;
   std::vector<double> backward_time_per_layer;
   std::vector<double> update_time_per_layer;
+#ifdef USE_MLSL
+  std::vector<double> startcomm_time_per_layer;
+  std::vector<double> waitcomm_time_per_layer;
+#endif
 
   std::vector<double> forward_time_per_layer_total;
   std::vector<double> backward_time_per_layer_total;
   std::vector<double> update_time_per_layer_total;
+#ifdef USE_MLSL
+  std::vector<double> startcomm_time_per_layer_total;
+  std::vector<double> waitcomm_time_per_layer_total;
+#endif
 
   void InitTimers();
   void ResetTimers();

src/caffe/multinode/multi_solver.cpp

@@ -46,77 +46,167 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 namespace caffe {
 
+#define START_ITER 1
+
+
+#ifdef CAFFE_PER_LAYER_TIMINGS
+#define LAYER_TIMING_START() do { \
+  timer.Start(); \
+}while(0)
+
+#define LAYER_TIMING_STOP(name, index) do { \
+  name##_time_per_layer[index] += timer.MicroSeconds(); \
+}while(0)
+#else
+#define LAYER_TIMING_START()
+
+#define LAYER_TIMING_STOP(name,index)
+#endif
+
+template <typename Dtype>
+inline bool MultiSolver<Dtype>::IsSkipWaitGradient(int layer_id) {
+  Net<Dtype>& net = *root_solver_->net();
+  const std::vector<shared_ptr<Layer<Dtype>>>& layers{ net.layers() };
+  const std::vector<bool>& layer_need_backward{ net.layer_need_backward() };
+
+  if (!layer_need_backward[layer_id] || ((layers[layer_id]->layerOp != nullptr)
+        && !layers[layer_id]->layerOp->HasParameterSets())) {
+      DLOG(INFO) << "ForwardBackwardImpl: no need for apply_updates for layer # "
+        << layer_id << ", skip on_delwt_wait, apply_updates, on_wtinc_ready";
+      return true;
+  }
+  return false;
+}
+
+template <typename Dtype>
+inline void MultiSolver<Dtype>::WaitAndUpdateGradient(int layer_id) {
+  LAYER_TIMING_START();
+  for (int j = 0; j < callbacks_.size(); ++j) {
+    callbacks_[j]->on_delwt_wait(layer_id);
+  }
+  LAYER_TIMING_STOP(waitcomm, layer_id);
+
+#ifdef FW_OVERLAP_OPT
+  if (layer_finished_flags_[layer_id]) {
+#endif
+    LAYER_TIMING_START();
+    for (int j = 0; j < callbacks_.size(); ++j) {
+      callbacks_[j]->apply_updates(layer_id);
+    }
+    LAYER_TIMING_STOP(update, layer_id);
+#ifdef FW_OVERLAP_OPT
+  }
+#endif
+}
+
 template <typename Dtype>
 Dtype MultiSolver<Dtype>::ForwardBackwardImpl(bool first, bool last) {
 
   Dtype loss = 0;
   Net<Dtype>& net = *root_solver_->net();
   const std::vector<shared_ptr<Layer<Dtype>>>& layers{ net.layers() };
   const std::vector<bool>& layer_need_backward{ net.layer_need_backward() };
+#ifdef FW_OVERLAP_OPT
+  int iter = root_solver_->iter();
+#endif
 
 #ifdef CAFFE_PER_LAYER_TIMINGS
   Timer& timer = root_solver_->timer;
   std::vector<double>& forward_time_per_layer = root_solver_->forward_time_per_layer;
   std::vector<double>& backward_time_per_layer = root_solver_->backward_time_per_layer;
   std::vector<double>& update_time_per_layer = root_solver_->update_time_per_layer;
+  std::vector<double>& startcomm_time_per_layer = root_solver_->startcomm_time_per_layer;
+  std::vector<double>& waitcomm_time_per_layer = root_solver_->waitcomm_time_per_layer;
 #endif /* CAFFE_PER_LAYER_TIMINGS */
 
+
   for (int i = 0; i < layers.size(); ++i) {
-#ifdef CAFFE_PER_LAYER_TIMINGS
-    timer.Start();
+#ifdef FW_OVERLAP_OPT
+    if (first && iter >= START_ITER + 1) {
+      while (layer_finished_flags_[i] == false) {
+        if (IsSkipWaitGradient(i)) {
+          break;
+        }
+
+        WaitAndUpdateGradient(i);
+        if (layer_finished_flags_[i]) {
+          break;
+        }
+
+        for (int k=i+1; k<layers.size(); k++) {
+          if (layer_finished_flags_[k] || IsSkipWaitGradient(k)) {
+            layer_finished_flags_[k] = true;
+            continue;
+          }
+
+          WaitAndUpdateGradient(k);
+          if (layer_finished_flags_[k])
+            break;
+        }
+      }
+      layer_finished_flags_[i] = false;
+    }
 #endif
-    loss += net.ForwardFromTo(i, i);
 
-#ifdef CAFFE_PER_LAYER_TIMINGS
-    forward_time_per_layer[i] += timer.MicroSeconds();
-#endif
+    LAYER_TIMING_START();
+    loss += net.ForwardFromTo(i, i);
+    LAYER_TIMING_STOP(forward, i);
   }
 
   for (int i = layers.size() - 1; i >= 0; --i) {
-#ifdef CAFFE_PER_LAYER_TIMINGS
-    timer.Start();
-#endif
-
     if (!layer_need_backward[i]) {
       continue;
     }
+    
+    LAYER_TIMING_START();
 
     net.BackwardFromTo(i, i);
 
+    LAYER_TIMING_STOP(backward, i);
+
     if (last && (layers[i]->layerOp != nullptr) && layers[i]->layerOp->HasParameterSets()) {
+      LAYER_TIMING_START();
       for (int j = 0; j < callbacks_.size(); ++j) {
-          callbacks_[j]->on_iter_finished(i);
+        callbacks_[j]->on_iter_finished(i);
       }
+      LAYER_TIMING_STOP(startcomm, i);
     }
-
-#ifdef CAFFE_PER_LAYER_TIMINGS
-    backward_time_per_layer[i] += timer.MicroSeconds();
-#endif
   }
 
+#ifdef FW_OVERLAP_OPT
+  int max_iter = root_solver_->param().max_iter();
+  bool test = (root_solver_->param().test_interval()
+          && ((iter + 1) % root_solver_->param().test_interval() == 0));
+  if (last && (test || (iter == max_iter - 1))) {
+    int finished_count = 0;
+    while (finished_count < layers.size()) {
+#else
   if (last) {
+#endif
 
     for (int i = 0; i < layers.size(); ++i) {
-#ifdef CAFFE_PER_LAYER_TIMINGS
-      timer.Start();
+#ifdef FW_OVERLAP_OPT
+        if (layer_finished_flags_[i])
+          continue;
 #endif
-      if (!layer_need_backward[i] || ((layers[i]->layerOp != nullptr) && !layers[i]->layerOp->HasParameterSets())) {
-        DLOG(INFO) << "ForwardBackwardImpl: no need for apply_updates for layer # " << i
-          << ", skip on_delwt_wait, apply_updates, on_wtinc_ready";
-        continue;
-      }
+        if (IsSkipWaitGradient(i)) {
+#ifdef FW_OVERLAP_OPT
+          finished_count++;
+          layer_finished_flags_[i] = true;
+#endif
+          continue;
+        }
 
-      for (int j = 0; j < callbacks_.size(); ++j) {
-        callbacks_[j]->on_delwt_wait(i);
-      }
+        WaitAndUpdateGradient(i);
 
-      for (int j = 0; j < callbacks_.size(); ++j) {
-        callbacks_[j]->apply_updates(i);
-      }
-#ifdef CAFFE_PER_LAYER_TIMINGS
-      update_time_per_layer[i] += timer.MicroSeconds();
+#ifdef FW_OVERLAP_OPT
+        if (layer_finished_flags_[i])
+          finished_count++;
 #endif
+      }
+#ifdef FW_OVERLAP_OPT
     }
+#endif
   }
 
   DLOG(WARNING) << "iter " << root_solver_->iter() << ", loss " << loss;