diff --git a/lib/models/include/models/dlrm.h b/lib/models/include/models/dlrm.h deleted file mode 100644 index 5da3c615e..000000000 --- a/lib/models/include/models/dlrm.h +++ /dev/null @@ -1,23 +0,0 @@ -#ifndef _FLEXFLOW_LIB_MODELS_INCLUDE_MODELS_DLRM_H -#define _FLEXFLOW_LIB_MODELS_INCLUDE_MODELS_DLRM_H - -#include "models/dlrm_config.dtg.h" -#include "pcg/computation_graph_builder.h" - -namespace FlexFlow { - -// Helper functions to construct the DLRM model - -DLRMConfig get_default_dlrm_config(); - -/** - * @brief Get the DLRM computation graph. - * - * @param DLRMConfig The config of DLRM model. - * @return ComputationGraph The PCG of a DLRM model. - */ -ComputationGraph get_dlrm_computation_graph(DLRMConfig const &); - -} // namespace FlexFlow - -#endif diff --git a/lib/models/include/models/dlrm/dlrm.h b/lib/models/include/models/dlrm/dlrm.h new file mode 100644 index 000000000..6f8d3a1f0 --- /dev/null +++ b/lib/models/include/models/dlrm/dlrm.h @@ -0,0 +1,40 @@ +#ifndef _FLEXFLOW_LIB_MODELS_INCLUDE_MODELS_DLRM_H +#define _FLEXFLOW_LIB_MODELS_INCLUDE_MODELS_DLRM_H + +#include "models/dlrm/dlrm_config.dtg.h" +#include "pcg/computation_graph_builder.h" + +namespace FlexFlow { + +// Helper functions to construct the DLRM model + +DLRMConfig get_default_dlrm_config(); + +tensor_guid_t create_dlrm_mlp(ComputationGraphBuilder &, + DLRMConfig const &, + tensor_guid_t const &, + std::vector const &, + int const &); + +tensor_guid_t create_dlrm_emb(ComputationGraphBuilder &, + DLRMConfig const &, + tensor_guid_t const &, + int const &, + int const &); + +tensor_guid_t create_dlrm_interact_features(ComputationGraphBuilder &, + DLRMConfig const &, + tensor_guid_t const &, + std::vector const &); + +/** + * @brief Get the DLRM computation graph. + * + * @param DLRMConfig The config of DLRM model. + * @return ComputationGraph The computation graph of a DLRM model. + */ +ComputationGraph get_dlrm_computation_graph(DLRMConfig const &); + +} // namespace FlexFlow + +#endif diff --git a/lib/models/include/models/dlrm/dlrm_config.struct.toml b/lib/models/include/models/dlrm/dlrm_config.struct.toml new file mode 100644 index 000000000..b60533c69 --- /dev/null +++ b/lib/models/include/models/dlrm/dlrm_config.struct.toml @@ -0,0 +1,69 @@ +namespace = "FlexFlow" +name = "DLRMConfig" + +features = [ + "eq", + "ord", + "hash", + "json", + "rapidcheck", + "fmt", +] + +includes = [ + "", + "", +] + +src_includes = [ + "utils/fmt/vector.h", + "utils/hash/vector.h", +] + +[[fields]] +name = "sparse_feature_size" +type = "int" + +[[fields]] +name = "sigmoid_bot" +type = "int" + +[[fields]] +name = "sigmoid_top" +type = "int" + +[[fields]] +name = "embedding_bag_size" +type = "size_t" + +[[fields]] +name = "loss_threshold" +type = "float" + +[[fields]] +name = "embedding_size" +type = "std::vector" + +[[fields]] +name = "mlp_bot" +type = "std::vector" + +[[fields]] +name = "mlp_top" +type = "std::vector" + +[[fields]] +name = "arch_interaction_op" +type = "std::string" + +[[fields]] +name = "dataset_path" +type = "std::string" + +[[fields]] +name = "data_size" +type = "int" + +[[fields]] +name = "batch_size" +type = "size_t" diff --git a/lib/models/include/models/dlrm_config.struct.toml b/lib/models/include/models/dlrm_config.struct.toml deleted file mode 100644 index d7e875b13..000000000 --- a/lib/models/include/models/dlrm_config.struct.toml +++ /dev/null @@ -1,15 +0,0 @@ -namespace = "FlexFlow" -name = "DLRMConfig" - -features = [ - "eq", - "ord", - "hash", - "json", - "rapidcheck", - "fmt", -] - -[[fields]] -name = "batch_size" -type = "size_t" diff --git a/lib/models/src/models/dlrm.cc b/lib/models/src/models/dlrm.cc deleted file mode 100644 index 1cf601408..000000000 --- a/lib/models/src/models/dlrm.cc +++ /dev/null @@ -1,15 +0,0 @@ -#include "models/dlrm.h" -#include "pcg/computation_graph.h" - -namespace FlexFlow { - -DLRMConfig get_default_dlrm_config() { - return DLRMConfig(64); -} - -ComputationGraph get_dlrm_computation_graph(DLRMConfig const &config) { - ComputationGraphBuilder cgb; - return cgb.computation_graph; -} - -} // namespace FlexFlow diff --git a/lib/models/src/models/dlrm/dlrm.cc b/lib/models/src/models/dlrm/dlrm.cc new file mode 100644 index 000000000..6b699993f --- /dev/null +++ b/lib/models/src/models/dlrm/dlrm.cc @@ -0,0 +1,157 @@ +#include "models/dlrm/dlrm.h" +#include "pcg/computation_graph.h" +#include "utils/containers/concat_vectors.h" + +namespace FlexFlow { + +DLRMConfig get_default_dlrm_config() { + DLRMConfig config{/*sparse_feature_size=*/64, + /*sigmoid_bot=*/-1, + /*sigmoid_top=*/-1, + /*embedding_bag_size=*/1, + /*loss_threshold=*/0, + /*embedding_size=*/std::vector{}, + /*mlp_bot=*/std::vector{}, + /*mlp_top=*/std::vector{}, + /*arch_interaction_op=*/"cat", + /*dataset_path=*/"", + /*data_size=*/-1, + /*batch_size=*/64}; + + config.embedding_size.emplace_back(1000000); + config.embedding_size.emplace_back(1000000); + config.embedding_size.emplace_back(1000000); + config.embedding_size.emplace_back(1000000); + + config.mlp_bot.emplace_back(4); + config.mlp_bot.emplace_back(64); + config.mlp_bot.emplace_back(64); + + config.mlp_top.emplace_back(64); + config.mlp_top.emplace_back(64); + config.mlp_top.emplace_back(2); + + return config; +} + +tensor_guid_t create_dlrm_mlp(ComputationGraphBuilder &cgb, + DLRMConfig const &config, + tensor_guid_t const &input, + std::vector const &mlp_layers, + int const &sigmoid_layer) { + tensor_guid_t t = input; + for (size_t i = 0; i < mlp_layers.size() - 1; i++) { + float std_dev = sqrt(2.0f / (mlp_layers[i + 1] + mlp_layers[i])); + InitializerAttrs projection_initializer = + InitializerAttrs{NormInitializerAttrs{ + /*seed=*/std::rand(), + /*mean=*/0, + /*stddev=*/std_dev, + }}; + + std_dev = sqrt(2.0f / mlp_layers[i + 1]); + InitializerAttrs bias_initializer = InitializerAttrs{NormInitializerAttrs{ + /*seed=*/std::rand(), + /*mean=*/0, + /*stddev=*/std_dev, + }}; + + Activation activation = + (i == sigmoid_layer) ? Activation::SIGMOID : Activation::RELU; + + t = cgb.dense(/*input=*/t, + /*outDim=*/mlp_layers[i + 1], + /*activation=*/activation, + /*use_bias=*/true, + /*data_type=*/DataType::FLOAT, + /*projection_initializer=*/projection_initializer, + /*bias_initializer=*/bias_initializer); + } + return t; +} + +tensor_guid_t create_dlrm_emb(ComputationGraphBuilder &cgb, + DLRMConfig const &config, + tensor_guid_t const &input, + int const &input_dim, + int const &output_dim) { + float range = sqrt(1.0f / input_dim); + InitializerAttrs embed_initializer = InitializerAttrs{UniformInitializerAttrs{ + /*seed=*/std::rand(), + /*min_val=*/-range, + /*max_val=*/range, + }}; + + tensor_guid_t t = cgb.embedding(input, + /*num_entries=*/input_dim, + /*outDim=*/output_dim, + /*aggr=*/AggregateOp::SUM, + /*dtype=*/DataType::HALF, + /*kernel_initializer=*/embed_initializer); + return cgb.cast(t, DataType::FLOAT); +} + +tensor_guid_t create_dlrm_interact_features( + ComputationGraphBuilder &cgb, + DLRMConfig const &config, + tensor_guid_t const &bottom_mlp_output, + std::vector const &emb_outputs) { + if (config.arch_interaction_op != "cat") { + throw mk_runtime_error(fmt::format( + "Currently only arch_interaction_op=cat is supported, but found " + "arch_interaction_op={}. If you need support for additional " + "arch_interaction_op value, please create an issue.", + config.arch_interaction_op)); + } + + return cgb.concat( + /*tensors=*/concat_vectors({bottom_mlp_output}, emb_outputs), + /*axis=*/1); +} + +ComputationGraph get_dlrm_computation_graph(DLRMConfig const &config) { + ComputationGraphBuilder cgb; + + auto create_input_tensor = [&](FFOrdered const &dims, + DataType const &data_type) -> tensor_guid_t { + TensorShape input_shape = TensorShape{ + TensorDims{dims}, + data_type, + }; + return cgb.create_input(input_shape, CreateGrad::YES); + }; + + // Create input tensors + std::vector sparse_inputs( + config.embedding_size.size(), + create_input_tensor({config.batch_size, config.embedding_bag_size}, + DataType::INT64)); + + tensor_guid_t dense_input = create_input_tensor( + {config.batch_size, config.mlp_bot.front()}, DataType::FLOAT); + + // Construct the model + tensor_guid_t bottom_mlp_output = create_dlrm_mlp( + cgb, config, dense_input, config.mlp_bot, config.sigmoid_bot); + + std::vector emb_outputs; + for (size_t i = 0; i < config.embedding_size.size(); i++) { + int input_dim = config.embedding_size[i]; + int output_dim = config.sparse_feature_size; + emb_outputs.emplace_back( + create_dlrm_emb(cgb, config, sparse_inputs[i], input_dim, output_dim)); + } + + tensor_guid_t interacted_features = create_dlrm_interact_features( + cgb, config, bottom_mlp_output, emb_outputs); + + tensor_guid_t output = create_dlrm_mlp(cgb, + config, + interacted_features, + config.mlp_top, + config.mlp_top.size() - 2); + + return cgb.computation_graph; +} + +} // namespace FlexFlow diff --git a/lib/models/test/src/models/dlrm.cc b/lib/models/test/src/models/dlrm/dlrm.cc similarity index 87% rename from lib/models/test/src/models/dlrm.cc rename to lib/models/test/src/models/dlrm/dlrm.cc index 94986b21b..ce199c843 100644 --- a/lib/models/test/src/models/dlrm.cc +++ b/lib/models/test/src/models/dlrm/dlrm.cc @@ -1,4 +1,4 @@ -#include "models/dlrm.h" +#include "models/dlrm/dlrm.h" #include "pcg/computation_graph.h" #include @@ -12,7 +12,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("num layers") { int result_num_layers = get_layers(result).size(); - int correct_num_layers = 0; + int correct_num_layers = 23; CHECK(result_num_layers == correct_num_layers); } } diff --git a/lib/pcg/src/pcg/computation_graph_builder.cc b/lib/pcg/src/pcg/computation_graph_builder.cc index 4a565476b..fe001a4ce 100644 --- a/lib/pcg/src/pcg/computation_graph_builder.cc +++ b/lib/pcg/src/pcg/computation_graph_builder.cc @@ -164,7 +164,8 @@ tensor_guid_t ComputationGraphBuilder::cast(tensor_guid_t const &input, DataType dtype, std::optional const &name) { - NOT_IMPLEMENTED() + // NOT_IMPLEMENTED() + return input; } static std::string get_default_name(OperatorType op_type) {