Add tool for exporting and visualizing model architectures and SP decompositions

.github/workflows/per-lib-check.yml

@@ -62,71 +62,79 @@ jobs:
 
       - name: Build utils
         run: |
-          build_libs.sh utils
+          build_target.sh utils
 
       - name: Build op-attrs
         run: |
-          build_libs.sh op-attrs
+          build_target.sh op-attrs
 
       - name: Build pcg
         run: |
-          build_libs.sh pcg
+          build_target.sh pcg
 
       - name: Build kernels
         run: |
-          build_libs.sh kernels
+          build_target.sh kernels
 
       - name: Build substitutions
         run: |
-          build_libs.sh substitutions
+          build_target.sh substitutions
 
       - name: Build compiler
         run: |
-          build_libs.sh compiler
+          build_target.sh compiler
 
       - name: Build substitution-generator
         run: |
-          build_libs.sh substitution-generator
+          build_target.sh substitution-generator
 
       - name: Build local-execution
         run: |
-          build_libs.sh local-execution
+          build_target.sh local-execution
 
       - name: Build models
         run: |
-          build_libs.sh models
+          build_target.sh models
+
+      - name: Build substitution-to-dot
+        run: |
+          build_target.sh substitution-to-dot
+
+      - name: Build export-model-arch
+        run: |
+          build_target.sh export-model-arch
 
       - name: Test utils
         run: |
-          test_libs.sh utils
+          test_target.sh utils
 
       - name: Test op-attrs
         run: |
-          test_libs.sh op-attrs
+          test_target.sh op-attrs
 
       - name: Test pcg
         run: |
-          test_libs.sh pcg
+          test_target.sh pcg
 
       - name: Test substitutions
         run: |
-          test_libs.sh substitutions
+          test_target.sh substitutions
 
       # - name: Test compiler
       #   run: |
-      #     test_libs.sh compiler
+      #     test_target.sh compiler
 
       - name: Test substitution-generator
         run: |
-          test_libs.sh substitution-generator
+          test_target.sh substitution-generator
 
       - name: Test local-execution
         run: |
-          test_libs.sh local-execution
+          test_target.sh local-execution
 
       - name: Test models
         run: |
-          test_libs.sh models
+          test_target.sh models
 
       - name: Generate code coverage
         run: |

.proj.toml

@@ -13,6 +13,8 @@ build_targets = [
   "substitution-generator",
   "local-execution",
   "models",
+  "export-model-arch",
+  "substitution-to-dot",
 ]
 
 test_targets = [

CMakeLists.txt

@@ -74,6 +74,7 @@ option(FF_BUILD_UNIT_TESTS "build non-operator unit tests" OFF)
 option(FF_BUILD_SUBSTITUTION_TOOL "build substitution conversion tool" OFF)
 option(FF_BUILD_VISUALIZATION_TOOL "build substitution visualization tool" ON)
 option(FF_BUILD_ARG_PARSER "build command line argument parser" OFF)
+option(FF_BUILD_BIN_EXPORT_MODEL_ARCH "build export-model-arch utility" ON)
 
 set(FF_CUDA_ARCH "autodetect" CACHE STRING "Target CUDA Arch")
 if (FF_CUDA_ARCH STREQUAL "")

bin/CMakeLists.txt

@@ -7,9 +7,13 @@ if(FF_BUILD_SUBSTITUTION_TOOL)
 endif()
 
 if(FF_BUILD_VISUALIZATION_TOOL)
-  add_subdirectory(substitutions-to-dot)
+  add_subdirectory(substitution-to-dot)
 endif()
 
 if(FF_BUILD_ARG_PARSER)
   add_subdirectory(arg_parser)
 endif()
+
+if(FF_BUILD_BIN_EXPORT_MODEL_ARCH)
+  add_subdirectory(export-model-arch)
+endif()

bin/export-model-arch/CMakeLists.txt

@@ -0,0 +1,12 @@
+ff_add_executable(
+  NAME
+    export-model-arch
+  SRC_PATTERNS
+    src/*.cc
+  PRIVATE_INCLUDE
+    include/
+  DEPS
+    utils
+    models
+    compiler
+)

bin/export-model-arch/include/export_model_arch/json_sp_model_export.struct.toml

@@ -0,0 +1,27 @@
+namespace = "FlexFlow"
+name = "JsonSPModelExport"
+features = [
+  "eq",
+  "hash",
+  "json",
+  "fmt",
+]
+
+includes = [
+  "pcg/file_format/v1/v1_computation_graph.dtg.h",
+  "utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree.h",
+]
+
+src_includes = [
+  "utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/hash.h",
+  "utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/fmt.h",
+  "utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/json.h",
+]
+
+[[fields]]
+name = "sp_decomposition"
+type = "::FlexFlow::GenericBinarySPDecompositionTree<int>"
+
+[[fields]]
+name = "computation_graph"
+type = "::FlexFlow::V1ComputationGraph"

bin/export-model-arch/src/export_model_arch.cc

@@ -0,0 +1,208 @@
+#include "compiler/series_parallel/computation_graph_binary_sp_decomposition.h"
+#include "compiler/series_parallel/get_computation_graph_series_parallel_decomposition.h"
+#include "export_model_arch/json_sp_model_export.dtg.h"
+#include "models/split_test/split_test.h"
+#include "models/transformer/transformer.h"
+#include "op-attrs/computation_graph_op_attrs.h"
+#include "pcg/computation_graph.h"
+#include "pcg/file_format/v1/v1_computation_graph.h"
+#include "utils/cli/cli_get_help_message.h"
+#include "utils/cli/cli_parse.h"
+#include "utils/cli/cli_parse_result.h"
+#include "utils/cli/cli_spec.h"
+#include "utils/graph/series_parallel/binary_sp_decomposition_tree/generic_binary_sp_decomposition_tree/transform.h"
+#include "utils/graph/series_parallel/binary_sp_decomposition_tree/right_associative_binary_sp_tree_from_nary.h"
+#include "utils/graph/series_parallel/get_series_parallel_decomposition.h"
+
+using namespace ::FlexFlow;
+
+ComputationGraph get_single_operator_computation_graph() {
+  ComputationGraphBuilder b;
+
+  size_t batch_size = 8;
+  size_t in_channels = 16;
+  size_t out_channels = 12;
+  TensorShape input_shape = TensorShape{
+      TensorDims{FFOrdered<size_t>{
+          batch_size,
+          in_channels,
+          out_channels,
+      }},
+      DataType::FLOAT,
+  };
+
+  tensor_guid_t input = b.create_input(input_shape, CreateGrad::YES);
+
+  InitializerAttrs kernel_initializer =
+      InitializerAttrs{GlorotUniformAttrs{/*seed=*/12}};
+  InitializerAttrs bias_initializer = InitializerAttrs{ZeroInitializerAttrs{}};
+  tensor_guid_t output = b.dense(input,
+                                 in_channels,
+                                 Activation::RELU,
+                                 /*use_bias=*/true,
+                                 DataType::FLOAT,
+                                 kernel_initializer,
+                                 bias_initializer,
+                                 "my_example_operator");
+
+  return b.computation_graph;
+}
+
+ComputationGraph get_default_transformer_computation_graph() {
+  TransformerConfig config = get_default_transformer_config();
+  ComputationGraph cg = get_transformer_computation_graph(config);
+
+  return cg;
+}
+
+tl::expected<ComputationGraph, std::string>
+    get_model_computation_graph(std::string const &model_name) {
+  if (model_name == "transformer") {
+    return get_default_transformer_computation_graph();
+  } else if (model_name == "split_test") {
+    int batch_size = 8;
+    return get_split_test_computation_graph(batch_size);
+  } else if (model_name == "single_operator") {
+    return get_single_operator_computation_graph();
+  } else {
+    return tl::unexpected(fmt::format("Unknown model name: {}", model_name));
+  }
+}
+
+tl::expected<JsonSPModelExport, std::string>
+    get_sp_model_export(std::string const &model_name) {
+  ComputationGraph computation_graph = ({
+    tl::expected<ComputationGraph, std::string> result =
+        get_model_computation_graph(model_name);
+    if (!result.has_value()) {
+      return tl::unexpected(result.error());
+    }
+    result.value();
+  });
+
+  ComputationGraphBinarySPDecomposition sp_decomposition = ({
+    std::optional<ComputationGraphBinarySPDecomposition> result =
+        get_computation_graph_right_assoc_binary_sp_decomposition(
+            computation_graph);
+    if (!result.has_value()) {
+      return tl::unexpected("Failed to generate series-parallel decomposition "
+                            "of computation graph.");
+    }
+    result.value();
+  });
+
+  std::pair<V1ComputationGraph, bidict<int, layer_guid_t>> v1_result =
+      to_v1_including_node_numbering(computation_graph);
+  V1ComputationGraph v1_cg = v1_result.first;
+  bidict<int, layer_guid_t> layer_numbering = v1_result.second;
+  GenericBinarySPDecompositionTree<int> v1_sp_decomposition =
+      transform(sp_decomposition.raw_tree,
+                [&](layer_guid_t const &l) { return layer_numbering.at_r(l); });
+
+  return JsonSPModelExport{
+      v1_sp_decomposition,
+      v1_cg,
+  };
+}
+
+int main(int argc, char **argv) {
+  CLISpec cli = empty_cli_spec();
+
+  CLIArgumentKey arg_key_help = cli_add_help_flag(cli);
+
+  CLIArgumentKey key_sp_decomposition =
+      cli_add_flag(cli,
+                   CLIFlagSpec{"sp-decomposition",
+                               std::nullopt,
+                               "also output a series parallel decomposition of "
+                               "the model's computation graph"});
+
+  CLIArgumentKey key_dot = cli_add_flag(
+      cli,
+      CLIFlagSpec{
+          "dot",
+          std::nullopt,
+          "output a dot representation of the model's computation graph"});
+
+  CLIArgumentKey key_preprocessed_dot = cli_add_flag(
+      cli,
+      CLIFlagSpec{"preprocessed-dot",
+                  std::nullopt,
+                  "output a dot representation of model's computation graph "
+                  "for preprocessed to help check series-parallel structure"});
+
+  std::vector<std::string> model_options = {
+      "transformer", "split_test", "single_operator"};
+  CLIArgumentKey key_model_name = cli_add_positional_argument(
+      cli,
+      CLIPositionalArgumentSpec{
+          "model", model_options, "name of the model to export"});
+
+  assert(argc >= 1);
+  std::string prog_name = argv[0];
+
+  CLIParseResult parsed = ({
+    tl::expected<CLIParseResult, std::string> result =
+        cli_parse(cli, argc, argv);
+    if (!result.has_value()) {
+      std::string error_msg = result.error();
+      std::cerr << cli_get_help_message(prog_name, cli);
+      std::cerr << std::endl;
+      std::cerr << "error: " << error_msg << std::endl;
+      return 1;
+    }
+
+    result.value();
+  });
+
+  bool help = cli_get_flag(parsed, arg_key_help);
+  if (help) {
+    std::cerr << cli_get_help_message(prog_name, cli);
+    return 1;
+  }
+
+  std::string model_name = cli_get_argument(parsed, key_model_name);
+  bool sp_decompositition = cli_get_flag(parsed, key_sp_decomposition);
+  bool dot = cli_get_flag(parsed, key_dot);
+  bool preprocessed_dot = cli_get_flag(parsed, key_preprocessed_dot);
+
+  auto handle_error = [](auto const &result) {
+    if (!result.has_value()) {
+      std::cerr << "error: " << result.error() << std::endl;
+      exit(1);
+    }
+
+    return result.value();
+  };
+
+  if (dot) {
+    ComputationGraph cg = handle_error(get_model_computation_graph(model_name));
+
+    std::cout << as_dot(cg) << std::endl;
+    return 0;
+  }
+
+  if (preprocessed_dot) {
+    ComputationGraph cg = handle_error(get_model_computation_graph(model_name));
+    std::string rendered =
+        render_preprocessed_computation_graph_for_sp_decomposition(cg);
+
+    std::cout << rendered << std::endl;
+    return 0;
+  }
+
+  nlohmann::json json_output;
+  if (sp_decompositition) {
+    JsonSPModelExport model_export =
+        handle_error(get_sp_model_export(model_name));
+
+    json_output = model_export;
+  } else {
+    ComputationGraph cg = handle_error(get_model_computation_graph(model_name));
+
+    json_output = to_v1(cg);
+  }
+  std::cout << json_output.dump(2) << std::endl;
+
+  return 0;
+}