Merge remote-tracking branch 'origin/repo-refactor' into wmdi-compile…

…r-unity-dp

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

@@ -23,6 +23,9 @@ jobs:
       - name: Add helpers directory to path
         run: echo "${PWD}/.github/workflows/helpers" >> $GITHUB_PATH
 
+      - name: Free additional space on runner
+        run: free_space_on_runner.sh
+
       - name: Install nix
         uses: cachix/install-nix-action@v25
         with:

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

@@ -1,6 +1,8 @@
 #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/bert/bert.h"
+#include "models/candle_uno/candle_uno.h"
 #include "models/inception_v3/inception_v3.h"
 #include "models/split_test/split_test.h"
 #include "models/transformer/transformer.h"
@@ -63,6 +65,10 @@ tl::expected<ComputationGraph, std::string>
   } else if (model_name == "inception_v3") {
     return get_inception_v3_computation_graph(
         get_default_inception_v3_training_config());
+  } else if (model_name == "candle_uno") {
+    return get_candle_uno_computation_graph(get_default_candle_uno_config());
+  } else if (model_name == "bert") {
+    return get_bert_computation_graph(get_default_bert_config());
   } else if (model_name == "split_test") {
     int batch_size = 8;
     return get_split_test_computation_graph(batch_size);
@@ -135,8 +141,12 @@ int main(int argc, char **argv) {
                   "output a dot representation of model's computation graph "
                   "for preprocessed to help check series-parallel structure"});
 
-  std::vector<std::string> model_options = {
-      "transformer", "inception_v3", "split_test", "single_operator"};
+  std::vector<std::string> model_options = {"transformer",
+                                            "inception_v3",
+                                            "candle_uno",
+                                            "bert",
+                                            "split_test",
+                                            "single_operator"};
   CLIArgumentKey key_model_name = cli_add_positional_argument(
       cli,
       CLIPositionalArgumentSpec{

lib/compiler/test/src/compiler/series_parallel/get_computation_graph_series_parallel_decomposition.cc

@@ -1,4 +1,6 @@
 #include "compiler/series_parallel/get_computation_graph_series_parallel_decomposition.h"
+#include "models/bert/bert.h"
+#include "models/candle_uno/candle_uno.h"
 #include "models/inception_v3/inception_v3.h"
 #include "models/split_test/split_test.h"
 #include "models/transformer/transformer.h"
@@ -302,6 +304,26 @@ TEST_SUITE(FF_TEST_SUITE) {
 
         CHECK(sp_decomposition.has_value());
       }
+
+      SUBCASE("candle_uno") {
+        ComputationGraph cg =
+            get_candle_uno_computation_graph(get_default_candle_uno_config());
+
+        std::optional<SeriesParallelDecomposition> sp_decomposition =
+            get_computation_graph_series_parallel_decomposition(cg);
+
+        CHECK(sp_decomposition.has_value());
+      }
+
+      SUBCASE("bert") {
+        ComputationGraph cg =
+            get_bert_computation_graph(get_default_bert_config());
+
+        std::optional<SeriesParallelDecomposition> sp_decomposition =
+            get_computation_graph_series_parallel_decomposition(cg);
+
+        CHECK(sp_decomposition.has_value());
+      }
     }
   }
 
@@ -347,5 +369,29 @@ TEST_SUITE(FF_TEST_SUITE) {
       std::string result =
           render_preprocessed_computation_graph_for_sp_decomposition(cg);
     }
+
+    SUBCASE("inception_v3") {
+      ComputationGraph cg = get_inception_v3_computation_graph(
+          get_default_inception_v3_training_config());
+
+      std::string result =
+          render_preprocessed_computation_graph_for_sp_decomposition(cg);
+    }
+
+    SUBCASE("candle_uno") {
+      ComputationGraph cg =
+          get_candle_uno_computation_graph(get_default_candle_uno_config());
+
+      std::string result =
+          render_preprocessed_computation_graph_for_sp_decomposition(cg);
+    }
+
+    SUBCASE("bert") {
+      ComputationGraph cg =
+          get_bert_computation_graph(get_default_bert_config());
+
+      std::string result =
+          render_preprocessed_computation_graph_for_sp_decomposition(cg);
+    }
   }
 }

lib/models/include/models/bert/bert.h

@@ -0,0 +1,41 @@
+#ifndef _FLEXFLOW_LIB_MODELS_INCLUDE_MODELS_BERT_H
+#define _FLEXFLOW_LIB_MODELS_INCLUDE_MODELS_BERT_H
+
+#include "models/bert/bert_config.dtg.h"
+#include "pcg/computation_graph_builder.h"
+
+namespace FlexFlow {
+
+// Helper functions to construct the BERT model
+tensor_guid_t create_bert_feedforward_network(ComputationGraphBuilder &,
+                                              BertConfig const &,
+                                              tensor_guid_t const &);
+tensor_guid_t create_bert_encoder_layer(ComputationGraphBuilder &,
+                                        BertConfig const &,
+                                        tensor_guid_t const &);
+tensor_guid_t create_bert_encoder(ComputationGraphBuilder &,
+                                  BertConfig const &,
+                                  tensor_guid_t const &);
+
+/**
+ * @brief Get the base config of the BERT model.
+ *
+ * @details Refer to
+ * https://huggingface.co/docs/transformers/v4.18.0/en/model_doc/bert#transformers.BertConfig
+ * for default configs.
+ */
+BertConfig get_default_bert_config();
+
+/**
+ * @brief Get the BERT computation graph.
+ *
+ * @note This is a plain encoder-only model for pre-training.
+ *
+ * @param BertConfig The config of BERT model.
+ * @return ComputationGraph The computation graph of a BERT model.
+ */
+ComputationGraph get_bert_computation_graph(BertConfig const &);
+
+} // namespace FlexFlow
+
+#endif

lib/models/include/models/bert/bert_config.struct.toml

@@ -0,0 +1,71 @@
+namespace = "FlexFlow"
+name = "BertConfig"
+
+features = [
+  "eq",
+  "ord",
+  "hash",
+  "json",
+  "rapidcheck",
+  "fmt",
+]
+
+includes = [
+  "op-attrs/activation.dtg.h",
+]
+
+[[fields]]
+name = "vocab_size"
+type = "size_t"
+
+[[fields]]
+name = "hidden_size"
+type = "size_t"
+
+[[fields]]
+name = "num_encoder_layers"
+type = "size_t"
+
+[[fields]]
+name = "num_heads"
+type = "size_t"
+
+[[fields]]
+name = "dim_feedforward"
+type = "size_t"
+
+[[fields]]
+name = "hidden_act"
+type = "::FlexFlow::Activation"
+
+[[fields]]
+name = "hidden_dropout_prob"
+type = "float"
+
+[[fields]]
+name = "attention_probs_dropout_prob"
+type = "float"
+
+[[fields]]
+name = "initializer_range"
+type = "float"
+
+[[fields]]
+name = "layer_norm_eps"
+type = "float"
+
+[[fields]]
+name = "position_embedding_type"
+type = "std::string"
+
+[[fields]]
+name = "classifier_dropout"
+type = "float"
+
+[[fields]]
+name = "sequence_length"
+type = "size_t"
+
+[[fields]]
+name = "batch_size"
+type = "size_t"

lib/models/include/models/candle_uno/candle_uno.h

@@ -0,0 +1,41 @@
+#ifndef _FLEXFLOW_LIB_MODELS_INCLUDE_MODELS_CANDLE_UNO_H
+#define _FLEXFLOW_LIB_MODELS_INCLUDE_MODELS_CANDLE_UNO_H
+
+#include "candle_uno_config.dtg.h"
+#include "pcg/computation_graph_builder.h"
+#include <map>
+#include <string>
+#include <vector>
+
+namespace FlexFlow {
+
+// Helper functions to construct the Candle Uno model
+tensor_guid_t create_candle_uno_feature_model(ComputationGraphBuilder &,
+                                              CandleUnoConfig const &,
+                                              tensor_guid_t const &);
+
+/**
+ * @brief Get the default configs of Candle Uno model.
+ *
+ * @details The default configs come from the dataset used by the original
+ * model:
+ * https://github.com/ECP-CANDLE/Benchmarks/tree/f6a3da8818308c9edcd9fedbcf831dd5968efcdd/Pilot1/Uno
+ */
+CandleUnoConfig get_default_candle_uno_config();
+
+/**
+ * @brief Get the Candle Uno computation graph.
+ *
+ * @details CandleUnoConfig.feature_shapes is a map from feature name to the
+ * number of channels for the feature, and CandleUnoConfig.input_features is a
+ * map from specific data identifier in the dataset to the feature name used in
+ * this model.
+ *
+ * @param CandleUnoConfig The config of the Candle Uno model.
+ * @return ComputationGraph The PCG of a Transformer model.
+ */
+ComputationGraph get_candle_uno_computation_graph(CandleUnoConfig const &);
+
+} // namespace FlexFlow
+
+#endif

lib/models/include/models/candle_uno/candle_uno_config.struct.toml

@@ -0,0 +1,52 @@
+namespace = "FlexFlow"
+name = "CandleUnoConfig"
+
+features = [
+  "eq",
+  "ord",
+  "hash",
+  "json",
+  "rapidcheck",
+  "fmt",
+]
+
+includes = [
+  "<vector>",
+  "<map>",
+  "<string>",
+]
+
+src_includes = [
+  "utils/fmt/vector.h",
+  "utils/fmt/map.h",
+  "utils/hash/vector.h",
+  "utils/hash/map.h",
+]
+
+[[fields]]
+name = "batch_size"
+type = "size_t"
+
+[[fields]]
+name = "dense_layers"
+type = "std::vector<int>"
+
+[[fields]]
+name = "dense_feature_layers"
+type = "std::vector<int>"
+
+[[fields]]
+name = "feature_shapes"
+type = "std::map<std::string, int>"
+
+[[fields]]
+name = "input_features"
+type = "std::map<std::string, std::string>"
+
+[[fields]]
+name = "dropout"
+type = "float"
+
+[[fields]]
+name = "residual"
+type = "bool"