Insert gradient reduction during pass expansion when multiple consumers use the same input

lib/realm-execution/src/realm-execution/pcg_instance.cc

@@ -341,6 +341,7 @@ static Realm::Event spawn_dynamic_node_invocation(
       },
       [&](LossAttrs const &) { return spawn_task(); },
       [&](CopyAttrs const &) { return issue_copy(); },
+      [&](GradientReductionAttrs const &) { return issue_reduction(); },
   });
 }
 

lib/realm-execution/test/src/realm-execution/test_e2e.cc

@@ -411,6 +411,144 @@ MappedParallelComputationGraph
   return mpcg;
 }
 
+MappedParallelComputationGraph
+    make_test_gradient_reduction_mpcg_for_device_type(DeviceType device_type) {
+  positive_int batch_size = 10_p;
+  positive_int data_dim = 16_p;
+  positive_int hidden_dim = 32_p;
+  positive_int output_dim = 1_p;
+
+  TensorShape output_tensor_shape = TensorShape{
+      TensorDims{FFOrdered{batch_size, output_dim}}, DataType::FLOAT};
+
+  TensorShape label_tensor_shape = TensorShape{
+      TensorDims{FFOrdered{batch_size, output_dim}}, DataType::FLOAT};
+
+  ParallelComputationGraph pcg = empty_parallel_computation_graph();
+
+  TensorShape input_tensor_shape =
+      TensorShape{TensorDims{FFOrdered{batch_size, data_dim}}, DataType::FLOAT};
+
+  ParallelLayerAddedResult inputs_layer =
+      pcg_add_input_layer(pcg, input_tensor_shape, CreateGrad::YES);
+  parallel_tensor_guid_t t_input =
+      require_only_key(inputs_layer.outputs, TensorSlotName::OUTPUT);
+
+  ParallelLayerAddedResult relu_operator_1 =
+      add_parallel_layer(pcg,
+                         make_layer_attrs(make_relu_attrs()),
+                         {
+                             {
+                                 TensorSlotName::INPUT,
+                                 t_input,
+                             },
+                         },
+                         /*weights=*/{});
+
+  parallel_tensor_guid_t t_relu_1 =
+      require_only_key(relu_operator_1.outputs, TensorSlotName::OUTPUT);
+
+  ParallelLayerAddedResult relu_operator_2 =
+      add_parallel_layer(pcg,
+                         make_layer_attrs(make_relu_attrs()),
+                         {
+                             {
+                                 TensorSlotName::INPUT,
+                                 t_relu_1,
+                             },
+                         },
+                         /*weights=*/{});
+
+  parallel_tensor_guid_t t_relu_2 =
+      require_only_key(relu_operator_2.outputs, TensorSlotName::OUTPUT);
+
+  ParallelLayerAddedResult relu_operator_3 =
+      add_parallel_layer(pcg,
+                         make_layer_attrs(make_relu_attrs()),
+                         {
+                             {
+                                 TensorSlotName::INPUT,
+                                 t_relu_1,
+                             },
+                         },
+                         /*weights=*/{});
+
+  parallel_tensor_guid_t t_relu_3 =
+      require_only_key(relu_operator_3.outputs, TensorSlotName::OUTPUT);
+
+  MachineSpaceCoordinate mc0{0_n, 0_n};
+  MachineSpaceCoordinate mc1{0_n, 1_n};
+  MachineSpaceCoordinate mc2{0_n, 2_n};
+
+  ParallelTensorSpaceCoordinate tensor_coord0{
+      /*sum_component=*/0_n,
+      /*discard_copy_component=*/0_n,
+      /*shard_component=*/FFOrdered{0_n}};
+
+  MappedParallelComputationGraph mpcg =
+      mapped_pcg_from_pcg_and_mapped_op_task_groups(
+          /*pcg=*/pcg,
+          /*mapped_op_task_groups=*/{
+              {
+                  inputs_layer.parallel_layer,
+                  MappedOperatorTaskGroup{
+                      {
+                          {
+                              mc0,
+                              OperatorAtomicTaskShardBinding{{
+                                  {TensorSlotName::OUTPUT, tensor_coord0},
+                              }},
+                          },
+                      },
+                  },
+              },
+              {
+                  relu_operator_1.parallel_layer,
+                  MappedOperatorTaskGroup{
+                      {
+                          {
+                              mc0,
+                              OperatorAtomicTaskShardBinding{{
+                                  {TensorSlotName::INPUT, tensor_coord0},
+                                  {TensorSlotName::OUTPUT, tensor_coord0},
+                              }},
+                          },
+                      },
+                  },
+              },
+              {
+                  relu_operator_2.parallel_layer,
+                  MappedOperatorTaskGroup{
+                      {
+                          {
+                              mc1,
+                              OperatorAtomicTaskShardBinding{{
+                                  {TensorSlotName::INPUT, tensor_coord0},
+                                  {TensorSlotName::OUTPUT, tensor_coord0},
+                              }},
+                          },
+                      },
+                  },
+              },
+              {
+                  relu_operator_3.parallel_layer,
+                  MappedOperatorTaskGroup{
+                      {
+                          {
+                              mc2,
+                              OperatorAtomicTaskShardBinding{{
+                                  {TensorSlotName::INPUT, tensor_coord0},
+                                  {TensorSlotName::OUTPUT, tensor_coord0},
+                              }},
+                          },
+                      },
+                  },
+              },
+          });
+
+  return mpcg;
+}
+
 TEST_SUITE(FF_TEST_SUITE) {
   TEST_CASE("RealmBackend e2e Training (CPU Model Parallelism)") {
     std::vector<char *> fake_args =
@@ -537,6 +675,60 @@ TEST_SUITE(FF_TEST_SUITE) {
         });
     result.wait();
   }
+
+  TEST_CASE("RealmBackend e2e Training Gradient Reduction Op (CPU Model "
+            "Parallelism)") {
+    std::vector<char *> fake_args =
+        make_fake_realm_args(/*num_cpus=*/3_p, /*num_gpus=*/0_n);
+    int fake_argc = fake_args.size();
+    char **fake_argv = fake_args.data();
+
+    RealmManager manager = RealmManager{&fake_argc, &fake_argv};
+    ControllerTaskResult result =
+        manager.start_controller([](RealmContext &ctx) {
+          Allocator allocator = ctx.get_current_device_allocator();
+
+          MappedParallelComputationGraph mpcg =
+              make_test_gradient_reduction_mpcg_for_device_type(
+                  DeviceType::CPU);
+
+          std::map<DynamicValueAttrs, DynamicTensorAccessor> input_tensors;
+
+          OptimizerAttrs optimizer_attrs = OptimizerAttrs{
+              SGDOptimizerAttrs{
+                  /*lr=*/0.001,
+                  /*momentum=*/0.9,
+                  /*nesterov=*/false,
+                  /*weight_decay=*/0.001,
+              },
+          };
+
+          DistributedFfHandle device_handle = create_distributed_ff_handle(
+              ctx,
+              /*workSpaceSize=*/1024 * 1024,
+              /*allowTensorOpMathConversion=*/true);
+
+          PCGInstance pcg_instance = create_pcg_instance(
+              /*ctx=*/ctx,
+              /*mpcg=*/mpcg,
+              /*optimizer=*/optimizer_attrs,
+              /*loss=*/std::nullopt,
+              /*input_tensors=*/input_tensors,
+              /*profiling_settings=*/ProfilingSettings{0, 0},
+              /*device_handle=*/device_handle,
+              /*device_type=*/DeviceType::CPU);
+
+          // begin training loop
+          int num_epochs = 1;
+          for (int i = 0; i < num_epochs; i++) {
+            perform_all_passes_for_pcg_instance(
+                /*instance=*/pcg_instance,
+                /*profiling_settings=*/ProfilingSettings{0, 0},
+                /*device_handle=*/device_handle);
+          }
+        });
+    result.wait();
+  }
 }
 
 TEST_SUITE(FF_CUDA_TEST_SUITE) {
@@ -672,6 +864,61 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) {
         });
     result.wait();
   }
+
+  TEST_CASE("RealmBackend e2e Training Gradient Reduction Op (GPU Model "
+            "Parallelism)") {
+    std::vector<char *> fake_args =
+        make_fake_realm_args(/*num_cpus=*/1_p, /*num_gpus=*/3_n);
+    int fake_argc = fake_args.size();
+    char **fake_argv = fake_args.data();
+
+    RealmManager manager = RealmManager{&fake_argc, &fake_argv};
+
+    ControllerTaskResult result =
+        manager.start_controller([](RealmContext &ctx) {
+          Allocator allocator = ctx.get_current_device_allocator();
+
+          MappedParallelComputationGraph mpcg =
+              make_test_gradient_reduction_mpcg_for_device_type(
+                  DeviceType::GPU);
+
+          OptimizerAttrs optimizer_attrs = OptimizerAttrs{
+              SGDOptimizerAttrs{
+                  /*lr=*/0.001,
+                  /*momentum=*/0.9,
+                  /*nesterov=*/false,
+                  /*weight_decay=*/0.001,
+              },
+          };
+
+          std::map<DynamicValueAttrs, DynamicTensorAccessor> input_tensors;
+
+          DistributedFfHandle device_handle = create_distributed_ff_handle(
+              ctx,
+              /*workSpaceSize=*/1024 * 1024,
+              /*allowTensorOpMathConversion=*/true);
+
+          PCGInstance pcg_instance = create_pcg_instance(
+              /*ctx=*/ctx,
+              /*mpcg=*/mpcg,
+              /*optimizer=*/optimizer_attrs,
+              /*loss=*/std::nullopt,
+              /*input_tensors=*/input_tensors,
+              /*profiling_settings=*/ProfilingSettings{0, 0},
+              /*device_handle=*/device_handle,
+              /*device_type=*/DeviceType::GPU);
+
+          // begin training loop
+          int num_epochs = 1;
+          for (int i = 0; i < num_epochs; i++) {
+            perform_all_passes_for_pcg_instance(
+                /*instance=*/pcg_instance,
+                /*profiling_settings=*/ProfilingSettings{0, 0},
+                /*device_handle=*/device_handle);
+          }
+        });
+    result.wait();
+  }
 }
 
 } // namespace test

lib/task-spec/include/task-spec/dynamic_graph/dynamic_gradient_reduction_layer_guid_t.dtg.toml

@@ -0,0 +1,13 @@
+namespace = "FlexFlow"
+name = "dynamic_gradient_reduction_layer_guid_t"
+type = "struct"
+features = [
+  "eq",
+  "ord",
+  "hash",
+  "json",
+  "fmt",
+  "rapidcheck",
+]
+
+fields = []

lib/task-spec/include/task-spec/dynamic_graph/dynamic_layer_guid_t.dtg.toml

@@ -14,6 +14,7 @@ includes = [
   "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h",
   "task-spec/dynamic_graph/dynamic_loss_layer_guid_t.dtg.h",
   "task-spec/dynamic_graph/dynamic_copy_layer_guid_t.dtg.h",
+  "task-spec/dynamic_graph/dynamic_gradient_reduction_layer_guid_t.dtg.h",
 ]
 
 [[values]]
@@ -31,3 +32,7 @@ key = "loss_layer_guid"
 [[values]]
 type = "::FlexFlow::dynamic_copy_layer_guid_t"
 key = "copy_layer_guid"
+
+[[values]]
+type = "::FlexFlow::dynamic_gradient_reduction_layer_guid_t"
+key = "gradient_reduction_layer_guid"

lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h

@@ -2,6 +2,8 @@
 #define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_NODE_MAPPING_H
 
 #include "task-spec/dynamic_graph/dynamic_node_mapping.dtg.h"
+#include "task-spec/dynamic_graph/dynamic_tensor_slot.dtg.h"
+#include "task-spec/dynamic_graph/parallel_tensor_mapping.dtg.h"
 #include "task-spec/global_device_id_t.dtg.h"
 
 namespace FlexFlow {

lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml

@@ -13,6 +13,7 @@ includes = [
   "<optional>",
   "task-spec/dynamic_graph/dynamic_tensor_guid_t.dtg.h",
   "op-attrs/parallel_tensor_shape.dtg.h",
+  "task-spec/dynamic_graph/subgradient_id_t.dtg.h",
   "op-attrs/parallel_tensor_space_coordinate.dtg.h",
   "pcg/machine_space_coordinate.dtg.h",
   "utils/bidict/bidict.h",
@@ -48,6 +49,15 @@ For a \ref DynamicOpenDataflowGraph originating form a \ref MappedParallelComput
 name = "create_grad"
 type = "std::optional<bool>"
 
+[[fields]]
+name = "subgradient_id"
+type = "std::optional<::FlexFlow::subgradient_id_t>"
+docstring = '''
+\brief The unique identity of the subgradient represented by this value.
+
+This field is filled in by \ref pass_expansion.h when a value is consumed multiple times, requiring a \ref GradientReductionAttrs in the backward pass.
+'''
+
 [[fields]]
 name = "shard_coord"
 type = "std::optional<::FlexFlow::ParallelTensorSpaceCoordinate>"

lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.h

@@ -3,12 +3,17 @@
 
 #include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h"
 #include "task-spec/dynamic_graph/parallel_tensor_mapping.dtg.h"
+#include "task-spec/dynamic_graph/subgradient_id_t.dtg.h"
 
 namespace FlexFlow {
 
 DynamicValueAttrs decide_dynamic_value_attrs_role(DynamicValueAttrs const &,
                                                   DynamicTensorRole);
 
+DynamicValueAttrs
+    decide_dynamic_value_attrs_subgradient_id(DynamicValueAttrs const &,
+                                              subgradient_id_t const &);
+
 DynamicValueAttrs
     decide_dynamic_value_attrs_mapping(DynamicValueAttrs const &,
                                        ParallelTensorMapping const &);

lib/task-spec/include/task-spec/dynamic_graph/gradient_reduction_attrs.dtg.toml

@@ -0,0 +1,13 @@
+namespace = "FlexFlow"
+name = "GradientReductionAttrs"
+type = "struct"
+features = [
+  "eq",
+  "ord",
+  "hash",
+  "json",
+  "fmt",
+  "rapidcheck",
+]
+
+fields = []

lib/task-spec/include/task-spec/dynamic_graph/internal_dynamic_slot_site.dtg.toml

@@ -3,7 +3,7 @@ name = "InternalDynamicSlotSite"
 type = "struct"
 features = [
   "eq",
-	"ord",
+  "ord",
   "hash",
   "fmt",
   "json",

lib/task-spec/include/task-spec/dynamic_graph/pass_expansion.h

@@ -41,8 +41,15 @@ DynamicNodeAttrs pass_expand_node(DynamicNodeAttrs const &, DynamicTaskType);
 
 DynamicNodeInvocation
     perform_fwd_pass_expansion_for_invocation(DynamicNodeInvocation const &);
-DynamicNodeInvocation
-    perform_bwd_pass_expansion_for_invocation(DynamicNodeInvocation const &);
+DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation(
+    dynamic_invocation_id_t const &,
+    DynamicNodeInvocation const &,
+    std::map<DynamicValueAttrs,
+             std::map<dynamic_invocation_id_t, subgradient_id_t>> const &);
+DynamicNodeInvocation create_gradient_reduction_for_value(
+    DynamicOpenDataflowGraph const &,
+    DynamicValueAttrs const &,
+    std::map<dynamic_invocation_id_t, subgradient_id_t> const &);
 
 DynamicOpenDataflowGraph
     perform_pass_expansion(DynamicOpenDataflowGraph const &);