enable CreateCommonWorld to bootstrap from existing common world

Summary: Use romain-intel's ContextFactory to create common worlds from existing common worlds, thus bypassing KV store completely. Changed data_parallel_model to automatically find if there is already a CW we can work. CreateCommonWorldOp takes optional second parameter, which is existing CW.

Reviewed By: andrewwdye

Differential Revision: D5494956

fbshipit-source-id: 5f7a840bcd5fe4ea756fafeacc746bc2cf5078b0

Author: Aapo Kyrola

caffe2/contrib/gloo/common_world_ops.h

@@ -52,9 +52,24 @@ class CreateCommonWorld final : public Operator<Context> {
 
     try {
       // Create context and connect everyone to everyone
-      auto context =
-          std::make_shared<::gloo::rendezvous::Context>(rank_, size_);
-      context->connectFullMesh(store, device_);
+      std::shared_ptr<::gloo::Context> context;
+
+      if (InputSize() == 1) {
+        auto new_context =
+            std::make_shared<::gloo::rendezvous::Context>(rank_, size_);
+        new_context->connectFullMesh(store, device_);
+        context = std::move(new_context);
+      } else {
+        VLOG(1) << "Creating new common world by forking existing one.";
+        auto backingCommonWorld =
+            OperatorBase::Input<std::shared_ptr<::gloo::Context>>(EXISTING_CW);
+
+        // Confirm the backing common world is compatible with this op
+        CAFFE_ENFORCE_EQ(rank_, backingCommonWorld->rank);
+        CAFFE_ENFORCE_EQ(size_, backingCommonWorld->size);
+        ::gloo::rendezvous::ContextFactory factory(backingCommonWorld);
+        context = factory.makeContext(device_);
+      }
 
       // Switch pairs to synchronous mode if configured to do so
       if (sync_) {
@@ -100,7 +115,7 @@ class CreateCommonWorld final : public Operator<Context> {
   Workspace* ws_;
   std::string status_blob_;
 
-  INPUT_TAGS(STORE_HANDLER);
+  INPUT_TAGS(STORE_HANDLER, EXISTING_CW);
   OUTPUT_TAGS(COMM);
 };
 

caffe2/contrib/gloo/gloo_test.py

@@ -87,34 +87,40 @@ def run_test_distributed(self, fn, device_option=None, **kwargs):
             fn(**kwargs)
             workspace.ResetWorkspace()
 
-    def create_common_world(self, comm_rank, comm_size, tmpdir=None):
+    def create_common_world(self, comm_rank, comm_size, tmpdir=None, existing_cw=None):
         store_handler = "store_handler"
 
         # If REDIS_HOST is set, use RedisStoreHandler for rendezvous.
-        redis_host = os.getenv("REDIS_HOST")
-        redis_port = int(os.getenv("REDIS_PORT", 6379))
-        if redis_host is not None:
-            workspace.RunOperatorOnce(
-                core.CreateOperator(
-                    "RedisStoreHandlerCreate",
-                    [],
-                    [store_handler],
-                    prefix=str(TestCase.test_counter) + "/",
-                    host=redis_host,
-                    port=redis_port))
+        if existing_cw is None:
+            redis_host = os.getenv("REDIS_HOST")
+            redis_port = int(os.getenv("REDIS_PORT", 6379))
+            if redis_host is not None:
+                workspace.RunOperatorOnce(
+                    core.CreateOperator(
+                        "RedisStoreHandlerCreate",
+                        [],
+                        [store_handler],
+                        prefix=str(TestCase.test_counter) + "/",
+                        host=redis_host,
+                        port=redis_port))
+            else:
+                workspace.RunOperatorOnce(
+                    core.CreateOperator(
+                        "FileStoreHandlerCreate",
+                        [],
+                        [store_handler],
+                        path=tmpdir))
+            common_world = "common_world"
         else:
-            workspace.RunOperatorOnce(
-                core.CreateOperator(
-                    "FileStoreHandlerCreate",
-                    [],
-                    [store_handler],
-                    path=tmpdir))
+            common_world = str(existing_cw) + ".forked"
 
-        common_world = "common_world"
+        inputs = [store_handler]
+        if existing_cw is not None:
+            inputs.append(existing_cw)
         workspace.RunOperatorOnce(
             core.CreateOperator(
                 "CreateCommonWorld",
-                [store_handler],
+                inputs,
                 [common_world],
                 size=comm_size,
                 rank=comm_rank,
@@ -269,6 +275,45 @@ def _test_allreduce(self,
         for _tmp in range(4):
             workspace.RunNet(net.Name())
 
+    def _test_allreduce_multicw(self,
+                                comm_rank=None,
+                                comm_size=None,
+                                tmpdir=None
+                                ):
+        _store_handler, common_world = self.create_common_world(
+            comm_rank=comm_rank,
+            comm_size=comm_size,
+            tmpdir=tmpdir)
+
+        _, common_world2 = self.create_common_world(
+            comm_rank=comm_rank,
+            comm_size=comm_size,
+            tmpdir=tmpdir,
+            existing_cw=common_world)
+
+        blob_size = 1e4
+        num_blobs = 4
+
+        for cw in [common_world, common_world2]:
+            blobs = []
+            for i in range(num_blobs):
+                blob = "blob_{}".format(i)
+                value = np.full(blob_size, (comm_rank * num_blobs) + i, np.float32)
+                workspace.FeedBlob(blob, value)
+                blobs.append(blob)
+
+            net = core.Net("allreduce_multicw")
+            net.Allreduce(
+                [cw] + blobs,
+                blobs,
+                engine=op_engine)
+
+            workspace.RunNetOnce(net)
+            for i in range(num_blobs):
+                np.testing.assert_array_equal(
+                    workspace.FetchBlob(blobs[i]),
+                    (num_blobs * comm_size) * (num_blobs * comm_size - 1) / 2)
+
     @given(comm_size=st.integers(min_value=2, max_value=8),
            blob_size=st.integers(min_value=1e3, max_value=1e6),
            num_blobs=st.integers(min_value=1, max_value=4),
@@ -295,6 +340,21 @@ def test_allreduce(self, comm_size, blob_size, num_blobs, device_option,
                     tmpdir=tmpdir,
                     use_float16=use_float16)
 
+    @given(device_option=st.sampled_from([hu.cpu_do]))
+    def test_forked_cw(self, device_option):
+        TestCase.test_counter += 1
+        if os.getenv('COMM_RANK') is not None:
+            self.run_test_distributed(
+                self._test_allreduce_multicw,
+                device_option=device_option)
+        else:
+            with TemporaryDirectory() as tmpdir:
+                self.run_test_locally(
+                    self._test_allreduce_multicw,
+                    comm_size=8,
+                    device_option=device_option,
+                    tmpdir=tmpdir)
+
     def _test_barrier(
         self,
         comm_rank=None,

caffe2/operators/communicator_op.cc

@@ -4,12 +4,16 @@
 namespace caffe2 {
 
 OPERATOR_SCHEMA(CreateCommonWorld)
-    .NumInputs(0, 1)
+    .NumInputs(0, 2)
     .NumOutputs(1)
     .SetDoc(R"DOC(
 Creates a common world for communication operators.
 )DOC")
     .Input(0, "kv_handler", "Key/value handler for rendezvous (optional).")
+    .Input(
+        1,
+        "existing_common_world",
+        "existing c-w that can be used to fork new one faster (optional).")
     .Output(0, "comm_world", "A common world for collective operations.")
     .Arg("size", "(int) size of the common world.")
     .Arg("rank", "(int) rank of this node in the common world.");

caffe2/python/data_parallel_model.py

@@ -544,7 +544,8 @@ def Synchronize(model, timeout_sec=30):
     barrier_instance += 1
     barrier_net = core.Net("sync_barrier_net_" + str(instance))
     comm_world = barrier_net.CreateCommonWorld(
-        model._rendezvous['kv_handler'],
+        [model._rendezvous['kv_handler']] +
+        _GetCommonWorldToFork(model.param_init_net),
         "sync_barrier_cw_" + str(instance),
         name="sync_barrier_cw_op_" + str(instance),
         size=model._rendezvous['num_shards'],
@@ -931,7 +932,8 @@ def get_control_and_context(self, control_output_blob):
         current_slot = self.counter % self.max_concurrent_context
         if len(self.common_worlds) < self.max_concurrent_context:
             common_world = self.param_init_net.CreateCommonWorld(
-                self.rendezvous['kv_handler'],
+                [self.rendezvous['kv_handler']] +
+                _GetCommonWorldToFork(self.param_init_net),
                 "{}_{}_cw".format(self.name, current_slot),
                 name="{}_{}_cw_op".format(self.name, current_slot),
                 size=self.rendezvous['num_shards'],
@@ -1366,6 +1368,17 @@ def OptimizeGradientMemory(model,
         )
 
 
+def _GetCommonWorldToFork(param_init_net):
+    '''
+    We can fork common worlds from existing ones. So inspect the param_init_net
+    for an already created commonworld
+    '''
+    for op in param_init_net.Proto().op:
+        if op.type == "CreateCommonWorld":
+            return [op.output[0]]
+    return []
+
+
 barrier_instance = 0