From ae32dfad4cdee5f76a0b09a119c25f5b6b9ef71d Mon Sep 17 00:00:00 2001 From: easyeasydev Date: Wed, 30 Oct 2024 21:40:06 -0400 Subject: [PATCH 1/5] Fix typo in Candle Uno --- lib/models/include/models/candle_uno/candle_uno.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/models/include/models/candle_uno/candle_uno.h b/lib/models/include/models/candle_uno/candle_uno.h index a2d21f2830..e3093972d7 100644 --- a/lib/models/include/models/candle_uno/candle_uno.h +++ b/lib/models/include/models/candle_uno/candle_uno.h @@ -32,7 +32,7 @@ CandleUnoConfig get_default_candle_uno_config(); * this model. * * @param CandleUnoConfig The config of the Candle Uno model. - * @return ComputationGraph The PCG of a Transformer model. + * @return ComputationGraph The computation graph of a Candle Uno model. */ ComputationGraph get_candle_uno_computation_graph(CandleUnoConfig const &); From 49723e11d2dbccad9251e23a203b3b3bf1ed4145 Mon Sep 17 00:00:00 2001 From: easyeasydev Date: Wed, 30 Oct 2024 21:40:48 -0400 Subject: [PATCH 2/5] Add simvp placeholder --- lib/models/include/models/simvp/simvp.h | 26 +++++++++++++++++ .../models/simvp/simvp_config.struct.toml | 28 +++++++++++++++++++ lib/models/src/models/simvp/simvp.cc | 14 ++++++++++ lib/models/test/src/models/simvp/simvp.cc | 19 +++++++++++++ 4 files changed, 87 insertions(+) create mode 100644 lib/models/include/models/simvp/simvp.h create mode 100644 lib/models/include/models/simvp/simvp_config.struct.toml create mode 100644 lib/models/src/models/simvp/simvp.cc create mode 100644 lib/models/test/src/models/simvp/simvp.cc diff --git a/lib/models/include/models/simvp/simvp.h b/lib/models/include/models/simvp/simvp.h new file mode 100644 index 0000000000..4eed320a7b --- /dev/null +++ b/lib/models/include/models/simvp/simvp.h @@ -0,0 +1,26 @@ +#ifndef _FLEXFLOW_LIB_MODELS_INCLUDE_MODELS_SIMVP_H +#define _FLEXFLOW_LIB_MODELS_INCLUDE_MODELS_SIMVP_H + +#include "pcg/computation_graph_builder.h" +#include "simvp_config.dtg.h" + +namespace FlexFlow { + +// Helper functions to construct the SimVP model + +/** + * @brief Get the default configs of SimVP model. + */ +SimVPConfig get_default_simvp_config(); + +/** + * @brief Get the SimVP computation graph. + * + * @param SimVPConfig The config of the SimVP model. + * @return ComputationGraph The computation graph of a SimVP model. + */ +ComputationGraph get_simvp_computation_graph(SimVPConfig const &config); + +} // namespace FlexFlow + +#endif diff --git a/lib/models/include/models/simvp/simvp_config.struct.toml b/lib/models/include/models/simvp/simvp_config.struct.toml new file mode 100644 index 0000000000..583990c42f --- /dev/null +++ b/lib/models/include/models/simvp/simvp_config.struct.toml @@ -0,0 +1,28 @@ +namespace = "FlexFlow" +name = "SimVPConfig" + +features = [ + "eq", + "ord", + "hash", + "json", + "rapidcheck", + "fmt", +] + +includes = [ + "", + "", + "", +] + +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" diff --git a/lib/models/src/models/simvp/simvp.cc b/lib/models/src/models/simvp/simvp.cc new file mode 100644 index 0000000000..4ae8695174 --- /dev/null +++ b/lib/models/src/models/simvp/simvp.cc @@ -0,0 +1,14 @@ +#include "models/simvp/simvp.h" + +namespace FlexFlow { + +SimVPConfig get_default_simvp_config() { + return SimVPConfig{64}; +} + +ComputationGraph get_simvp_computation_graph(SimVPConfig const &config) { + ComputationGraphBuilder cgb; + return cgb.computation_graph; +} + +} // namespace FlexFlow diff --git a/lib/models/test/src/models/simvp/simvp.cc b/lib/models/test/src/models/simvp/simvp.cc new file mode 100644 index 0000000000..a75944916d --- /dev/null +++ b/lib/models/test/src/models/simvp/simvp.cc @@ -0,0 +1,19 @@ +#include "models/simvp/simvp.h" +#include "pcg/computation_graph.h" +#include + +using namespace ::FlexFlow; + +TEST_SUITE(FF_TEST_SUITE) { + TEST_CASE("get_simvp_computation_graph") { + SimVPConfig config = get_default_simvp_config(); + + ComputationGraph result = get_simvp_computation_graph(config); + + SUBCASE("num layers") { + int result_num_layers = get_layers(result).size(); + int correct_num_layers = 0; + CHECK(result_num_layers == correct_num_layers); + } + } +} From 6fe279c6d6f38c110171474707aec90e93a150a7 Mon Sep 17 00:00:00 2001 From: easyeasydev Date: Wed, 13 Nov 2024 20:06:35 -0500 Subject: [PATCH 3/5] Add main structue and default configs --- lib/models/include/models/simvp/simvp.h | 9 +++ .../models/simvp/simvp_config.struct.toml | 44 ++++++++++++++ lib/models/src/models/simvp/simvp.cc | 57 ++++++++++++++++++- lib/models/test/src/models/simvp/simvp.cc | 2 +- 4 files changed, 110 insertions(+), 2 deletions(-) diff --git a/lib/models/include/models/simvp/simvp.h b/lib/models/include/models/simvp/simvp.h index 4eed320a7b..e527344670 100644 --- a/lib/models/include/models/simvp/simvp.h +++ b/lib/models/include/models/simvp/simvp.h @@ -13,6 +13,15 @@ namespace FlexFlow { */ SimVPConfig get_default_simvp_config(); +std::pair + create_simvp_encoder(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &input); + +tensor_guid_t create_simvp_middle_net(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &embed); + /** * @brief Get the SimVP computation graph. * diff --git a/lib/models/include/models/simvp/simvp_config.struct.toml b/lib/models/include/models/simvp/simvp_config.struct.toml index 583990c42f..a753f1f237 100644 --- a/lib/models/include/models/simvp/simvp_config.struct.toml +++ b/lib/models/include/models/simvp/simvp_config.struct.toml @@ -26,3 +26,47 @@ src_includes = [ [[fields]] name = "batch_size" type = "size_t" + +[[fields]] +name = "hid_S" +type = "size_t" + +[[fields]] +name = "hid_T" +type = "size_t" + +[[fields]] +name = "N_S" +type = "size_t" + +[[fields]] +name = "N_T" +type = "size_t" + +[[fields]] +name = "model_type" +type = "std::string" + +[[fields]] +name = "mlp_ratio" +type = "float" + +[[fields]] +name = "drop" +type = "float" + +[[fields]] +name = "drop_path" +type = "float" + +[[fields]] +name = "spatio_kernel_enc" +type = "size_t" + +[[fields]] +name = "spatio_kernel_dec" +type = "size_t" + +[[fields]] +name = "in_shape" +type = "std::vector" diff --git a/lib/models/src/models/simvp/simvp.cc b/lib/models/src/models/simvp/simvp.cc index 4ae8695174..53b927fd99 100644 --- a/lib/models/src/models/simvp/simvp.cc +++ b/lib/models/src/models/simvp/simvp.cc @@ -3,11 +3,66 @@ namespace FlexFlow { SimVPConfig get_default_simvp_config() { - return SimVPConfig{64}; + return SimVPConfig{ + /*batch_size=*/1, + /*hid_S=*/16, + /*hid_T=*/256, + /*N_S=*/4, + /*N_T=*/4, + /*model_type=*/"gSTA", + /*mlp_ratio=*/8.0, + /*drop=*/0.0, + /*drop_path=*/0.0, + /*spatio_kernel_enc=*/3, + /*spatio_kernel_dec=*/3, + /*in_shape=*/ + {10, 3, 32, 32}, + }; +} + +std::pair + create_simvp_encoder(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &input) { + return {input, input}; +} + +tensor_guid_t create_simvp_middle_net(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &embed) { + return embed; +} + +tensor_guid_t create_simvp_decoder(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &hid, + tensor_guid_t const &skip) { + return hid; } ComputationGraph get_simvp_computation_graph(SimVPConfig const &config) { ComputationGraphBuilder cgb; + + // Create input tensor + size_t B = config.batch_size; // Number of samples + size_t T = config.in_shape.at(0); // Number of frames in each sample + size_t C = config.in_shape.at(1); // Channel + size_t H = config.in_shape.at(2); // Height + size_t W = config.in_shape.at(3); // Width + + TensorShape input_shape = TensorShape{ + TensorDims{FFOrdered{B * T, C, H, W}}, + DataType::FLOAT, + }; + tensor_guid_t input = cgb.create_input(input_shape, CreateGrad::YES); + + // Create the model + auto [embed, skip] = create_simvp_encoder(cgb, config, input); + + tensor_guid_t hid = create_simvp_middle_net(cgb, config, embed); + + tensor_guid_t output = create_simvp_decoder(cgb, config, hid, skip); + return cgb.computation_graph; } diff --git a/lib/models/test/src/models/simvp/simvp.cc b/lib/models/test/src/models/simvp/simvp.cc index a75944916d..7bc6d74f98 100644 --- a/lib/models/test/src/models/simvp/simvp.cc +++ b/lib/models/test/src/models/simvp/simvp.cc @@ -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 = 1; CHECK(result_num_layers == correct_num_layers); } } From 635bf0b45535be9503c0730c129f4644e31aedcb Mon Sep 17 00:00:00 2001 From: easyeasydev Date: Wed, 13 Nov 2024 20:38:37 -0500 Subject: [PATCH 4/5] Add encoder --- lib/models/include/models/simvp/simvp.h | 14 ++++++++++ lib/models/src/models/simvp/simvp.cc | 37 ++++++++++++++++++++++++- 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/lib/models/include/models/simvp/simvp.h b/lib/models/include/models/simvp/simvp.h index e527344670..bcb89f5a94 100644 --- a/lib/models/include/models/simvp/simvp.h +++ b/lib/models/include/models/simvp/simvp.h @@ -13,6 +13,15 @@ namespace FlexFlow { */ SimVPConfig get_default_simvp_config(); +std::vector create_simvp_samplings(size_t N_S, bool reverse = false); + +tensor_guid_t create_simvp_convsc(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + size_t in_dim, + size_t out_dim, + bool downsampling, + tensor_guid_t const &input); + std::pair create_simvp_encoder(ComputationGraphBuilder &cgb, SimVPConfig const &config, @@ -22,6 +31,11 @@ tensor_guid_t create_simvp_middle_net(ComputationGraphBuilder &cgb, SimVPConfig const &config, tensor_guid_t const &embed); +tensor_guid_t create_simvp_decoder(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &hid, + tensor_guid_t const &skip); + /** * @brief Get the SimVP computation graph. * diff --git a/lib/models/src/models/simvp/simvp.cc b/lib/models/src/models/simvp/simvp.cc index 53b927fd99..3ea0f6fd03 100644 --- a/lib/models/src/models/simvp/simvp.cc +++ b/lib/models/src/models/simvp/simvp.cc @@ -20,19 +20,54 @@ SimVPConfig get_default_simvp_config() { }; } +// TODO: improve with a transform call +std::vector create_simvp_samplings(size_t N_S, bool reverse) { + std::vector samplings(N_S, false); + for (size_t i = 0; i < N_S; i++) { + if (i % 2 == 1) { + samplings[i] = true; + } + } + return samplings; +} + +// TODO +tensor_guid_t create_simvp_convsc(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + size_t in_dim, + size_t out_dim, + bool downsampling, + tensor_guid_t const &input) { + return input; +} + std::pair create_simvp_encoder(ComputationGraphBuilder &cgb, SimVPConfig const &config, tensor_guid_t const &input) { - return {input, input}; + size_t C = config.in_shape.at(1); // Channel + std::vector samplings = create_simvp_samplings(config.N_S); + + tensor_guid_t enc1 = + create_simvp_convsc(cgb, config, C, config.hid_S, samplings[0], input); + tensor_guid_t latent = enc1; + + for (size_t i = 1; i < config.N_S; i++) { + latent = create_simvp_convsc( + cgb, config, config.hid_S, config.hid_S, samplings[i], latent); + } + + return {latent, enc1}; } +// TODO tensor_guid_t create_simvp_middle_net(ComputationGraphBuilder &cgb, SimVPConfig const &config, tensor_guid_t const &embed) { return embed; } +// TODO tensor_guid_t create_simvp_decoder(ComputationGraphBuilder &cgb, SimVPConfig const &config, tensor_guid_t const &hid, From ccc8b2c0efff9c8754b6a5db7834f1bae2bb42a8 Mon Sep 17 00:00:00 2001 From: easyeasydev Date: Wed, 4 Dec 2024 20:49:19 -0500 Subject: [PATCH 5/5] Implement midnet and decoder --- lib/models/include/models/simvp/simvp.h | 32 +++- lib/models/src/models/simvp/simvp.cc | 184 +++++++++++++++++++--- lib/models/test/src/models/simvp/simvp.cc | 2 +- 3 files changed, 195 insertions(+), 23 deletions(-) diff --git a/lib/models/include/models/simvp/simvp.h b/lib/models/include/models/simvp/simvp.h index bcb89f5a94..8485ed58cf 100644 --- a/lib/models/include/models/simvp/simvp.h +++ b/lib/models/include/models/simvp/simvp.h @@ -17,10 +17,31 @@ std::vector create_simvp_samplings(size_t N_S, bool reverse = false); tensor_guid_t create_simvp_convsc(ComputationGraphBuilder &cgb, SimVPConfig const &config, + tensor_guid_t const &input, size_t in_dim, size_t out_dim, - bool downsampling, - tensor_guid_t const &input); + int kernel_size = 3, + bool downsampling = false, + bool upsampling = false); + +tensor_guid_t create_simvp_gsta_meta_block(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &input, + int in_channels, + int out_channels, + float mlp_ratio = 8.0, + float drop = 0.0, + float drop_path = 0.0); + +tensor_guid_t create_simvp_ga_sub_block(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &input, + int dim, + int kernel_size = 21, + float mlp_ratio = 4.0, + float drop = 0.0, + float drop_path = 0.1, + float init_value = 1e-2); std::pair create_simvp_encoder(ComputationGraphBuilder &cgb, @@ -29,7 +50,12 @@ std::pair tensor_guid_t create_simvp_middle_net(ComputationGraphBuilder &cgb, SimVPConfig const &config, - tensor_guid_t const &embed); + tensor_guid_t const &embed, + int channel_in, + int channel_hid, + float mlp_ratio = 4.0, + float drop = 0.0, + float drop_path = 0.1); tensor_guid_t create_simvp_decoder(ComputationGraphBuilder &cgb, SimVPConfig const &config, diff --git a/lib/models/src/models/simvp/simvp.cc b/lib/models/src/models/simvp/simvp.cc index 3ea0f6fd03..da2753cc3f 100644 --- a/lib/models/src/models/simvp/simvp.cc +++ b/lib/models/src/models/simvp/simvp.cc @@ -20,25 +20,45 @@ SimVPConfig get_default_simvp_config() { }; } -// TODO: improve with a transform call std::vector create_simvp_samplings(size_t N_S, bool reverse) { - std::vector samplings(N_S, false); - for (size_t i = 0; i < N_S; i++) { - if (i % 2 == 1) { + size_t N_S_even_floor = (N_S / 2) * 2; + + auto const change_to_true = [&](size_t idx) -> bool { + return (reverse == false) ? (idx % 2 == 1) : (idx % 2 == 0); + }; + + std::vector samplings(N_S_even_floor, false); + for (size_t i = 0; i < N_S_even_floor; i++) { + if (change_to_true(i)) { samplings[i] = true; } } + return samplings; } -// TODO tensor_guid_t create_simvp_convsc(ComputationGraphBuilder &cgb, SimVPConfig const &config, + tensor_guid_t const &input, size_t in_dim, size_t out_dim, + int kernel_size, bool downsampling, - tensor_guid_t const &input) { - return input; + bool upsampling) { + int stride = (downsampling == true) ? 2 : 1; + int padding = (kernel_size - stride + 1) / 2; + int out_channels = upsampling ? out_dim * 4 : out_dim; + + tensor_guid_t conv_out = cgb.conv2d(input, + out_channels, + kernel_size, + kernel_size, + stride, + stride, + padding, + padding); + + return conv_out; } std::pair @@ -48,31 +68,130 @@ std::pair size_t C = config.in_shape.at(1); // Channel std::vector samplings = create_simvp_samplings(config.N_S); - tensor_guid_t enc1 = - create_simvp_convsc(cgb, config, C, config.hid_S, samplings[0], input); + tensor_guid_t enc1 = create_simvp_convsc(cgb, + config, + input, + C, + config.hid_S, + config.spatio_kernel_enc, + samplings[0], + false); tensor_guid_t latent = enc1; - for (size_t i = 1; i < config.N_S; i++) { - latent = create_simvp_convsc( - cgb, config, config.hid_S, config.hid_S, samplings[i], latent); + for (size_t i = 1; i < samplings.size(); i++) { + latent = create_simvp_convsc(cgb, + config, + latent, + config.hid_S, + config.hid_S, + config.spatio_kernel_enc, + samplings[i], + false); } return {latent, enc1}; } // TODO +tensor_guid_t create_simvp_ga_sub_block(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &input, + int dim, + int kernel_size, + float mlp_ratio, + float drop, + float drop_path, + float init_value) { + return input; +} + +tensor_guid_t create_simvp_gsta_meta_block(ComputationGraphBuilder &cgb, + SimVPConfig const &config, + tensor_guid_t const &input, + int in_channels, + int out_channels, + float mlp_ratio, + float drop, + float drop_path) { + tensor_guid_t z = create_simvp_ga_sub_block( + cgb, config, input, in_channels, 21, mlp_ratio, drop, drop_path); + + if (in_channels == out_channels) { + return z; + } else { + return cgb.conv2d(z, out_channels, 1, 1, 1, 1, 0, 0); + } +} + tensor_guid_t create_simvp_middle_net(ComputationGraphBuilder &cgb, SimVPConfig const &config, - tensor_guid_t const &embed) { - return embed; + tensor_guid_t const &embed, + int channel_in, + int channel_hid, + float mlp_ratio, + float drop, + float drop_path) { + if (config.model_type != "gSTA") { + throw mk_runtime_error( + fmt::format("Currently only model_type=gSTA is " + "supported, but found model_type={}. " + "If you need support for additional " + "model_type values, please create an issue.", + config.model_type)); + } + + tensor_guid_t z = embed; + + // Downsample + z = create_simvp_gsta_meta_block( + cgb, config, z, channel_in, channel_hid, mlp_ratio, drop, drop_path); + + // Middle layers + for (size_t i = 1; i < config.N_T - 1; i++) { + z = create_simvp_gsta_meta_block( + cgb, config, z, channel_hid, channel_hid, mlp_ratio, drop, drop_path); + } + + // Upsample + z = create_simvp_gsta_meta_block( + cgb, config, z, channel_hid, channel_in, mlp_ratio, drop, drop_path); + + return z; } -// TODO tensor_guid_t create_simvp_decoder(ComputationGraphBuilder &cgb, SimVPConfig const &config, tensor_guid_t const &hid, tensor_guid_t const &skip) { - return hid; + + std::cout << "hid shape: " << cgb.get_shape(hid) << std::endl; + std::cout << "skip shape: " << cgb.get_shape(skip) << std::endl; + + size_t C = config.in_shape.at(1); // Channel + std::vector samplings = create_simvp_samplings(config.N_S, true); + + tensor_guid_t latent = hid; + for (size_t i = 0; i < samplings.size() - 1; i++) { + latent = create_simvp_convsc(cgb, + config, + latent, + config.hid_S, + config.hid_S, + config.spatio_kernel_dec, + false, + samplings[i]); + } + + tensor_guid_t out = create_simvp_convsc(cgb, + config, + cgb.add(latent, skip), + config.hid_S, + config.hid_S, + config.spatio_kernel_dec, + false, + samplings[samplings.size() - 1]); + + return cgb.conv2d(out, C, 1, 1, 1, 1, 0, 0, std::nullopt); } ComputationGraph get_simvp_computation_graph(SimVPConfig const &config) { @@ -85,6 +204,9 @@ ComputationGraph get_simvp_computation_graph(SimVPConfig const &config) { size_t H = config.in_shape.at(2); // Height size_t W = config.in_shape.at(3); // Width + // std::cout << "B T C H W: " << B << " " << T << " " << C << " " << H << " " + // << W << std::endl; + TensorShape input_shape = TensorShape{ TensorDims{FFOrdered{B * T, C, H, W}}, DataType::FLOAT, @@ -94,9 +216,33 @@ ComputationGraph get_simvp_computation_graph(SimVPConfig const &config) { // Create the model auto [embed, skip] = create_simvp_encoder(cgb, config, input); - tensor_guid_t hid = create_simvp_middle_net(cgb, config, embed); - - tensor_guid_t output = create_simvp_decoder(cgb, config, hid, skip); + // std::cout << "embed shape: " << cgb.get_shape(embed) << std::endl; + + // TODO: need to reshape embed here + + tensor_guid_t hid = create_simvp_middle_net(cgb, + config, + embed, + T * config.hid_S, + config.hid_T, + config.mlp_ratio, + config.drop, + config.drop_path); + + // TODO: need to reshape hid here + // std::cout << "hid shape: " << cgb.get_shape(hid) << std::endl; + // auto const embed_shape = cgb.get_shape(embed).dims.ff_ordered; + // std::vector hid_shape = { + // static_cast(B * T), + // static_cast(embed_shape.at(ff_dim_t{1})), + // static_cast(embed_shape.at(ff_dim_t{2})), + // static_cast(embed_shape.at(ff_dim_t{3})), + // }; + // hid = cgb.reshape(hid, hid_shape); + // std::cout << "hid new shape: " << cgb.get_shape(hid) << std::endl; + + // TODO: enable below after reshaping + // tensor_guid_t output = create_simvp_decoder(cgb, config, hid, skip); return cgb.computation_graph; } diff --git a/lib/models/test/src/models/simvp/simvp.cc b/lib/models/test/src/models/simvp/simvp.cc index 7bc6d74f98..c6939b7d20 100644 --- a/lib/models/test/src/models/simvp/simvp.cc +++ b/lib/models/test/src/models/simvp/simvp.cc @@ -12,7 +12,7 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("num layers") { int result_num_layers = get_layers(result).size(); - int correct_num_layers = 1; + int correct_num_layers = 19; CHECK(result_num_layers == correct_num_layers); } }