[DPER] Introduce barrier operation to force synchronization of threads in async execution (#49322)

kennyhorror · facebook-github-bot · commit 46debe7f2361 · 2020-12-15T16:13:42.000-08:00
Summary: Pull Request resolved: #49322 In some cases async execution might loose dependencies (Alias like ops) or produce suboptimal scheduling when there is an option which parts to schedule first. Example of the later behavior can happen in ModelParallel training where copy can get lower priority compared to the rest of the execution on the given GPU, which will caused other GPUs to starve. This operator allows to address these issues by introducing extra explicit dependencies between ops. Test Plan: Unit-test/ E2E testing in the future diffs. Reviewed By: xianjiec Differential Revision: D24933471 fbshipit-source-id: 1668994c7856d73926cde022378a99e1e8db3567
diff --git a/caffe2/operators/async_net_barrier_op.cc b/caffe2/operators/async_net_barrier_op.cc
@@ -0,0 +1,50 @@
+#include "caffe2/operators/async_net_barrier_op.h"
+
+namespace caffe2 {
+
+namespace {
+std::pair<std::vector<DeviceOption>, std::vector<DeviceOption>>
+asyncBarrierOpDevInfer(const OperatorDef& def) {
+  auto op_device =
+      def.has_device_option() ? def.device_option() : DeviceOption();
+  ArgumentHelper helper(def);
+  auto cross_device = helper.GetSingleArgument<int>("cross_device", 0);
+  std::vector<DeviceOption> opt;
+  for (int i = 0; i < def.input().size(); ++i) {
+    if (cross_device == 1) {
+      DeviceOption dev;
+      dev.set_device_type(op_device.device_type());
+      dev.set_device_id(i);
+      opt.push_back(dev);
+    } else {
+      opt.push_back(op_device);
+    }
+  }
+  return std::make_pair(opt, opt);
+}
+}
+
+OPERATOR_SCHEMA(AsyncNetBarrier)
+    .NumInputs(1, INT_MAX)
+    .NumOutputs(1, INT_MAX)
+    .IdenticalTypeAndShape()
+    .InputsCanCrossDevices()
+    .AllowOneToOneInplace()
+    .DeviceInferenceFunction(asyncBarrierOpDevInfer)
+    .SetDoc(R"DOC(
+This is a pretty much no-op operator, since it's only purposes is make sure that
+async_scheduling will schedule certian operations earlier than others.
+
+Exaple where this operator can work well - mixture of data-parallel and model-
+parallel training, where one wants to force that all copies are started before
+data-parallel part starts.
+)DOC")
+    .Arg(
+        "cross_device",
+        "Specifies either inputs should be across different devices in dev inference options");
+
+SHOULD_NOT_DO_GRADIENT(AsyncNetBarrier);
+REGISTER_CPU_OPERATOR(AsyncNetBarrier, AsyncNetBarrierOp<CPUContext>);
+
+
+} // namespace caffe2
diff --git a/caffe2/operators/async_net_barrier_op.cu b/caffe2/operators/async_net_barrier_op.cu
@@ -0,0 +1,8 @@
+#include "caffe2/core/context_gpu.h"
+#include "caffe2/operators/async_net_barrier_op.h"
+
+namespace caffe2 {
+
+REGISTER_CUDA_OPERATOR(AsyncNetBarrier, AsyncNetBarrierOp<CUDAContext>);
+
+} // namespace caffe2
diff --git a/caffe2/operators/async_net_barrier_op.h b/caffe2/operators/async_net_barrier_op.h
@@ -0,0 +1,30 @@
+#ifndef CAFFE2_OPERATORS_ASYNC_BARRIER_OP_H_
+#define CAFFE2_OPERATORS_ASYNC_BARRIER_OP_H_
+
+#include "caffe2/core/context.h"
+#include "caffe2/core/export_caffe2_op_to_c10.h"
+#include "caffe2/core/operator.h"
+
+namespace caffe2 {
+
+template <class Context>
+class AsyncNetBarrierOp : public Operator<Context> {
+ public:
+  USE_OPERATOR_CONTEXT_FUNCTIONS;
+  USE_SIMPLE_CTOR_DTOR(AsyncNetBarrierOp)
+
+  bool RunOnDevice() override {
+    // This is a pretty much no-op operator, since it's only purposes is make
+    // sure that async_scheduling will schedule certian operations earlier than
+    // others.
+    //
+    // Exaple where this operator can work well - mixture of data-parallel and
+    // model parallel training, where one wants to force that all copies are
+    // started before data-parallel part starts.
+    return true;
+  }
+};
+
+} // namespace caffe2
+
+#endif // CAFFE2_OPERATORS_ASYNC_BARRIER_OP_H_
diff --git a/caffe2/python/operator_test/async_net_barrier_test.py b/caffe2/python/operator_test/async_net_barrier_test.py
@@ -0,0 +1,31 @@
+#!/usr/bin/env python3
+
+import caffe2.python.hypothesis_test_util as hu
+import hypothesis.strategies as st
+import numpy as np
+from caffe2.python import core
+from hypothesis import given
+
+
+class TestAsyncNetBarrierOp(hu.HypothesisTestCase):
+    @given(
+        n=st.integers(1, 5),
+        shape=st.lists(st.integers(0, 5), min_size=1, max_size=3),
+        **hu.gcs
+    )
+    def test_async_net_barrier_op(self, n, shape, dc, gc):
+        test_inputs = [(100 * np.random.random(shape)).astype(np.float32) for _ in range(n)]
+        test_input_blobs = ["x_{}".format(i) for i in range(n)]
+
+        barrier_op = core.CreateOperator(
+            "AsyncNetBarrier",
+            test_input_blobs,
+            test_input_blobs,
+            device_option=gc,
+        )
+
+        def reference_func(*args):
+            self.assertEquals(len(args), n)
+            return args
+
+        self.assertReferenceChecks(gc, barrier_op, test_inputs, reference_func)
