Langjian/distributed json

CMakeLists.txt

@@ -226,12 +226,12 @@ message(STATUS "NGRAPH_PLAIDML_ENABLE:      ${NGRAPH_PLAIDML_ENABLE}")
 message(STATUS "NGRAPH_TARGET_ARCH:         ${NGRAPH_TARGET_ARCH}")
 message(STATUS "NGRAPH_TUNE_ARCH:           ${NGRAPH_TUNE_ARCH}")
 
-if(NGRAPH_DISTRIBUTED_ENABLE)
-    find_package(MPI REQUIRED)
-    add_definitions(-DNGRAPH_DISTRIBUTED)
-    include_directories(SYSTEM ${MPI_C_INCLUDE_PATH} ${MPI_CXX_INCLUDE_PATH})
-    link_directories(${MPI_C_LIBRARIES} ${MPI_CXX_LIBRARIES})
-endif()
+#if(NGRAPH_DISTRIBUTED_ENABLE)
+#    find_package(MPI REQUIRED)
+#    add_definitions(-DNGRAPH_DISTRIBUTED)
+#    include_directories(SYSTEM ${MPI_C_INCLUDE_PATH} ${MPI_CXX_INCLUDE_PATH})
+#    link_directories(${MPI_C_LIBRARIES} ${MPI_CXX_LIBRARIES})
+#endif()
 
 # Find and build ngraph - if not using pre-built one
 if (NOT USE_PRE_BUILT_NGRAPH)

build_ngtf.py

@@ -400,6 +400,11 @@ def main():
         action="store_true"
         )
 
+    parser.add_argument(
+        '--distributed_build',
+        help="Builds a distributed version of the nGraph components\n",
+        action="store_true")
+
     arguments = parser.parse_args()
 
     if (arguments.debug_build):
@@ -484,7 +489,6 @@ def main():
 
         ngraph_cmake_flags = [
             "-DNGRAPH_INSTALL_PREFIX=" + artifacts_location,
-            "-DNGRAPH_DISTRIBUTED_ENABLE=FALSE",
             "-DNGRAPH_USE_CXX_ABI=" + cxx_abi,
             "-DNGRAPH_UNIT_TEST_ENABLE=NO",
             "-DNGRAPH_DEX_ONLY=TRUE",
@@ -502,6 +506,11 @@ def main():
         if (arguments.debug_build):
             ngraph_cmake_flags.extend(["-DCMAKE_BUILD_TYPE=Debug"])
 
+        if (arguments.distributed_build): 
+            ngraph_cmake_flags.extend(["-DNGRAPH_DISTRIBUTED_ENABLE=TRUE"])
+        else:
+            ngraph_cmake_flags.extend(["-DNGRAPH_DISTRIBUTED_ENABLE=FALSE"])
+
         build_ngraph("./ngraph", ngraph_cmake_flags, verbosity)
 
     # Next build CMAKE options for the bridge
@@ -512,13 +521,17 @@ def main():
         "-DNGRAPH_TUNE_ARCH=" + target_arch,
         "-DNGRAPH_ARTIFACTS_DIR=" + artifacts_location,
         "-DUNIT_TEST_ENABLE=ON",
-        "-DNGRAPH_DISTRIBUTED_ENABLE=TRUE",
         "-DTF_SRC_DIR=" + tf_src_dir, "-DUNIT_TEST_TF_CC_DIR=" + os.path.join(
             artifacts_location, "tensorflow")
     ]
     if (arguments.debug_build):
         ngraph_tf_cmake_flags.extend(["-DCMAKE_BUILD_TYPE=Debug"])
 
+    if (arguments.distributed_build):
+        ngraph_tf_cmake_flags.extend(["-DNGRAPH_DISTRIBUTED_ENABLE=TRUE"])
+    else:
+        ngraph_tf_cmake_flags.extend(["-DNGRAPH_DISTRIBUTED_ENABLE=FALSE"])
+
     # Now build the bridge
     ng_tf_whl = build_ngraph_tf(artifacts_location, "../", venv_dir,
                                 ngraph_tf_cmake_flags, verbosity)

examples/mnist/mnist_softmax_distributed.py

@@ -16,6 +16,11 @@
 
 See extensive documentation at
 https://www.tensorflow.org/get_started/mnist/beginners
+Add distributed fetaure with horovod
+1. hvd.init()
+2. Add distributed wrapper from hvd.DistributedOptimizer
+3. Broadcast the variables from root rank to the rest processors: hvd.BroadcastGlobalVariablesHook(0)
+4. Print the output for root rank only
 """
 from __future__ import absolute_import
 from __future__ import division

src/ngraph_encapsulate_op.cc

@@ -37,7 +37,7 @@
 #include "ngraph/runtime/interpreter/int_backend.hpp"
 
 #ifdef NGRAPH_DISTRIBUTED
-#include <mpi.h>
+#include "ngraph/distributed.hpp"
 #endif
 
 using namespace std;
@@ -267,13 +267,9 @@ class NGraphEncapsulateOp : public OpKernel {
         std::string file_name =
             "tf_function_" + ctx->op_kernel().name() + ".json";
 #ifdef NGRAPH_DISTRIBUTED
-        int flag = 0;
-        MPI_Initialized(&flag);
-        if (!flag) {
-          MPI_Init(NULL, NULL);
-        }
+        ngraph::Distributed dist;
         int Rank_ID;
-        MPI_Comm_rank(MPI_COMM_WORLD, &Rank_ID);
+        Rank_ID = dist->get_rank();
         NgraphSerialize("tf_function_" + ctx->op_kernel().name() + "_" +
                             to_string(Rank_ID) + ".json",
                         ng_function);

src/ngraph_rewrite_pass.cc

@@ -30,7 +30,7 @@
 #include <iomanip>
 
 #ifdef NGRAPH_DISTRIBUTED
-#include <mpi.h>
+#include "ngraph/distributed.hpp"
 #endif
 
 using namespace std;
@@ -107,13 +107,8 @@ class NGraphRewritePass : public GraphOptimizationPass {
     std::stringstream ss;
     ss << kind << "_" << std::setfill('0') << std::setw(4) << idx;
 #ifdef NGRAPH_DISTRIBUTED
-    int flag = 0;
-    MPI_Initialized(&flag);
-    if (!flag) {
-      MPI_Init(NULL, NULL);
-    }
-    int Rank_ID;
-    MPI_Comm_rank(MPI_COMM_WORLD, &Rank_ID);
+    ngraph::Distributed dist;
+    int Rank_ID = dist->get_rank();
     ss << "_" << std::setfill('0') << std::setw(4) << Rank_ID;
 #endif
     return ss.str();
@@ -124,13 +119,8 @@ class NGraphRewritePass : public GraphOptimizationPass {
     ss << GraphFilenamePrefix(kind, idx) << "_" << std::setfill('0')
        << std::setw(4) << sub_idx;
 #ifdef NGRAPH_DISTRIBUTED
-    int flag = 0;
-    MPI_Initialized(&flag);
-    if (!flag) {
-      MPI_Init(NULL, NULL);
-    }
-    int Rank_ID;
-    MPI_Comm_rank(MPI_COMM_WORLD, &Rank_ID);
+    ngraph::Distributed dist;
+    int Rank_ID = dist->get_rank();
     ss << "_" << std::setfill('0') << std::setw(4) << Rank_ID;
 #endif
     return ss.str();