From eeaa135d5a592d1f695244a81feb6b72dab885a6 Mon Sep 17 00:00:00 2001
From: Elliott Slaughter <slaughter@cs.stanford.edu>
Date: Thu, 20 Aug 2026 10:12:40 -0700
Subject: [PATCH 1/9] Writing some intermediate tests for the multi-consumer
 pattern.

---
 .../task-spec/dynamic_graph/pass_expansion.cc | 414 ++++++++++++++++++
 1 file changed, 414 insertions(+)

diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc
index d3a915fb5c..731df510cc 100644
--- a/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc
+++ b/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc
@@ -1,6 +1,9 @@
 #include "task-spec/dynamic_graph/pass_expansion.h"
 #include "op-attrs/initializer_attrs.h"
+#include "op-attrs/operator_type.dtg.h"
+#include "op-attrs/ops/element_binary_attrs.dtg.h"
 #include "op-attrs/ops/element_unary.h"
+#include "op-attrs/tensor_slot_name.dtg.h"
 #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.h"
 #include "task-spec/dynamic_graph/dynamic_tensor_role.h"
 #include "task-spec/dynamic_graph/serializable_dynamic_node_invocation.h"
@@ -370,6 +373,130 @@ TEST_SUITE(FF_TEST_SUITE) {
     DynamicValueAttrs v2_grad = mk_value_attrs(1, grad_role);
     DynamicValueAttrs v3_grad = mk_value_attrs(2, grad_role);
 
+    SUBCASE("relu operator") {
+      TrainingOperationAttrs op_attrs = TrainingOperationAttrs{
+          PCGOperatorAttrs{
+              make_relu_attrs(),
+          },
+      };
+
+      DynamicNodeInvocation invocation = [&]() -> DynamicNodeInvocation {
+        return DynamicNodeInvocation{
+            /*inputs=*/{
+                {mk_slot(TensorSlotName::INPUT, std::nullopt), v1},
+            },
+            /*node_attrs=*/
+            DynamicNodeAttrs{
+                /*task_type=*/std::nullopt,
+                /*device_coord=*/std::nullopt,
+                /*mapping=*/std::nullopt,
+                /*op_attrs=*/op_attrs,
+                /*layer_guid=*/layer_guid,
+                /*per_device_op_state=*/std::nullopt,
+            },
+            /*outputs=*/
+            {
+                {mk_slot(TensorSlotName::OUTPUT, std::nullopt), v2},
+            },
+        };
+      }();
+
+      DynamicNodeInvocation result =
+          perform_bwd_pass_expansion_for_invocation(invocation);
+
+      DynamicNodeInvocation correct = [&]() -> DynamicNodeInvocation {
+        return DynamicNodeInvocation{
+            /*inputs=*/{
+                {mk_slot(TensorSlotName::INPUT, fwd_role), v1_fwd},
+                {mk_slot(TensorSlotName::OUTPUT, fwd_role), v2_fwd},
+                {mk_slot(TensorSlotName::OUTPUT, grad_role), v2_grad},
+            },
+            /*node_attrs=*/
+            DynamicNodeAttrs{
+                /*pass_type=*/DynamicTaskType::BWD,
+                /*device_coord=*/std::nullopt,
+                /*mapping=*/std::nullopt,
+                /*op_attrs=*/op_attrs,
+                /*layer_guid=*/layer_guid,
+                /*per_device_op_state=*/std::nullopt,
+            },
+            /*outputs=*/
+            {
+                {mk_slot(TensorSlotName::INPUT, grad_role), v1_grad},
+            },
+        };
+      }();
+
+      ASSERT(dynamic_node_invocation_to_serializable(result) ==
+             dynamic_node_invocation_to_serializable(correct));
+    }
+
+    SUBCASE("add operator") {
+      TrainingOperationAttrs op_attrs = TrainingOperationAttrs{
+          PCGOperatorAttrs{
+              ElementBinaryAttrs{
+                  /*type=*/OperatorType::EW_ADD,
+                  /*compute_type=*/DataType::FLOAT,
+                  /*should_broadcast_lhs=*/false,
+                  /*should_broadcast_rhs=*/false,
+              },
+          },
+      };
+
+      DynamicNodeInvocation invocation = [&]() -> DynamicNodeInvocation {
+        return DynamicNodeInvocation{
+            /*inputs=*/{
+                {mk_slot(TensorSlotName::LHS_INPUT, std::nullopt), v1},
+                {mk_slot(TensorSlotName::RHS_INPUT, std::nullopt), v2},
+            },
+            /*node_attrs=*/
+            DynamicNodeAttrs{
+                /*task_type=*/std::nullopt,
+                /*device_coord=*/std::nullopt,
+                /*mapping=*/std::nullopt,
+                /*op_attrs=*/op_attrs,
+                /*layer_guid=*/layer_guid,
+                /*per_device_op_state=*/std::nullopt,
+            },
+            /*outputs=*/
+            {
+                {mk_slot(TensorSlotName::OUTPUT, std::nullopt), v3},
+            },
+        };
+      }();
+
+      DynamicNodeInvocation result =
+          perform_bwd_pass_expansion_for_invocation(invocation);
+
+      DynamicNodeInvocation correct = [&]() -> DynamicNodeInvocation {
+        return DynamicNodeInvocation{
+            /*inputs=*/{
+                {mk_slot(TensorSlotName::LHS_INPUT, fwd_role), v1_fwd},
+                {mk_slot(TensorSlotName::RHS_INPUT, fwd_role), v2_fwd},
+                {mk_slot(TensorSlotName::OUTPUT, fwd_role), v3_fwd},
+                {mk_slot(TensorSlotName::OUTPUT, grad_role), v3_grad},
+            },
+            /*node_attrs=*/
+            DynamicNodeAttrs{
+                /*pass_type=*/DynamicTaskType::BWD,
+                /*device_coord=*/std::nullopt,
+                /*mapping=*/std::nullopt,
+                /*op_attrs=*/op_attrs,
+                /*layer_guid=*/layer_guid,
+                /*per_device_op_state=*/std::nullopt,
+            },
+            /*outputs=*/
+            {
+                {mk_slot(TensorSlotName::LHS_INPUT, grad_role), v1_grad},
+                {mk_slot(TensorSlotName::RHS_INPUT, grad_role), v2_grad},
+            },
+        };
+      }();
+
+      ASSERT(dynamic_node_invocation_to_serializable(result) ==
+             dynamic_node_invocation_to_serializable(correct));
+    }
+
     SUBCASE("normal operator") {
       TrainingOperationAttrs op_attrs = TrainingOperationAttrs{
           PCGOperatorAttrs{
@@ -894,4 +1021,291 @@ TEST_SUITE(FF_TEST_SUITE) {
 
     CHECK(dynamic_open_dataflow_graphs_are_isomorphic(result, correct));
   }
+
+  TEST_CASE("perform_pass_expansion(DynamicOpenDataflowGraph) with multiple "
+            "consumers") {
+    auto mk_node_attrs = [](size_t layer_id,
+                            TrainingOperationAttrs const &op_attrs,
+                            std::optional<DynamicTaskType> const &pass_type)
+        -> DynamicNodeAttrs {
+      return DynamicNodeAttrs{
+          /*pass_type=*/pass_type,
+          /*device_coord=*/std::nullopt,
+          /*mapping=*/std::nullopt,
+          /*op_attrs=*/op_attrs,
+          /*layer_guid=*/
+          dynamic_layer_guid_t{parallel_layer_guid_t{Node{layer_id}}},
+          /*per_device_op_state=*/std::nullopt,
+      };
+    };
+
+    auto mk_value_attrs =
+        [](size_t node_id, std::optional<DynamicTensorRole> const &tensor_type)
+        -> DynamicValueAttrs {
+      return DynamicValueAttrs{
+          /*tensor_guid=*/dynamic_tensor_guid_t{parallel_tensor_guid_t{
+              KwargDataflowOutput{
+                  Node{node_id},
+                  TensorSlotName::OUTPUT,
+              },
+          }},
+          /*parallel_tensor_shape=*/std::nullopt,
+          /*create_grad=*/true,
+          /*shard_coord=*/std::nullopt,
+          /*mapping=*/std::nullopt,
+          /*accessor=*/std::nullopt,
+          /*role=*/tensor_type,
+      };
+    };
+
+    TrainingOperationAttrs input_op_attrs = TrainingOperationAttrs{
+        PCGOperatorAttrs{
+            InputAttrs{
+                TensorShape{
+                    TensorDims{
+                        FFOrdered<positive_int>{
+                            4_p,
+                            8_p,
+                        },
+                    },
+                    DataType::FLOAT,
+                },
+            },
+        },
+    };
+
+    TrainingOperationAttrs relu_op_attrs = TrainingOperationAttrs{
+        PCGOperatorAttrs{
+            make_relu_attrs(),
+        },
+    };
+
+    DynamicOpenDataflowGraph input = [&]() -> DynamicOpenDataflowGraph {
+      DynamicNodeAttrs input_node =
+          mk_node_attrs(10, input_op_attrs, std::nullopt);
+      DynamicNodeAttrs relu1_node =
+          mk_node_attrs(11, relu_op_attrs, std::nullopt);
+      DynamicNodeAttrs relu2_node =
+          mk_node_attrs(12, relu_op_attrs, std::nullopt);
+
+      DynamicValueAttrs input_tensor = mk_value_attrs(0, std::nullopt);
+      DynamicValueAttrs relu1_output = mk_value_attrs(1, std::nullopt);
+      DynamicValueAttrs relu2_output = mk_value_attrs(2, std::nullopt);
+
+      auto mk_dynamic_slot =
+          [](TensorSlotName const &slot_name) -> DynamicTensorSlot {
+        return DynamicTensorSlot{
+            /*slot_name=*/slot_name,
+            /*slot_tensor_role=*/std::nullopt,
+            /*task_shard=*/std::nullopt,
+        };
+      };
+
+      std::set<DynamicNodeInvocation> invocation_set = {
+          DynamicNodeInvocation{
+              /*inputs=*/std::map<DynamicTensorSlot, DynamicValueAttrs>{},
+              /*node_attrs=*/input_node,
+              /*outputs=*/
+              std::map<DynamicTensorSlot, DynamicValueAttrs>{
+                  {
+                      mk_dynamic_slot(TensorSlotName::OUTPUT),
+                      input_tensor,
+                  },
+              },
+          },
+          DynamicNodeInvocation{
+              /*inputs=*/std::map<DynamicTensorSlot, DynamicValueAttrs>{
+                  {
+                      mk_dynamic_slot(TensorSlotName::INPUT),
+                      input_tensor,
+                  },
+              },
+              /*node_attrs=*/relu1_node,
+              /*outputs=*/
+              std::map<DynamicTensorSlot, DynamicValueAttrs>{
+                  {
+                      mk_dynamic_slot(TensorSlotName::OUTPUT),
+                      relu1_output,
+                  },
+              },
+          },
+          DynamicNodeInvocation{
+              /*inputs=*/std::map<DynamicTensorSlot, DynamicValueAttrs>{
+                  {
+                      mk_dynamic_slot(TensorSlotName::INPUT),
+                      input_tensor,
+                  },
+              },
+              /*node_attrs=*/relu2_node,
+              /*outputs=*/
+              std::map<DynamicTensorSlot, DynamicValueAttrs>{
+                  {
+                      mk_dynamic_slot(TensorSlotName::OUTPUT),
+                      relu2_output,
+                  },
+              },
+          },
+      };
+
+      return dynamic_open_dataflow_graph_from_invocation_set(invocation_set);
+    }();
+
+    DynamicOpenDataflowGraph result = perform_pass_expansion(input);
+
+    debug_print_dynamic_open_dataflow_graph_as_dot(result);
+
+    DynamicOpenDataflowGraph correct = [&]() -> DynamicOpenDataflowGraph {
+      DynamicNodeAttrs input_node_fwd =
+          mk_node_attrs(10, input_op_attrs, DynamicTaskType::FWD);
+      DynamicNodeAttrs relu1_node_fwd =
+          mk_node_attrs(11, relu_op_attrs, DynamicTaskType::FWD);
+      DynamicNodeAttrs relu2_node_fwd =
+          mk_node_attrs(12, relu_op_attrs, DynamicTaskType::FWD);
+
+      DynamicNodeAttrs relu1_node_bwd =
+          mk_node_attrs(13, relu_op_attrs, DynamicTaskType::BWD);
+      DynamicNodeAttrs relu2_node_bwd =
+          mk_node_attrs(14, relu_op_attrs, DynamicTaskType::BWD);
+
+      DynamicValueAttrs input_tensor_activation =
+          mk_value_attrs(0, mk_dynamic_tensor_role_fwd());
+      DynamicValueAttrs input_tensor_gradient =
+          mk_value_attrs(0, mk_dynamic_tensor_role_bwd());
+      DynamicValueAttrs relu1_output_tensor_activation =
+          mk_value_attrs(1, mk_dynamic_tensor_role_fwd());
+      DynamicValueAttrs relu1_output_tensor_gradient =
+          mk_value_attrs(1, mk_dynamic_tensor_role_bwd());
+      DynamicValueAttrs relu2_output_tensor_activation =
+          mk_value_attrs(2, mk_dynamic_tensor_role_fwd());
+      DynamicValueAttrs relu2_output_tensor_gradient =
+          mk_value_attrs(2, mk_dynamic_tensor_role_bwd());
+
+      auto mk_fwd_slot = [&](TensorSlotName slot_name) -> DynamicTensorSlot {
+        return DynamicTensorSlot{
+            /*slot_name=*/slot_name,
+            /*slot_tensor_role=*/mk_dynamic_tensor_role_fwd(),
+            /*task_shard=*/std::nullopt,
+        };
+      };
+
+      auto mk_grad_slot = [&](TensorSlotName slot_name) -> DynamicTensorSlot {
+        return DynamicTensorSlot{
+            /*slot_name=*/slot_name,
+            /*slot_tensor_role=*/mk_dynamic_tensor_role_bwd(),
+            /*task_shard=*/std::nullopt,
+        };
+      };
+
+      std::set<DynamicNodeInvocation> invocation_set = {
+          DynamicNodeInvocation{
+              /*inputs=*/std::map<DynamicTensorSlot, DynamicValueAttrs>{},
+              /*node_attrs=*/input_node_fwd,
+              /*outputs=*/
+              std::map{
+                  std::pair{
+                      mk_fwd_slot(TensorSlotName::OUTPUT),
+                      input_tensor_activation,
+                  },
+              },
+          },
+          DynamicNodeInvocation{
+              /*inputs=*/std::map{
+                  std::pair{
+                      mk_fwd_slot(TensorSlotName::INPUT),
+                      input_tensor_activation,
+                  },
+              },
+              /*node_attrs=*/relu1_node_fwd,
+              /*outputs=*/
+              std::map{
+                  std::pair{
+                      mk_fwd_slot(TensorSlotName::OUTPUT),
+                      relu1_output_tensor_activation,
+                  },
+              },
+          },
+          DynamicNodeInvocation{
+              /*inputs=*/std::map{
+                  std::pair{
+                      mk_fwd_slot(TensorSlotName::INPUT),
+                      input_tensor_activation,
+                  },
+              },
+              /*node_attrs=*/relu2_node_fwd,
+              /*outputs=*/
+              std::map{
+                  std::pair{
+                      mk_fwd_slot(TensorSlotName::OUTPUT),
+                      relu2_output_tensor_activation,
+                  },
+              },
+          },
+          DynamicNodeInvocation{
+              /*inputs=*/std::map{
+                  std::pair{
+                      mk_fwd_slot(TensorSlotName::INPUT),
+                      input_tensor_activation,
+                  },
+                  std::pair{
+                      mk_fwd_slot(TensorSlotName::OUTPUT),
+                      relu1_output_tensor_activation,
+                  },
+                  std::pair{
+                      mk_fwd_slot(TensorSlotName::OUTPUT),
+                      relu1_output_tensor_gradient,
+                  },
+              },
+              /*node_attrs=*/relu1_node_bwd,
+              /*outputs=*/
+              std::map{
+                  std::pair{
+                      mk_fwd_slot(TensorSlotName::INPUT),
+                      input_tensor_gradient,
+                  },
+              },
+          },
+          DynamicNodeInvocation{
+              /*inputs=*/std::map{
+                  std::pair{
+                      mk_fwd_slot(TensorSlotName::INPUT),
+                      input_tensor_activation,
+                  },
+                  std::pair{
+                      mk_fwd_slot(TensorSlotName::OUTPUT),
+                      relu2_output_tensor_activation,
+                  },
+                  std::pair{
+                      mk_fwd_slot(TensorSlotName::OUTPUT),
+                      relu2_output_tensor_gradient,
+                  },
+              },
+              /*node_attrs=*/relu1_node_bwd,
+              /*outputs=*/
+              std::map{
+                  std::pair{
+                      mk_fwd_slot(TensorSlotName::INPUT),
+                      input_tensor_gradient,
+                  },
+              },
+          },
+      };
+
+      return dynamic_open_dataflow_graph_from_invocation_set(invocation_set);
+    }();
+
+    CHECK(get_dynamic_invocation_set(result).size() ==
+          correct.invocations.size());
+
+    nlohmann::json result_json =
+        dynamic_open_dataflow_graph_to_serializable(result);
+    nlohmann::json correct_json =
+        dynamic_open_dataflow_graph_to_serializable(correct);
+
+    CHECK_MESSAGE(get_dynamic_invocation_set(result) ==
+                      get_dynamic_invocation_set(correct),
+                  check_kv("result", result_json.dump()),
+                  check_kv("correct", correct_json.dump()));
+
+    CHECK(dynamic_open_dataflow_graphs_are_isomorphic(result, correct));
+  }
 }

From 237c38ddf0e74ddcfbcf751fe1c153b1080b1c3b Mon Sep 17 00:00:00 2001
From: Elliott Slaughter <slaughter@cs.stanford.edu>
Date: Thu, 20 Aug 2026 12:03:50 -0700
Subject: [PATCH 2/9] Add the subgradient ID.

---
 .../dynamic_graph/dynamic_value_attrs.dtg.toml         | 10 ++++++++++
 .../dynamic_graph/internal_dynamic_slot_site.dtg.toml  |  2 +-
 .../serializable_dynamic_value_attrs.dtg.toml          |  5 +++++
 .../src/task-spec/dynamic_graph/loss_insertion.cc      |  4 ++++
 .../make_dynamic_open_dataflow_graph_from_cg.cc        |  2 ++
 ...make_dynamic_open_dataflow_graph_from_mapped_pcg.cc |  1 +
 .../src/task-spec/dynamic_graph/pass_expansion.cc      |  1 +
 .../dynamic_graph/serializable_dynamic_value_attrs.cc  |  2 ++
 .../test/src/task-spec/dynamic_graph/copy_insertion.cc |  2 ++
 .../dynamic_graph/dynamic_open_dataflow_graph.cc       |  5 +++++
 .../src/task-spec/dynamic_graph/machine_slicing.cc     |  1 +
 ...make_dynamic_open_dataflow_graph_from_mapped_pcg.cc |  4 ++++
 .../test/src/task-spec/dynamic_graph/pass_expansion.cc |  5 +++++
 .../src/task-spec/dynamic_graph/shard_expansion.cc     |  3 +++
 .../src/task-spec/dynamic_graph/update_insertion.cc    |  4 ++++
 15 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml
index 41638fd4ef..4c9fb58fcd 100644
--- a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml
+++ b/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/internal_dynamic_slot_site.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::InternalDynamicSlotSite>"
+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 GradientReduction in the backward pass.
+'''
+
 [[fields]]
 name = "shard_coord"
 type = "std::optional<::FlexFlow::ParallelTensorSpaceCoordinate>"
diff --git a/lib/task-spec/include/task-spec/dynamic_graph/internal_dynamic_slot_site.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/internal_dynamic_slot_site.dtg.toml
index b0164a6086..ec440b5c9f 100644
--- a/lib/task-spec/include/task-spec/dynamic_graph/internal_dynamic_slot_site.dtg.toml
+++ b/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",
diff --git a/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_value_attrs.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_value_attrs.dtg.toml
index a72c674fda..d2166497c1 100644
--- a/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_value_attrs.dtg.toml
+++ b/lib/task-spec/include/task-spec/dynamic_graph/serializable_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/internal_dynamic_slot_site.dtg.h",
   "op-attrs/parallel_tensor_space_coordinate.dtg.h",
   "pcg/machine_space_coordinate.dtg.h",
   "utils/bidict/bidict.h",
@@ -37,6 +38,10 @@ type = "std::optional<::FlexFlow::ParallelTensorShape>"
 name = "create_grad"
 type = "std::optional<bool>"
 
+[[fields]]
+name = "subgradient_id"
+type = "std::optional<::FlexFlow::InternalDynamicSlotSite>"
+
 [[fields]]
 name = "shard_coord"
 type = "std::optional<::FlexFlow::ParallelTensorSpaceCoordinate>"
diff --git a/lib/task-spec/src/task-spec/dynamic_graph/loss_insertion.cc b/lib/task-spec/src/task-spec/dynamic_graph/loss_insertion.cc
index e7c4b34460..5c33c2de66 100644
--- a/lib/task-spec/src/task-spec/dynamic_graph/loss_insertion.cc
+++ b/lib/task-spec/src/task-spec/dynamic_graph/loss_insertion.cc
@@ -21,10 +21,13 @@ LossInsertionResult perform_loss_insertion(
   DynamicValueAttrs logit_value = assert_unwrap(
       find_output_value_attrs(dg, logit_tensor, mk_dynamic_tensor_role_fwd()));
 
+  ASSERT(!logit_value.subgradient_id.has_value());
+
   DynamicValueAttrs label_value{
       /*tensor_guid=*/mk_dynamic_tensor_guid_for_loss(),
       /*parallel_tensor_shape=*/logit_value.parallel_tensor_shape,
       /*create_grad=*/false,
+      /*subgradient_id=*/std::nullopt,
       /*shard_coord=*/logit_value.shard_coord,
       /*mapping=*/std::nullopt,
       /*accessor=*/std::nullopt,
@@ -35,6 +38,7 @@ LossInsertionResult perform_loss_insertion(
       /*tensor_guid=*/logit_value.tensor_guid,
       /*parallel_tensor_shape=*/logit_value.parallel_tensor_shape,
       /*create_grad=*/logit_value.create_grad,
+      /*subgradient_id=*/std::nullopt,
       /*shard_coord=*/logit_value.shard_coord,
       /*mapping=*/std::nullopt,
       /*accessor=*/std::nullopt,
diff --git a/lib/task-spec/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_cg.cc b/lib/task-spec/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_cg.cc
index 3ce603f5ff..6aa90b81a8 100644
--- a/lib/task-spec/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_cg.cc
+++ b/lib/task-spec/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_cg.cc
@@ -45,6 +45,7 @@ DynamicOpenDataflowGraph
                   /*tensor_guid=*/dynamic_tensor_guid_t{tensor},
                   /*parallel_tensor_shape=*/lift_to_parallel(attrs.shape),
                   /*create_grad=*/(attrs.create_grad == CreateGrad::YES),
+                  /*subgradient_id=*/std::nullopt,
                   /*shard_coord=*/std::nullopt,
                   /*mapping=*/std::nullopt,
                   /*accessor=*/std::nullopt,
@@ -67,6 +68,7 @@ DynamicOpenDataflowGraph
                   /*tensor_guid=*/dynamic_tensor_guid_t{tensor},
                   /*parallel_tensor_shape=*/lift_to_parallel(attrs.shape),
                   /*create_grad=*/(attrs.create_grad == CreateGrad::YES),
+                  /*subgradient_id=*/std::nullopt,
                   /*shard_coord=*/std::nullopt,
                   /*mapping=*/std::nullopt,
                   /*accessor=*/std::nullopt,
diff --git a/lib/task-spec/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.cc b/lib/task-spec/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.cc
index 13e9196416..913f295275 100644
--- a/lib/task-spec/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.cc
+++ b/lib/task-spec/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.cc
@@ -51,6 +51,7 @@ DynamicNodeInvocation make_dynamic_node_invocation_from_mapped(
             /*tensor_guid=*/dynamic_tensor_guid_t{tensor.guid},
             /*parallel_tensor_shape=*/tensor.attrs.shape,
             /*create_grad=*/(tensor.attrs.create_grad == CreateGrad::YES),
+            /*subgradient_id=*/std::nullopt,
             /*shard_coord=*/std::nullopt,
             /*mapping=*/std::nullopt,
             /*accessor=*/std::nullopt,
diff --git a/lib/task-spec/src/task-spec/dynamic_graph/pass_expansion.cc b/lib/task-spec/src/task-spec/dynamic_graph/pass_expansion.cc
index a095880118..0bd14a85a0 100644
--- a/lib/task-spec/src/task-spec/dynamic_graph/pass_expansion.cc
+++ b/lib/task-spec/src/task-spec/dynamic_graph/pass_expansion.cc
@@ -39,6 +39,7 @@ void require_value_is_pass_expanded(DynamicValueAttrs const &v) {
 
 void require_value_is_not_pass_expanded(DynamicValueAttrs const &v) {
   ASSERT(!v.role.has_value(), v);
+  ASSERT(!v.subgradient_id.has_value(), v);
 }
 
 void require_invocation_is_fully_pass_expanded(
diff --git a/lib/task-spec/src/task-spec/dynamic_graph/serializable_dynamic_value_attrs.cc b/lib/task-spec/src/task-spec/dynamic_graph/serializable_dynamic_value_attrs.cc
index ee53d01f21..0c69b850ac 100644
--- a/lib/task-spec/src/task-spec/dynamic_graph/serializable_dynamic_value_attrs.cc
+++ b/lib/task-spec/src/task-spec/dynamic_graph/serializable_dynamic_value_attrs.cc
@@ -9,6 +9,7 @@ SerializableDynamicValueAttrs
       /*tensor_guid=*/attrs.tensor_guid,
       /*parallel_tensor_shape=*/attrs.parallel_tensor_shape,
       /*create_grad=*/attrs.create_grad,
+      /*subgradient_id=*/attrs.subgradient_id,
       /*shard_coord=*/attrs.shard_coord,
       /*mapping=*/attrs.mapping,
       /*role=*/attrs.role,
@@ -21,6 +22,7 @@ DynamicValueAttrs dynamic_value_attrs_from_serializable(
       /*tensor_guid=*/attrs.tensor_guid,
       /*parallel_tensor_shape=*/attrs.parallel_tensor_shape,
       /*create_grad=*/attrs.create_grad,
+      /*subgradient_id=*/attrs.subgradient_id,
       /*shard_coord=*/attrs.shard_coord,
       /*mapping=*/attrs.mapping,
       /*accessor=*/std::nullopt,
diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/copy_insertion.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/copy_insertion.cc
index 055586664a..cfde7dff87 100644
--- a/lib/task-spec/test/src/task-spec/dynamic_graph/copy_insertion.cc
+++ b/lib/task-spec/test/src/task-spec/dynamic_graph/copy_insertion.cc
@@ -34,6 +34,7 @@ static DynamicValueAttrs
       },
       /*parallel_tensor_shape=*/std::nullopt,
       /*create_grad=*/false,
+      /*subgradient_id=*/std::nullopt,
       /*shard_coord=*/std::nullopt,
       /*mapping=*/mapping,
       /*accessor=*/std::nullopt,
@@ -792,6 +793,7 @@ TEST_SUITE(FF_TEST_SUITE) {
         },
         /*parallel_tensor_shape=*/std::nullopt,
         /*create_grad=*/std::nullopt,
+        /*subgradient_id=*/std::nullopt,
         /*shard_coord=*/std::nullopt,
         /*mapping=*/std::nullopt,
         /*accessor=*/std::nullopt,
diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/dynamic_open_dataflow_graph.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/dynamic_open_dataflow_graph.cc
index af0f799754..739d2d9547 100644
--- a/lib/task-spec/test/src/task-spec/dynamic_graph/dynamic_open_dataflow_graph.cc
+++ b/lib/task-spec/test/src/task-spec/dynamic_graph/dynamic_open_dataflow_graph.cc
@@ -23,6 +23,7 @@ TEST_SUITE(FF_TEST_SUITE) {
           }},
           /*parallel_tensor_shape=*/std::nullopt,
           /*create_grad=*/std::nullopt,
+          /*subgradient_id=*/std::nullopt,
           /*shard_coord=*/std::nullopt,
           /*mapping=*/std::nullopt,
           /*accessor=*/std::nullopt,
@@ -209,6 +210,7 @@ TEST_SUITE(FF_TEST_SUITE) {
           }},
           /*parallel_tensor_shape=*/std::nullopt,
           /*create_grad=*/std::nullopt,
+          /*subgradient_id=*/std::nullopt,
           /*shard_coord=*/std::nullopt,
           /*mapping=*/std::nullopt,
           /*accessor=*/std::nullopt,
@@ -424,6 +426,7 @@ TEST_SUITE(FF_TEST_SUITE) {
         /*tensor_guid=*/tensor_guid,
         /*parallel_tensor_shape=*/std::nullopt,
         /*create_grad=*/std::nullopt,
+        /*subgradient_id=*/std::nullopt,
         /*shard_coord=*/std::nullopt,
         /*mapping=*/std::nullopt,
         /*accessor=*/std::nullopt,
@@ -461,6 +464,7 @@ TEST_SUITE(FF_TEST_SUITE) {
         /*tensor_guid=*/tensor_guid,
         /*parallel_tensor_shape=*/std::nullopt,
         /*create_grad=*/std::nullopt,
+        /*subgradient_id=*/std::nullopt,
         /*shard_coord=*/std::nullopt,
         /*mapping=*/std::nullopt,
         /*accessor=*/std::nullopt,
@@ -478,6 +482,7 @@ TEST_SUITE(FF_TEST_SUITE) {
         /*tensor_guid=*/tensor_guid,
         /*parallel_tensor_shape=*/std::nullopt,
         /*create_grad=*/std::nullopt,
+        /*subgradient_id=*/std::nullopt,
         /*shard_coord=*/std::nullopt,
         /*mapping=*/std::nullopt,
         /*accessor=*/std::nullopt,
diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/machine_slicing.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/machine_slicing.cc
index f666473970..3eed847c37 100644
--- a/lib/task-spec/test/src/task-spec/dynamic_graph/machine_slicing.cc
+++ b/lib/task-spec/test/src/task-spec/dynamic_graph/machine_slicing.cc
@@ -77,6 +77,7 @@ TEST_SUITE(FF_TEST_SUITE) {
           }},
           /*parallel_tensor_shape=*/std::nullopt,
           /*create_grad=*/std::nullopt,
+          /*subgradient_id=*/std::nullopt,
           /*shard_coord=*/shard_coord,
           /*mapping=*/std::nullopt,
           /*accessor=*/std::nullopt,
diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.cc
index ccd8d2abfe..85bb3c21fd 100644
--- a/lib/task-spec/test/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.cc
+++ b/lib/task-spec/test/src/task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.cc
@@ -145,6 +145,7 @@ TEST_SUITE(FF_TEST_SUITE) {
                       /*tensor_guid=*/dynamic_tensor_guid_t{input_tensor_guid},
                       /*parallel_tensor_shape=*/input_shape,
                       /*create_grad=*/true,
+                      /*subgradient_id=*/std::nullopt,
                       /*shard_coord=*/std::nullopt,
                       /*mapping=*/std::nullopt,
                       /*accessor=*/std::nullopt,
@@ -177,6 +178,7 @@ TEST_SUITE(FF_TEST_SUITE) {
                       /*tensor_guid=*/dynamic_tensor_guid_t{output_tensor_guid},
                       /*parallel_tensor_shape=*/output_shape,
                       /*create_grad=*/true,
+                      /*subgradient_id=*/std::nullopt,
                       /*shard_coord=*/std::nullopt,
                       /*mapping=*/std::nullopt,
                       /*accessor=*/std::nullopt,
@@ -413,6 +415,7 @@ TEST_SUITE(FF_TEST_SUITE) {
               /*tensor_guid=*/dynamic_tensor_guid_t{tensor_guid},
               /*parallel_tensor_shape=*/shape,
               /*create_grad=*/true,
+              /*subgradient_id=*/std::nullopt,
               /*shard_coord=*/std::nullopt,
               /*mapping=*/std::nullopt,
               /*accessor=*/std::nullopt,
@@ -924,6 +927,7 @@ TEST_SUITE(FF_TEST_SUITE) {
             /*tensor_guid=*/dynamic_tensor_guid_t{tensor_guid},
             /*parallel_tensor_shape=*/shape,
             /*create_grad=*/create_grad,
+            /*subgradient_id=*/std::nullopt,
             /*shard_coord=*/std::nullopt,
             /*mapping=*/std::nullopt,
             /*accessor=*/std::nullopt,
diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc
index 731df510cc..01575e49ab 100644
--- a/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc
+++ b/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc
@@ -59,6 +59,7 @@ TEST_SUITE(FF_TEST_SUITE) {
           },
           /*parallel_tensor_shape=*/std::nullopt,
           /*create_grad=*/create_grad,
+          /*subgradient_id=*/std::nullopt,
           /*shard_coord=*/std::nullopt,
           /*mapping=*/std::nullopt,
           /*accessor=*/std::nullopt,
@@ -184,6 +185,7 @@ TEST_SUITE(FF_TEST_SUITE) {
           }},
           /*parallel_tensor_shape=*/std::nullopt,
           /*create_grad=*/std::nullopt,
+          /*subgradient_id=*/std::nullopt,
           /*shard_coord=*/std::nullopt,
           /*mapping=*/std::nullopt,
           /*accessor=*/std::nullopt,
@@ -340,6 +342,7 @@ TEST_SUITE(FF_TEST_SUITE) {
           }},
           /*parallel_tensor_shape=*/std::nullopt,
           /*create_grad=*/std::nullopt,
+          /*subgradient_id=*/std::nullopt,
           /*shard_coord=*/std::nullopt,
           /*mapping=*/std::nullopt,
           /*accessor=*/std::nullopt,
@@ -714,6 +717,7 @@ TEST_SUITE(FF_TEST_SUITE) {
           }},
           /*parallel_tensor_shape=*/std::nullopt,
           /*create_grad=*/false,
+          /*subgradient_id=*/std::nullopt,
           /*shard_coord=*/std::nullopt,
           /*mapping=*/std::nullopt,
           /*accessor=*/std::nullopt,
@@ -1051,6 +1055,7 @@ TEST_SUITE(FF_TEST_SUITE) {
           }},
           /*parallel_tensor_shape=*/std::nullopt,
           /*create_grad=*/true,
+          /*subgradient_id=*/std::nullopt,
           /*shard_coord=*/std::nullopt,
           /*mapping=*/std::nullopt,
           /*accessor=*/std::nullopt,
diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/shard_expansion.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/shard_expansion.cc
index 62437da13e..c0d9a6fd77 100644
--- a/lib/task-spec/test/src/task-spec/dynamic_graph/shard_expansion.cc
+++ b/lib/task-spec/test/src/task-spec/dynamic_graph/shard_expansion.cc
@@ -84,6 +84,7 @@ DynamicValueAttrs
       }},
       /*parallel_tensor_shape=*/std::nullopt,
       /*create_grad=*/std::nullopt,
+      /*subgradient_id=*/std::nullopt,
       /*shard_coord=*/shard_coord,
       /*mapping=*/ParallelTensorMapping{mapping},
       /*accessor=*/std::nullopt,
@@ -135,6 +136,7 @@ TEST_SUITE(FF_TEST_SUITE) {
           },
           /*parallel_tensor_shape=*/std::nullopt,
           /*create_grad=*/std::nullopt,
+          /*subgradient_id=*/std::nullopt,
           /*shard_coord=*/shard_coord,
           /*mapping=*/std::nullopt,
           /*accessor=*/std::nullopt,
@@ -599,6 +601,7 @@ TEST_SUITE(FF_TEST_SUITE) {
             }},
             /*parallel_tensor_shape=*/std::nullopt,
             /*create_grad=*/std::nullopt,
+            /*subgradient_id=*/std::nullopt,
             /*shard_coord=*/shard_coord,
             /*mapping=*/
             ParallelTensorMapping{tensor_binding},
diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/update_insertion.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/update_insertion.cc
index ae39d5afcb..5d8b59c1b4 100644
--- a/lib/task-spec/test/src/task-spec/dynamic_graph/update_insertion.cc
+++ b/lib/task-spec/test/src/task-spec/dynamic_graph/update_insertion.cc
@@ -63,6 +63,7 @@ TEST_SUITE(FF_TEST_SUITE) {
                     /*tensor_guid=*/tensor_guid,
                     /*parallel_tensor_shape=*/std::nullopt,
                     /*create_grad=*/std::nullopt,
+                    /*subgradient_id=*/std::nullopt,
                     /*shard_coord=*/std::nullopt,
                     /*mapping=*/std::nullopt,
                     /*accessor=*/std::nullopt,
@@ -104,6 +105,7 @@ TEST_SUITE(FF_TEST_SUITE) {
                                 /*tensor_guid=*/tensor_guid,
                                 /*parallel_tensor_shape=*/std::nullopt,
                                 /*create_grad=*/std::nullopt,
+                                /*subgradient_id=*/std::nullopt,
                                 /*shard_coord=*/std::nullopt,
                                 /*mapping=*/std::nullopt,
                                 /*accessor=*/std::nullopt,
@@ -121,6 +123,7 @@ TEST_SUITE(FF_TEST_SUITE) {
                                 /*tensor_guid=*/tensor_guid,
                                 /*parallel_tensor_shape=*/std::nullopt,
                                 /*create_grad=*/std::nullopt,
+                                /*subgradient_id=*/std::nullopt,
                                 /*shard_coord=*/std::nullopt,
                                 /*mapping=*/std::nullopt,
                                 /*accessor=*/std::nullopt,
@@ -139,6 +142,7 @@ TEST_SUITE(FF_TEST_SUITE) {
                                 /*tensor_guid=*/tensor_guid,
                                 /*parallel_tensor_shape=*/std::nullopt,
                                 /*create_grad=*/std::nullopt,
+                                /*subgradient_id=*/std::nullopt,
                                 /*shard_coord=*/std::nullopt,
                                 /*mapping=*/std::nullopt,
                                 /*accessor=*/std::nullopt,

From 40004997fc1c13f4d42406cd21364a3ecd69c10b Mon Sep 17 00:00:00 2001
From: Elliott Slaughter <slaughter@cs.stanford.edu>
Date: Thu, 20 Aug 2026 15:03:39 -0700
Subject: [PATCH 3/9] Tweak subgradient ID representation.

---
 .../dynamic_value_attrs.dtg.toml              |  4 +--
 .../dynamic_graph/gradient_reduction.dtg.toml | 13 ++++++++++
 .../serializable_dynamic_value_attrs.dtg.toml |  4 +--
 .../dynamic_graph/subgradient_id_t.dtg.toml   | 26 +++++++++++++++++++
 4 files changed, 43 insertions(+), 4 deletions(-)
 create mode 100644 lib/task-spec/include/task-spec/dynamic_graph/gradient_reduction.dtg.toml
 create mode 100644 lib/task-spec/include/task-spec/dynamic_graph/subgradient_id_t.dtg.toml

diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml
index 4c9fb58fcd..bd22171aa3 100644
--- a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml
+++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml
@@ -13,7 +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/internal_dynamic_slot_site.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",
@@ -51,7 +51,7 @@ type = "std::optional<bool>"
 
 [[fields]]
 name = "subgradient_id"
-type = "std::optional<::FlexFlow::InternalDynamicSlotSite>"
+type = "std::optional<::FlexFlow::subgradient_id_t>"
 docstring = '''
 \brief The unique identity of the subgradient represented by this value.
 
diff --git a/lib/task-spec/include/task-spec/dynamic_graph/gradient_reduction.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/gradient_reduction.dtg.toml
new file mode 100644
index 0000000000..82f759fd19
--- /dev/null
+++ b/lib/task-spec/include/task-spec/dynamic_graph/gradient_reduction.dtg.toml
@@ -0,0 +1,13 @@
+namespace = "FlexFlow"
+name = "GradientReductionAttrs"
+type = "struct"
+features = [
+  "eq",
+  "ord",
+  "hash",
+  "json",
+  "fmt",
+  "rapidcheck",
+]
+
+fields = []
diff --git a/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_value_attrs.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_value_attrs.dtg.toml
index d2166497c1..2d7cf335b9 100644
--- a/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_value_attrs.dtg.toml
+++ b/lib/task-spec/include/task-spec/dynamic_graph/serializable_dynamic_value_attrs.dtg.toml
@@ -13,7 +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/internal_dynamic_slot_site.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",
@@ -40,7 +40,7 @@ type = "std::optional<bool>"
 
 [[fields]]
 name = "subgradient_id"
-type = "std::optional<::FlexFlow::InternalDynamicSlotSite>"
+type = "std::optional<::FlexFlow::subgradient_id_t>"
 
 [[fields]]
 name = "shard_coord"
diff --git a/lib/task-spec/include/task-spec/dynamic_graph/subgradient_id_t.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/subgradient_id_t.dtg.toml
new file mode 100644
index 0000000000..a9e0db7054
--- /dev/null
+++ b/lib/task-spec/include/task-spec/dynamic_graph/subgradient_id_t.dtg.toml
@@ -0,0 +1,26 @@
+namespace = "FlexFlow"
+name = "subgradient_id_t"
+type = "struct"
+features = [
+  "eq",
+  "ord",
+  "hash",
+  "json",
+  "fmt",
+]
+
+docstring = """
+\brief The unique identity of a subgradient consumed by \ref GradientReduction.
+
+This field is filled in by \ref pass_expansion.h when a value is consumed multiple times, requiring a \ref GradientReduction in the backward pass. The subgradient ID is *only locally unique* among the inputs to a given \ref GradientReduction, and nothing else should be assumed about it. It is currently implemented by taking the \ref DynamicTensorSlot as it is used as an input to \ref GradientReduction.
+"""
+
+includes = [
+  "task-spec/dynamic_graph/dynamic_tensor_slot.dtg.h",
+]
+
+src_includes = []
+
+[[fields]]
+name = "gradient_reduction_slot"
+type = "::FlexFlow::DynamicTensorSlot"

From f60eac77e170374814d069ee7903daa7c21df1d6 Mon Sep 17 00:00:00 2001
From: Elliott Slaughter <slaughter@cs.stanford.edu>
Date: Thu, 20 Aug 2026 17:10:41 -0700
Subject: [PATCH 4/9] More work on gradient reduction.

---
 .../src/realm-execution/pcg_instance.cc       |   2 +
 ...c_gradient_reduction_layer_guid_t.dtg.toml |  13 ++
 .../dynamic_layer_guid_t.dtg.toml             |   5 +
 .../dynamic_graph/subgradient_id_t.dtg.toml   |   6 +-
 .../training_only_op_type.dtg.toml            |   3 +
 .../training_operation_attrs.dtg.toml         |   5 +
 .../dynamic_invocation_id_t.dtg.toml          |  20 ++
 .../task-spec/dynamic_graph/pass_expansion.cc | 190 ++++++++++++++++--
 .../dynamic_graph/training_operation_attrs.cc |   5 +
 .../task-spec/dynamic_graph/pass_expansion.cc |  80 ++++++--
 10 files changed, 296 insertions(+), 33 deletions(-)
 create mode 100644 lib/task-spec/include/task-spec/dynamic_graph/dynamic_gradient_reduction_layer_guid_t.dtg.toml
 create mode 100644 lib/task-spec/src/task-spec/dynamic_graph/dynamic_invocation_id_t.dtg.toml

diff --git a/lib/realm-execution/src/realm-execution/pcg_instance.cc b/lib/realm-execution/src/realm-execution/pcg_instance.cc
index 1d102f8cca..076ec2d740 100644
--- a/lib/realm-execution/src/realm-execution/pcg_instance.cc
+++ b/lib/realm-execution/src/realm-execution/pcg_instance.cc
@@ -14,6 +14,7 @@
 #include "task-spec/dynamic_graph/dynamic_task_type.dtg.h"
 #include "task-spec/dynamic_graph/dynamic_tensor_guid_t.dtg.h"
 #include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h"
+#include "task-spec/dynamic_graph/gradient_reduction.dtg.h"
 #include "task-spec/dynamic_graph/loss_insertion.h"
 #include "task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.h"
 #include "task-spec/dynamic_graph/pass_expansion.h"
@@ -341,6 +342,7 @@ static Realm::Event spawn_dynamic_node_invocation(
       },
       [&](LossAttrs const &) { return spawn_task(); },
       [&](CopyAttrs const &) { return issue_copy(); },
+      [&](GradientReductionAttrs const &) { return issue_reduction(); },
   });
 }
 
diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_gradient_reduction_layer_guid_t.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_gradient_reduction_layer_guid_t.dtg.toml
new file mode 100644
index 0000000000..bcaf9f3bbf
--- /dev/null
+++ b/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 = []
diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_layer_guid_t.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_layer_guid_t.dtg.toml
index 5200bfc6a6..ef37911069 100644
--- a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_layer_guid_t.dtg.toml
+++ b/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"
diff --git a/lib/task-spec/include/task-spec/dynamic_graph/subgradient_id_t.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/subgradient_id_t.dtg.toml
index a9e0db7054..61fd935996 100644
--- a/lib/task-spec/include/task-spec/dynamic_graph/subgradient_id_t.dtg.toml
+++ b/lib/task-spec/include/task-spec/dynamic_graph/subgradient_id_t.dtg.toml
@@ -12,15 +12,15 @@ features = [
 docstring = """
 \brief The unique identity of a subgradient consumed by \ref GradientReduction.
 
-This field is filled in by \ref pass_expansion.h when a value is consumed multiple times, requiring a \ref GradientReduction in the backward pass. The subgradient ID is *only locally unique* among the inputs to a given \ref GradientReduction, and nothing else should be assumed about it. It is currently implemented by taking the \ref DynamicTensorSlot as it is used as an input to \ref GradientReduction.
+This field is filled in by \ref pass_expansion.h when a value is consumed multiple times, requiring a \ref GradientReduction in the backward pass. The subgradient ID is *only locally unique* among the inputs to a given \ref GradientReduction, and nothing else should be assumed about it. It is currently implemented by taking the \ref TensorSlotName as it is used as an input to \ref GradientReduction.
 """
 
 includes = [
-  "task-spec/dynamic_graph/dynamic_tensor_slot.dtg.h",
+  "op-attrs/tensor_slot_name.dtg.h",
 ]
 
 src_includes = []
 
 [[fields]]
 name = "gradient_reduction_slot"
-type = "::FlexFlow::DynamicTensorSlot"
+type = "::FlexFlow::TensorSlotName"
diff --git a/lib/task-spec/include/task-spec/dynamic_graph/training_only_op_type.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/training_only_op_type.dtg.toml
index f1013771d1..1d20a99288 100644
--- a/lib/task-spec/include/task-spec/dynamic_graph/training_only_op_type.dtg.toml
+++ b/lib/task-spec/include/task-spec/dynamic_graph/training_only_op_type.dtg.toml
@@ -11,5 +11,8 @@ features = [
 [[values]]
 name = "COPY"
 
+[[values]]
+name = "GRADIENT_REDUCTION"
+
 [[values]]
 name = "LOSS"
diff --git a/lib/task-spec/include/task-spec/dynamic_graph/training_operation_attrs.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/training_operation_attrs.dtg.toml
index 2c4e739571..2381a3db3f 100644
--- a/lib/task-spec/include/task-spec/dynamic_graph/training_operation_attrs.dtg.toml
+++ b/lib/task-spec/include/task-spec/dynamic_graph/training_operation_attrs.dtg.toml
@@ -13,6 +13,7 @@ includes = [
   "op-attrs/ops/loss_functions/loss_attrs.dtg.h",
   "op-attrs/pcg_operator_attrs.dtg.h",
   "task-spec/dynamic_graph/copy_attrs.dtg.h",
+  "task-spec/dynamic_graph/gradient_reduction.dtg.h",
 ]
 
 [[values]]
@@ -26,3 +27,7 @@ key = "loss"
 [[values]]
 type = "::FlexFlow::CopyAttrs"
 key = "copy"
+
+[[values]]
+type = "::FlexFlow::GradientReductionAttrs"
+key = "gradient_reduction"
diff --git a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_invocation_id_t.dtg.toml b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_invocation_id_t.dtg.toml
new file mode 100644
index 0000000000..f600824ab3
--- /dev/null
+++ b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_invocation_id_t.dtg.toml
@@ -0,0 +1,20 @@
+namespace = "FlexFlow"
+name = "dynamic_invocation_id_t"
+type = "struct"
+features = [
+  "eq",
+  "ord",
+  "hash",
+  "json",
+  "fmt",
+]
+
+includes = [
+  "utils/nonnegative_int/nonnegative_int.h",
+]
+
+src_includes = []
+
+[[fields]]
+name = "idx"
+type = "::FlexFlow::nonnegative_int"
diff --git a/lib/task-spec/src/task-spec/dynamic_graph/pass_expansion.cc b/lib/task-spec/src/task-spec/dynamic_graph/pass_expansion.cc
index 0bd14a85a0..cedf51bed0 100644
--- a/lib/task-spec/src/task-spec/dynamic_graph/pass_expansion.cc
+++ b/lib/task-spec/src/task-spec/dynamic_graph/pass_expansion.cc
@@ -1,15 +1,33 @@
 #include "task-spec/dynamic_graph/pass_expansion.h"
+#include "op-attrs/tensor_slot_name.dtg.h"
+#include "task-spec/dynamic_graph/dynamic_node_attrs.dtg.h"
+#include "task-spec/dynamic_graph/dynamic_node_invocation.dtg.h"
 #include "task-spec/dynamic_graph/dynamic_node_invocation.h"
 #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.h"
 #include "task-spec/dynamic_graph/dynamic_tensor_role.h"
+#include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h"
+#include "task-spec/dynamic_graph/subgradient_id_t.dtg.h"
 #include "task-spec/dynamic_graph/training_operation_attrs.h"
 #include "utils/containers/are_all_same.h"
+#include "utils/containers/concat_vectors.h"
+#include "utils/containers/contains_duplicates.h"
+#include "utils/containers/filter.h"
 #include "utils/containers/flatmap.h"
+#include "utils/containers/generate_map.h"
 #include "utils/containers/get_only.h"
+#include "utils/containers/map_from_pairs.h"
 #include "utils/containers/map_values.h"
+#include "utils/containers/map_values2.h"
 #include "utils/containers/merge_disjoint_maps.h"
+#include "utils/containers/multiset_of.h"
+#include "utils/containers/range.h"
 #include "utils/containers/repeat_until_converged.h"
+#include "utils/containers/set_of.h"
 #include "utils/containers/transform.h"
+#include "utils/containers/zip_with.h"
+#include "utils/optional.h"
+#include <optional>
+#include <set>
 
 namespace FlexFlow {
 
@@ -27,6 +45,7 @@ void require_node_might_not_be_pass_expanded(DynamicNodeAttrs const &n) {
   }
 
   ASSERT(!n.task_type.has_value(), n);
+  ASSERT(!assert_unwrap(n.op_attrs).is_gradient_reduction());
 }
 
 void require_slot_is_not_pass_expanded(DynamicTensorSlot const &s) {
@@ -325,6 +344,144 @@ DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation(
   return result;
 }
 
+static std::map<int, subgradient_id_t>
+    choose_subgradient_ids(std::set<int> const &invocations) {
+  const std::vector<TensorSlotName> slot_names = {TensorSlotName::INPUT_00,
+                                                  TensorSlotName::INPUT_01,
+                                                  TensorSlotName::INPUT_02,
+                                                  TensorSlotName::INPUT_03,
+                                                  TensorSlotName::INPUT_04,
+                                                  TensorSlotName::INPUT_05,
+                                                  TensorSlotName::INPUT_06,
+                                                  TensorSlotName::INPUT_07,
+                                                  TensorSlotName::INPUT_08,
+                                                  TensorSlotName::INPUT_09,
+                                                  TensorSlotName::INPUT_10,
+                                                  TensorSlotName::INPUT_11,
+                                                  TensorSlotName::INPUT_12,
+                                                  TensorSlotName::INPUT_13,
+                                                  TensorSlotName::INPUT_14,
+                                                  TensorSlotName::INPUT_15};
+  std::map<int, subgradient_id_t> result = map_from_pairs(
+      zip_with(vector_of(invocations),
+               slot_names,
+               [](int idx, TensorSlotName slot_name) {
+                 return std::pair{idx, subgradient_id_t{slot_name}};
+               }));
+  ASSERT(result.size() == invocations.size());
+  return result;
+}
+
+static std::vector<DynamicNodeInvocation>
+    reduce_gradients(std::vector<DynamicNodeInvocation> const &invocations) {
+  std::multiset<DynamicValueAttrs> outputs =
+      multiset_of(flatmap(invocations, [](DynamicNodeInvocation const &i) {
+        return vector_of(values(i.outputs));
+      }));
+  std::set<DynamicValueAttrs> unique_outputs = set_of(outputs);
+
+  std::map<DynamicNodeInvocation, int> invocation_index =
+      map_from_keys_and_values(invocations, range(invocations.size()));
+
+  std::map<DynamicValueAttrs, std::set<int>> output_invocation_set;
+  for (int idx : range(invocations.size())) {
+    DynamicNodeInvocation const &i = invocations[idx];
+    for (DynamicValueAttrs const &v : values(i.outputs)) {
+      if (outputs.count(v) > 1) {
+        output_invocation_set[v].insert(idx);
+      }
+    }
+  }
+  std::map<DynamicValueAttrs, std::map<int, subgradient_id_t>>
+      output_subgradient_ids =
+          map_values(output_invocation_set, choose_subgradient_ids);
+
+  std::map<DynamicValueAttrs, DynamicNodeInvocation> gradient_reductions =
+      map_values2(output_invocation_set,
+                  [&](DynamicValueAttrs const &output,
+                      std::set<int> const &invocation_set) {
+                    return DynamicNodeInvocation{
+                        /*inputs=*/map_from_pairs(transform(
+                            values(output_subgradient_ids.at(output)),
+                            [&](subgradient_id_t id) {
+                              DynamicTensorSlot slot{
+                                  /*slot_name=*/id.gradient_reduction_slot,
+                                  /*slot_tensor_role=*/
+                                  DynamicTensorRole{FwbTensorType::GRADIENT},
+                                  /*task_shard=*/std::nullopt,
+                              };
+
+                              DynamicValueAttrs input = output;
+                              input.subgradient_id = id;
+
+                              return std::pair{slot, input};
+                            })),
+                        /*node_attrs=*/
+                        DynamicNodeAttrs{
+                            /*task_type=*/DynamicTaskType::BWD,
+                            /*device_ids=*/std::nullopt,
+                            /*mapping=*/std::nullopt,
+                            /*op_attrs=*/
+                            TrainingOperationAttrs{GradientReductionAttrs{}},
+                            /*layer_guid=*/
+                            dynamic_layer_guid_t{
+                                dynamic_gradient_reduction_layer_guid_t{}},
+                            /*per_device_op_state=*/std::nullopt,
+                        },
+                        /*outputs=*/
+                        std::map<DynamicTensorSlot, DynamicValueAttrs>{
+                            {
+                                DynamicTensorSlot{
+                                    /*slot_name=*/TensorSlotName::OUTPUT,
+                                    /*slot_tensor_role=*/
+                                    DynamicTensorRole{FwbTensorType::GRADIENT},
+                                    /*task_shard=*/std::nullopt,
+                                },
+                                output,
+                            },
+                        },
+                    };
+                  });
+
+  return concat_vectors(
+      transform(range(invocations.size()),
+                [&](int idx) {
+                  DynamicNodeInvocation mapped_invocation = invocations.at(idx);
+                  mapped_invocation.outputs = map_values(
+                      mapped_invocation.outputs,
+                      [&](DynamicValueAttrs const &output) {
+                        DynamicValueAttrs mapped_output = output;
+                        if (output_subgradient_ids.count(output)) {
+                          mapped_output.subgradient_id =
+                              output_subgradient_ids.at(output).at(idx);
+                        }
+                        return mapped_output;
+                      });
+                  return mapped_invocation;
+                }),
+      vector_of(values(gradient_reductions)));
+}
+
+// Like flatmap_dynamic_invocation_set except we replace the creation of
+// duplicate values with GradientReduction
+static DynamicOpenDataflowGraph
+    flatmap_dynamic_invocation_set_with_gradient_reduction(
+        DynamicOpenDataflowGraph const &g,
+        std::function<std::set<DynamicNodeInvocation>(
+            DynamicNodeInvocation const &)> const &f) {
+  std::set<DynamicNodeInvocation> current_invocation_set =
+      get_dynamic_invocation_set(g);
+  std::vector<DynamicNodeInvocation> new_invocation_set =
+      flatmap(vector_of(current_invocation_set), f);
+
+  ASSERT(!contains_duplicates(new_invocation_set));
+
+  new_invocation_set = reduce_gradients(new_invocation_set);
+
+  return dynamic_open_dataflow_graph_from_invocation_set(
+      set_of(new_invocation_set));
+}
+
 DynamicOpenDataflowGraph
     perform_pass_expansion(DynamicOpenDataflowGraph const &g) {
 
@@ -333,22 +490,23 @@ DynamicOpenDataflowGraph
   std::set<dynamic_invocation_id_t> needed_in_bwd_pass =
       determine_invocations_needed_in_backward_pass_for_gradient_computation(g);
 
-  DynamicOpenDataflowGraph result = flatmap_dynamic_invocation_set(
-      g, [&](DynamicNodeInvocation const &invocation) {
-        dynamic_invocation_id_t invocation_id =
-            dynamic_graph_get_id_for_invocation(g, invocation);
-
-        if (contains(needed_in_bwd_pass, invocation_id)) {
-          return std::set{
-              perform_fwd_pass_expansion_for_invocation(invocation),
-              perform_bwd_pass_expansion_for_invocation(invocation),
-          };
-        } else {
-          return std::set{
-              perform_fwd_pass_expansion_for_invocation(invocation),
-          };
-        }
-      });
+  DynamicOpenDataflowGraph result =
+      flatmap_dynamic_invocation_set_with_gradient_reduction(
+          g, [&](DynamicNodeInvocation const &invocation) {
+            dynamic_invocation_id_t invocation_id =
+                dynamic_graph_get_id_for_invocation(g, invocation);
+
+            if (contains(needed_in_bwd_pass, invocation_id)) {
+              return std::set{
+                  perform_fwd_pass_expansion_for_invocation(invocation),
+                  perform_bwd_pass_expansion_for_invocation(invocation),
+              };
+            } else {
+              return std::set{
+                  perform_fwd_pass_expansion_for_invocation(invocation),
+              };
+            }
+          });
 
   require_graph_is_fully_pass_expanded(result);
 
diff --git a/lib/task-spec/src/task-spec/dynamic_graph/training_operation_attrs.cc b/lib/task-spec/src/task-spec/dynamic_graph/training_operation_attrs.cc
index 3bdfe0d30a..ec6738a635 100644
--- a/lib/task-spec/src/task-spec/dynamic_graph/training_operation_attrs.cc
+++ b/lib/task-spec/src/task-spec/dynamic_graph/training_operation_attrs.cc
@@ -1,5 +1,6 @@
 #include "task-spec/dynamic_graph/training_operation_attrs.h"
 #include "op-attrs/pcg_operator_attrs.h"
+#include "task-spec/dynamic_graph/gradient_reduction.dtg.h"
 #include "utils/overload.h"
 
 namespace FlexFlow {
@@ -12,6 +13,7 @@ bool training_op_attrs_has_op_type(TrainingOperationAttrs const &op_attrs,
       },
       [](LossAttrs const &) -> bool { return false; },
       [](CopyAttrs const &) -> bool { return false; },
+      [](GradientReductionAttrs const &) -> bool { return false; },
   });
 }
 
@@ -29,6 +31,9 @@ TrainingOpType training_op_attrs_get_op_type(
       [](CopyAttrs const &) -> TrainingOpType {
         return TrainingOpType{TrainingOnlyOpType::COPY};
       },
+      [](GradientReductionAttrs const &) -> TrainingOpType {
+        return TrainingOpType{TrainingOnlyOpType::GRADIENT_REDUCTION};
+      },
   });
 }
 
diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc
index 01575e49ab..64de063bc9 100644
--- a/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc
+++ b/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc
@@ -4,10 +4,15 @@
 #include "op-attrs/ops/element_binary_attrs.dtg.h"
 #include "op-attrs/ops/element_unary.h"
 #include "op-attrs/tensor_slot_name.dtg.h"
+#include "task-spec/dynamic_graph/dynamic_gradient_reduction_layer_guid_t.dtg.h"
+#include "task-spec/dynamic_graph/dynamic_layer_guid_t.dtg.h"
 #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.h"
 #include "task-spec/dynamic_graph/dynamic_tensor_role.h"
+#include "task-spec/dynamic_graph/gradient_reduction.dtg.h"
 #include "task-spec/dynamic_graph/serializable_dynamic_node_invocation.h"
 #include "task-spec/dynamic_graph/serializable_dynamic_open_dataflow_graph.h"
+#include "task-spec/dynamic_graph/subgradient_id_t.dtg.h"
+#include "task-spec/dynamic_graph/training_operation_attrs.dtg.h"
 #include "test/utils/doctest/check_kv.h"
 #include <doctest/doctest.h>
 
@@ -1043,8 +1048,24 @@ TEST_SUITE(FF_TEST_SUITE) {
       };
     };
 
+    auto mk_gradient_node_attrs =
+        [](std::optional<DynamicTaskType> const &pass_type)
+        -> DynamicNodeAttrs {
+      return DynamicNodeAttrs{
+          /*pass_type=*/pass_type,
+          /*device_coord=*/std::nullopt,
+          /*mapping=*/std::nullopt,
+          /*op_attrs=*/TrainingOperationAttrs{GradientReductionAttrs{}},
+          /*layer_guid=*/
+          dynamic_layer_guid_t{dynamic_gradient_reduction_layer_guid_t{}},
+          /*per_device_op_state=*/std::nullopt,
+      };
+    };
+
     auto mk_value_attrs =
-        [](size_t node_id, std::optional<DynamicTensorRole> const &tensor_type)
+        [](size_t node_id,
+           std::optional<DynamicTensorRole> const &tensor_type,
+           std::optional<subgradient_id_t> const &subgradient_id)
         -> DynamicValueAttrs {
       return DynamicValueAttrs{
           /*tensor_guid=*/dynamic_tensor_guid_t{parallel_tensor_guid_t{
@@ -1055,7 +1076,7 @@ TEST_SUITE(FF_TEST_SUITE) {
           }},
           /*parallel_tensor_shape=*/std::nullopt,
           /*create_grad=*/true,
-          /*subgradient_id=*/std::nullopt,
+          /*subgradient_id=*/subgradient_id,
           /*shard_coord=*/std::nullopt,
           /*mapping=*/std::nullopt,
           /*accessor=*/std::nullopt,
@@ -1093,9 +1114,12 @@ TEST_SUITE(FF_TEST_SUITE) {
       DynamicNodeAttrs relu2_node =
           mk_node_attrs(12, relu_op_attrs, std::nullopt);
 
-      DynamicValueAttrs input_tensor = mk_value_attrs(0, std::nullopt);
-      DynamicValueAttrs relu1_output = mk_value_attrs(1, std::nullopt);
-      DynamicValueAttrs relu2_output = mk_value_attrs(2, std::nullopt);
+      DynamicValueAttrs input_tensor =
+          mk_value_attrs(0, std::nullopt, std::nullopt);
+      DynamicValueAttrs relu1_output =
+          mk_value_attrs(1, std::nullopt, std::nullopt);
+      DynamicValueAttrs relu2_output =
+          mk_value_attrs(2, std::nullopt, std::nullopt);
 
       auto mk_dynamic_slot =
           [](TensorSlotName const &slot_name) -> DynamicTensorSlot {
@@ -1157,8 +1181,6 @@ TEST_SUITE(FF_TEST_SUITE) {
 
     DynamicOpenDataflowGraph result = perform_pass_expansion(input);
 
-    debug_print_dynamic_open_dataflow_graph_as_dot(result);
-
     DynamicOpenDataflowGraph correct = [&]() -> DynamicOpenDataflowGraph {
       DynamicNodeAttrs input_node_fwd =
           mk_node_attrs(10, input_op_attrs, DynamicTaskType::FWD);
@@ -1171,19 +1193,29 @@ TEST_SUITE(FF_TEST_SUITE) {
           mk_node_attrs(13, relu_op_attrs, DynamicTaskType::BWD);
       DynamicNodeAttrs relu2_node_bwd =
           mk_node_attrs(14, relu_op_attrs, DynamicTaskType::BWD);
+      DynamicNodeAttrs gradient_reduction_node_bwd =
+          mk_gradient_node_attrs(DynamicTaskType::BWD);
 
       DynamicValueAttrs input_tensor_activation =
-          mk_value_attrs(0, mk_dynamic_tensor_role_fwd());
+          mk_value_attrs(0, mk_dynamic_tensor_role_fwd(), std::nullopt);
+      DynamicValueAttrs input_tensor_subgradient1 =
+          mk_value_attrs(0,
+                         mk_dynamic_tensor_role_bwd(),
+                         subgradient_id_t{TensorSlotName::INPUT_00});
+      DynamicValueAttrs input_tensor_subgradient2 =
+          mk_value_attrs(0,
+                         mk_dynamic_tensor_role_bwd(),
+                         subgradient_id_t{TensorSlotName::INPUT_01});
       DynamicValueAttrs input_tensor_gradient =
-          mk_value_attrs(0, mk_dynamic_tensor_role_bwd());
+          mk_value_attrs(0, mk_dynamic_tensor_role_bwd(), std::nullopt);
       DynamicValueAttrs relu1_output_tensor_activation =
-          mk_value_attrs(1, mk_dynamic_tensor_role_fwd());
+          mk_value_attrs(1, mk_dynamic_tensor_role_fwd(), std::nullopt);
       DynamicValueAttrs relu1_output_tensor_gradient =
-          mk_value_attrs(1, mk_dynamic_tensor_role_bwd());
+          mk_value_attrs(1, mk_dynamic_tensor_role_bwd(), std::nullopt);
       DynamicValueAttrs relu2_output_tensor_activation =
-          mk_value_attrs(2, mk_dynamic_tensor_role_fwd());
+          mk_value_attrs(2, mk_dynamic_tensor_role_fwd(), std::nullopt);
       DynamicValueAttrs relu2_output_tensor_gradient =
-          mk_value_attrs(2, mk_dynamic_tensor_role_bwd());
+          mk_value_attrs(2, mk_dynamic_tensor_role_bwd(), std::nullopt);
 
       auto mk_fwd_slot = [&](TensorSlotName slot_name) -> DynamicTensorSlot {
         return DynamicTensorSlot{
@@ -1265,7 +1297,7 @@ TEST_SUITE(FF_TEST_SUITE) {
               std::map{
                   std::pair{
                       mk_fwd_slot(TensorSlotName::INPUT),
-                      input_tensor_gradient,
+                      input_tensor_subgradient1,
                   },
               },
           },
@@ -1289,6 +1321,26 @@ TEST_SUITE(FF_TEST_SUITE) {
               std::map{
                   std::pair{
                       mk_fwd_slot(TensorSlotName::INPUT),
+                      input_tensor_subgradient2,
+                  },
+              },
+          },
+          DynamicNodeInvocation{
+              /*inputs=*/std::map{
+                  std::pair{
+                      mk_grad_slot(TensorSlotName::INPUT_00),
+                      input_tensor_subgradient1,
+                  },
+                  std::pair{
+                      mk_grad_slot(TensorSlotName::INPUT_01),
+                      input_tensor_subgradient2,
+                  },
+              },
+              /*node_attrs=*/gradient_reduction_node_bwd,
+              /*outputs=*/
+              std::map{
+                  std::pair{
+                      mk_grad_slot(TensorSlotName::OUTPUT),
                       input_tensor_gradient,
                   },
               },

From ced5738bf7d313268e066356ea40bce3c37b1f5c Mon Sep 17 00:00:00 2001
From: Elliott Slaughter <slaughter@cs.stanford.edu>
Date: Thu, 20 Aug 2026 17:25:22 -0700
Subject: [PATCH 5/9] Fix up test case.

---
 .../src/task-spec/dynamic_graph/pass_expansion.cc  | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc
index 64de063bc9..c8e553cce5 100644
--- a/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc
+++ b/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc
@@ -1190,9 +1190,9 @@ TEST_SUITE(FF_TEST_SUITE) {
           mk_node_attrs(12, relu_op_attrs, DynamicTaskType::FWD);
 
       DynamicNodeAttrs relu1_node_bwd =
-          mk_node_attrs(13, relu_op_attrs, DynamicTaskType::BWD);
+          mk_node_attrs(11, relu_op_attrs, DynamicTaskType::BWD);
       DynamicNodeAttrs relu2_node_bwd =
-          mk_node_attrs(14, relu_op_attrs, DynamicTaskType::BWD);
+          mk_node_attrs(12, relu_op_attrs, DynamicTaskType::BWD);
       DynamicNodeAttrs gradient_reduction_node_bwd =
           mk_gradient_node_attrs(DynamicTaskType::BWD);
 
@@ -1288,7 +1288,7 @@ TEST_SUITE(FF_TEST_SUITE) {
                       relu1_output_tensor_activation,
                   },
                   std::pair{
-                      mk_fwd_slot(TensorSlotName::OUTPUT),
+                      mk_grad_slot(TensorSlotName::OUTPUT),
                       relu1_output_tensor_gradient,
                   },
               },
@@ -1296,7 +1296,7 @@ TEST_SUITE(FF_TEST_SUITE) {
               /*outputs=*/
               std::map{
                   std::pair{
-                      mk_fwd_slot(TensorSlotName::INPUT),
+                      mk_grad_slot(TensorSlotName::INPUT),
                       input_tensor_subgradient1,
                   },
               },
@@ -1312,15 +1312,15 @@ TEST_SUITE(FF_TEST_SUITE) {
                       relu2_output_tensor_activation,
                   },
                   std::pair{
-                      mk_fwd_slot(TensorSlotName::OUTPUT),
+                      mk_grad_slot(TensorSlotName::OUTPUT),
                       relu2_output_tensor_gradient,
                   },
               },
-              /*node_attrs=*/relu1_node_bwd,
+              /*node_attrs=*/relu2_node_bwd,
               /*outputs=*/
               std::map{
                   std::pair{
-                      mk_fwd_slot(TensorSlotName::INPUT),
+                      mk_grad_slot(TensorSlotName::INPUT),
                       input_tensor_subgradient2,
                   },
               },

From 8c0a147efdd20c266490e2a5f42ba871310686c8 Mon Sep 17 00:00:00 2001
From: Elliott Slaughter <slaughter@cs.stanford.edu>
Date: Thu, 20 Aug 2026 17:36:06 -0700
Subject: [PATCH 6/9] Cleanup.

---
 .../src/realm-execution/pcg_instance.cc       |   2 +-
 ...toml => gradient_reduction_attrs.dtg.toml} |   0
 .../training_operation_attrs.dtg.toml         |   2 +-
 .../dynamic_graph/training_operation_attrs.cc |   1 -
 .../task-spec/dynamic_graph/pass_expansion.cc | 129 +-----------------
 5 files changed, 3 insertions(+), 131 deletions(-)
 rename lib/task-spec/include/task-spec/dynamic_graph/{gradient_reduction.dtg.toml => gradient_reduction_attrs.dtg.toml} (100%)

diff --git a/lib/realm-execution/src/realm-execution/pcg_instance.cc b/lib/realm-execution/src/realm-execution/pcg_instance.cc
index 076ec2d740..458e0e42d1 100644
--- a/lib/realm-execution/src/realm-execution/pcg_instance.cc
+++ b/lib/realm-execution/src/realm-execution/pcg_instance.cc
@@ -14,7 +14,7 @@
 #include "task-spec/dynamic_graph/dynamic_task_type.dtg.h"
 #include "task-spec/dynamic_graph/dynamic_tensor_guid_t.dtg.h"
 #include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h"
-#include "task-spec/dynamic_graph/gradient_reduction.dtg.h"
+#include "task-spec/dynamic_graph/gradient_reduction_attrs.dtg.h"
 #include "task-spec/dynamic_graph/loss_insertion.h"
 #include "task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.h"
 #include "task-spec/dynamic_graph/pass_expansion.h"
diff --git a/lib/task-spec/include/task-spec/dynamic_graph/gradient_reduction.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/gradient_reduction_attrs.dtg.toml
similarity index 100%
rename from lib/task-spec/include/task-spec/dynamic_graph/gradient_reduction.dtg.toml
rename to lib/task-spec/include/task-spec/dynamic_graph/gradient_reduction_attrs.dtg.toml
diff --git a/lib/task-spec/include/task-spec/dynamic_graph/training_operation_attrs.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/training_operation_attrs.dtg.toml
index 2381a3db3f..174d04f794 100644
--- a/lib/task-spec/include/task-spec/dynamic_graph/training_operation_attrs.dtg.toml
+++ b/lib/task-spec/include/task-spec/dynamic_graph/training_operation_attrs.dtg.toml
@@ -13,7 +13,7 @@ includes = [
   "op-attrs/ops/loss_functions/loss_attrs.dtg.h",
   "op-attrs/pcg_operator_attrs.dtg.h",
   "task-spec/dynamic_graph/copy_attrs.dtg.h",
-  "task-spec/dynamic_graph/gradient_reduction.dtg.h",
+  "task-spec/dynamic_graph/gradient_reduction_attrs.dtg.h",
 ]
 
 [[values]]
diff --git a/lib/task-spec/src/task-spec/dynamic_graph/training_operation_attrs.cc b/lib/task-spec/src/task-spec/dynamic_graph/training_operation_attrs.cc
index ec6738a635..ca437db9c4 100644
--- a/lib/task-spec/src/task-spec/dynamic_graph/training_operation_attrs.cc
+++ b/lib/task-spec/src/task-spec/dynamic_graph/training_operation_attrs.cc
@@ -1,6 +1,5 @@
 #include "task-spec/dynamic_graph/training_operation_attrs.h"
 #include "op-attrs/pcg_operator_attrs.h"
-#include "task-spec/dynamic_graph/gradient_reduction.dtg.h"
 #include "utils/overload.h"
 
 namespace FlexFlow {
diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc
index c8e553cce5..a1e6180249 100644
--- a/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc
+++ b/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc
@@ -1,14 +1,11 @@
 #include "task-spec/dynamic_graph/pass_expansion.h"
 #include "op-attrs/initializer_attrs.h"
-#include "op-attrs/operator_type.dtg.h"
-#include "op-attrs/ops/element_binary_attrs.dtg.h"
 #include "op-attrs/ops/element_unary.h"
-#include "op-attrs/tensor_slot_name.dtg.h"
 #include "task-spec/dynamic_graph/dynamic_gradient_reduction_layer_guid_t.dtg.h"
 #include "task-spec/dynamic_graph/dynamic_layer_guid_t.dtg.h"
 #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.h"
 #include "task-spec/dynamic_graph/dynamic_tensor_role.h"
-#include "task-spec/dynamic_graph/gradient_reduction.dtg.h"
+#include "task-spec/dynamic_graph/gradient_reduction_attrs.dtg.h"
 #include "task-spec/dynamic_graph/serializable_dynamic_node_invocation.h"
 #include "task-spec/dynamic_graph/serializable_dynamic_open_dataflow_graph.h"
 #include "task-spec/dynamic_graph/subgradient_id_t.dtg.h"
@@ -381,130 +378,6 @@ TEST_SUITE(FF_TEST_SUITE) {
     DynamicValueAttrs v2_grad = mk_value_attrs(1, grad_role);
     DynamicValueAttrs v3_grad = mk_value_attrs(2, grad_role);
 
-    SUBCASE("relu operator") {
-      TrainingOperationAttrs op_attrs = TrainingOperationAttrs{
-          PCGOperatorAttrs{
-              make_relu_attrs(),
-          },
-      };
-
-      DynamicNodeInvocation invocation = [&]() -> DynamicNodeInvocation {
-        return DynamicNodeInvocation{
-            /*inputs=*/{
-                {mk_slot(TensorSlotName::INPUT, std::nullopt), v1},
-            },
-            /*node_attrs=*/
-            DynamicNodeAttrs{
-                /*task_type=*/std::nullopt,
-                /*device_coord=*/std::nullopt,
-                /*mapping=*/std::nullopt,
-                /*op_attrs=*/op_attrs,
-                /*layer_guid=*/layer_guid,
-                /*per_device_op_state=*/std::nullopt,
-            },
-            /*outputs=*/
-            {
-                {mk_slot(TensorSlotName::OUTPUT, std::nullopt), v2},
-            },
-        };
-      }();
-
-      DynamicNodeInvocation result =
-          perform_bwd_pass_expansion_for_invocation(invocation);
-
-      DynamicNodeInvocation correct = [&]() -> DynamicNodeInvocation {
-        return DynamicNodeInvocation{
-            /*inputs=*/{
-                {mk_slot(TensorSlotName::INPUT, fwd_role), v1_fwd},
-                {mk_slot(TensorSlotName::OUTPUT, fwd_role), v2_fwd},
-                {mk_slot(TensorSlotName::OUTPUT, grad_role), v2_grad},
-            },
-            /*node_attrs=*/
-            DynamicNodeAttrs{
-                /*pass_type=*/DynamicTaskType::BWD,
-                /*device_coord=*/std::nullopt,
-                /*mapping=*/std::nullopt,
-                /*op_attrs=*/op_attrs,
-                /*layer_guid=*/layer_guid,
-                /*per_device_op_state=*/std::nullopt,
-            },
-            /*outputs=*/
-            {
-                {mk_slot(TensorSlotName::INPUT, grad_role), v1_grad},
-            },
-        };
-      }();
-
-      ASSERT(dynamic_node_invocation_to_serializable(result) ==
-             dynamic_node_invocation_to_serializable(correct));
-    }
-
-    SUBCASE("add operator") {
-      TrainingOperationAttrs op_attrs = TrainingOperationAttrs{
-          PCGOperatorAttrs{
-              ElementBinaryAttrs{
-                  /*type=*/OperatorType::EW_ADD,
-                  /*compute_type=*/DataType::FLOAT,
-                  /*should_broadcast_lhs=*/false,
-                  /*should_broadcast_rhs=*/false,
-              },
-          },
-      };
-
-      DynamicNodeInvocation invocation = [&]() -> DynamicNodeInvocation {
-        return DynamicNodeInvocation{
-            /*inputs=*/{
-                {mk_slot(TensorSlotName::LHS_INPUT, std::nullopt), v1},
-                {mk_slot(TensorSlotName::RHS_INPUT, std::nullopt), v2},
-            },
-            /*node_attrs=*/
-            DynamicNodeAttrs{
-                /*task_type=*/std::nullopt,
-                /*device_coord=*/std::nullopt,
-                /*mapping=*/std::nullopt,
-                /*op_attrs=*/op_attrs,
-                /*layer_guid=*/layer_guid,
-                /*per_device_op_state=*/std::nullopt,
-            },
-            /*outputs=*/
-            {
-                {mk_slot(TensorSlotName::OUTPUT, std::nullopt), v3},
-            },
-        };
-      }();
-
-      DynamicNodeInvocation result =
-          perform_bwd_pass_expansion_for_invocation(invocation);
-
-      DynamicNodeInvocation correct = [&]() -> DynamicNodeInvocation {
-        return DynamicNodeInvocation{
-            /*inputs=*/{
-                {mk_slot(TensorSlotName::LHS_INPUT, fwd_role), v1_fwd},
-                {mk_slot(TensorSlotName::RHS_INPUT, fwd_role), v2_fwd},
-                {mk_slot(TensorSlotName::OUTPUT, fwd_role), v3_fwd},
-                {mk_slot(TensorSlotName::OUTPUT, grad_role), v3_grad},
-            },
-            /*node_attrs=*/
-            DynamicNodeAttrs{
-                /*pass_type=*/DynamicTaskType::BWD,
-                /*device_coord=*/std::nullopt,
-                /*mapping=*/std::nullopt,
-                /*op_attrs=*/op_attrs,
-                /*layer_guid=*/layer_guid,
-                /*per_device_op_state=*/std::nullopt,
-            },
-            /*outputs=*/
-            {
-                {mk_slot(TensorSlotName::LHS_INPUT, grad_role), v1_grad},
-                {mk_slot(TensorSlotName::RHS_INPUT, grad_role), v2_grad},
-            },
-        };
-      }();
-
-      ASSERT(dynamic_node_invocation_to_serializable(result) ==
-             dynamic_node_invocation_to_serializable(correct));
-    }
-
     SUBCASE("normal operator") {
       TrainingOperationAttrs op_attrs = TrainingOperationAttrs{
           PCGOperatorAttrs{

From 19d2671088181f933bab985e567b6bb2f43519d5 Mon Sep 17 00:00:00 2001
From: Elliott Slaughter <slaughter@cs.stanford.edu>
Date: Thu, 20 Aug 2026 17:43:44 -0700
Subject: [PATCH 7/9] Cleanup.

---
 .../src/realm-execution/pcg_instance.cc       |  1 -
 .../dynamic_invocation_id_t.dtg.toml          | 20 -------------------
 .../task-spec/dynamic_graph/pass_expansion.cc | 15 +++-----------
 .../task-spec/dynamic_graph/pass_expansion.cc |  5 -----
 4 files changed, 3 insertions(+), 38 deletions(-)
 delete mode 100644 lib/task-spec/src/task-spec/dynamic_graph/dynamic_invocation_id_t.dtg.toml

diff --git a/lib/realm-execution/src/realm-execution/pcg_instance.cc b/lib/realm-execution/src/realm-execution/pcg_instance.cc
index 458e0e42d1..a037ca2654 100644
--- a/lib/realm-execution/src/realm-execution/pcg_instance.cc
+++ b/lib/realm-execution/src/realm-execution/pcg_instance.cc
@@ -14,7 +14,6 @@
 #include "task-spec/dynamic_graph/dynamic_task_type.dtg.h"
 #include "task-spec/dynamic_graph/dynamic_tensor_guid_t.dtg.h"
 #include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h"
-#include "task-spec/dynamic_graph/gradient_reduction_attrs.dtg.h"
 #include "task-spec/dynamic_graph/loss_insertion.h"
 #include "task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_mapped_pcg.h"
 #include "task-spec/dynamic_graph/pass_expansion.h"
diff --git a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_invocation_id_t.dtg.toml b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_invocation_id_t.dtg.toml
deleted file mode 100644
index f600824ab3..0000000000
--- a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_invocation_id_t.dtg.toml
+++ /dev/null
@@ -1,20 +0,0 @@
-namespace = "FlexFlow"
-name = "dynamic_invocation_id_t"
-type = "struct"
-features = [
-  "eq",
-  "ord",
-  "hash",
-  "json",
-  "fmt",
-]
-
-includes = [
-  "utils/nonnegative_int/nonnegative_int.h",
-]
-
-src_includes = []
-
-[[fields]]
-name = "idx"
-type = "::FlexFlow::nonnegative_int"
diff --git a/lib/task-spec/src/task-spec/dynamic_graph/pass_expansion.cc b/lib/task-spec/src/task-spec/dynamic_graph/pass_expansion.cc
index cedf51bed0..e536e3ebaa 100644
--- a/lib/task-spec/src/task-spec/dynamic_graph/pass_expansion.cc
+++ b/lib/task-spec/src/task-spec/dynamic_graph/pass_expansion.cc
@@ -1,19 +1,13 @@
 #include "task-spec/dynamic_graph/pass_expansion.h"
-#include "op-attrs/tensor_slot_name.dtg.h"
-#include "task-spec/dynamic_graph/dynamic_node_attrs.dtg.h"
-#include "task-spec/dynamic_graph/dynamic_node_invocation.dtg.h"
 #include "task-spec/dynamic_graph/dynamic_node_invocation.h"
 #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.h"
 #include "task-spec/dynamic_graph/dynamic_tensor_role.h"
-#include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h"
-#include "task-spec/dynamic_graph/subgradient_id_t.dtg.h"
 #include "task-spec/dynamic_graph/training_operation_attrs.h"
 #include "utils/containers/are_all_same.h"
 #include "utils/containers/concat_vectors.h"
 #include "utils/containers/contains_duplicates.h"
 #include "utils/containers/filter.h"
 #include "utils/containers/flatmap.h"
-#include "utils/containers/generate_map.h"
 #include "utils/containers/get_only.h"
 #include "utils/containers/map_from_pairs.h"
 #include "utils/containers/map_values.h"
@@ -25,9 +19,6 @@
 #include "utils/containers/set_of.h"
 #include "utils/containers/transform.h"
 #include "utils/containers/zip_with.h"
-#include "utils/optional.h"
-#include <optional>
-#include <set>
 
 namespace FlexFlow {
 
@@ -372,8 +363,8 @@ static std::map<int, subgradient_id_t>
   return result;
 }
 
-static std::vector<DynamicNodeInvocation>
-    reduce_gradients(std::vector<DynamicNodeInvocation> const &invocations) {
+static std::vector<DynamicNodeInvocation> reduce_gradients_for_invocations(
+    std::vector<DynamicNodeInvocation> const &invocations) {
   std::multiset<DynamicValueAttrs> outputs =
       multiset_of(flatmap(invocations, [](DynamicNodeInvocation const &i) {
         return vector_of(values(i.outputs));
@@ -476,7 +467,7 @@ static DynamicOpenDataflowGraph
 
   ASSERT(!contains_duplicates(new_invocation_set));
 
-  new_invocation_set = reduce_gradients(new_invocation_set);
+  new_invocation_set = reduce_gradients_for_invocations(new_invocation_set);
 
   return dynamic_open_dataflow_graph_from_invocation_set(
       set_of(new_invocation_set));
diff --git a/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc b/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc
index a1e6180249..c66777b7e6 100644
--- a/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc
+++ b/lib/task-spec/test/src/task-spec/dynamic_graph/pass_expansion.cc
@@ -1,15 +1,10 @@
 #include "task-spec/dynamic_graph/pass_expansion.h"
 #include "op-attrs/initializer_attrs.h"
 #include "op-attrs/ops/element_unary.h"
-#include "task-spec/dynamic_graph/dynamic_gradient_reduction_layer_guid_t.dtg.h"
-#include "task-spec/dynamic_graph/dynamic_layer_guid_t.dtg.h"
 #include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.h"
 #include "task-spec/dynamic_graph/dynamic_tensor_role.h"
-#include "task-spec/dynamic_graph/gradient_reduction_attrs.dtg.h"
 #include "task-spec/dynamic_graph/serializable_dynamic_node_invocation.h"
 #include "task-spec/dynamic_graph/serializable_dynamic_open_dataflow_graph.h"
-#include "task-spec/dynamic_graph/subgradient_id_t.dtg.h"
-#include "task-spec/dynamic_graph/training_operation_attrs.dtg.h"
 #include "test/utils/doctest/check_kv.h"
 #include <doctest/doctest.h>
 

From 9071af7fda9678c017759bc7f9daaa5f92c1c667 Mon Sep 17 00:00:00 2001
From: Elliott Slaughter <slaughter@cs.stanford.edu>
Date: Thu, 20 Aug 2026 20:14:12 -0700
Subject: [PATCH 8/9] Fix doxygen reference.

---
 .../task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml        | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml
index bd22171aa3..2da74edd71 100644
--- a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml
+++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml
@@ -55,7 +55,7 @@ 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 GradientReduction in the backward pass.
+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]]

From bacc8898afa8c3b769676d31b3c4e774c6a0634e Mon Sep 17 00:00:00 2001
From: Elliott Slaughter <slaughter@cs.stanford.edu>
Date: Fri, 21 Aug 2026 11:15:34 -0700
Subject: [PATCH 9/9] Add a failing e2e test for gradient reduction.

---
 .../test/src/realm-execution/test_e2e.cc      | 247 ++++++++++++++++++
 1 file changed, 247 insertions(+)

diff --git a/lib/realm-execution/test/src/realm-execution/test_e2e.cc b/lib/realm-execution/test/src/realm-execution/test_e2e.cc
index 9ba4886b4b..0d3f7cee2b 100644
--- a/lib/realm-execution/test/src/realm-execution/test_e2e.cc
+++ b/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