From eeaa135d5a592d1f695244a81feb6b72dab885a6 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Thu, 20 Aug 2026 10:12:40 -0700 Subject: [PATCH 01/21] 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 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 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{ + 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 invocation_set = { + DynamicNodeInvocation{ + /*inputs=*/std::map{}, + /*node_attrs=*/input_node, + /*outputs=*/ + std::map{ + { + mk_dynamic_slot(TensorSlotName::OUTPUT), + input_tensor, + }, + }, + }, + DynamicNodeInvocation{ + /*inputs=*/std::map{ + { + mk_dynamic_slot(TensorSlotName::INPUT), + input_tensor, + }, + }, + /*node_attrs=*/relu1_node, + /*outputs=*/ + std::map{ + { + mk_dynamic_slot(TensorSlotName::OUTPUT), + relu1_output, + }, + }, + }, + DynamicNodeInvocation{ + /*inputs=*/std::map{ + { + mk_dynamic_slot(TensorSlotName::INPUT), + input_tensor, + }, + }, + /*node_attrs=*/relu2_node, + /*outputs=*/ + std::map{ + { + 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 invocation_set = { + DynamicNodeInvocation{ + /*inputs=*/std::map{}, + /*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 Date: Thu, 20 Aug 2026 12:03:50 -0700 Subject: [PATCH 02/21] 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 = [ "", "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" +[[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 = [ "", "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" +[[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 Date: Thu, 20 Aug 2026 15:03:39 -0700 Subject: [PATCH 03/21] 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 = [ "", "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" [[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 = [ "", "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" [[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 Date: Thu, 20 Aug 2026 17:10:41 -0700 Subject: [PATCH 04/21] 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 +#include 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 + choose_subgradient_ids(std::set const &invocations) { + const std::vector 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 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 + reduce_gradients(std::vector const &invocations) { + std::multiset outputs = + multiset_of(flatmap(invocations, [](DynamicNodeInvocation const &i) { + return vector_of(values(i.outputs)); + })); + std::set unique_outputs = set_of(outputs); + + std::map invocation_index = + map_from_keys_and_values(invocations, range(invocations.size())); + + std::map> 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> + output_subgradient_ids = + map_values(output_invocation_set, choose_subgradient_ids); + + std::map gradient_reductions = + map_values2(output_invocation_set, + [&](DynamicValueAttrs const &output, + std::set 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{ + /*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( + DynamicNodeInvocation const &)> const &f) { + std::set current_invocation_set = + get_dynamic_invocation_set(g); + std::vector 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 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 @@ -1043,8 +1048,24 @@ TEST_SUITE(FF_TEST_SUITE) { }; }; + auto mk_gradient_node_attrs = + [](std::optional 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 const &tensor_type) + [](size_t node_id, + std::optional const &tensor_type, + std::optional 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 Date: Thu, 20 Aug 2026 17:25:22 -0700 Subject: [PATCH 05/21] 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 Date: Thu, 20 Aug 2026 17:36:06 -0700 Subject: [PATCH 06/21] 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 Date: Thu, 20 Aug 2026 17:43:44 -0700 Subject: [PATCH 07/21] 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 -#include namespace FlexFlow { @@ -372,8 +363,8 @@ static std::map return result; } -static std::vector - reduce_gradients(std::vector const &invocations) { +static std::vector reduce_gradients_for_invocations( + std::vector const &invocations) { std::multiset 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 From 9071af7fda9678c017759bc7f9daaa5f92c1c667 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Thu, 20 Aug 2026 20:14:12 -0700 Subject: [PATCH 08/21] 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 Date: Fri, 21 Aug 2026 11:15:34 -0700 Subject: [PATCH 09/21] 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 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 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 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 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 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 From 4d5598e20a2e73d4ea4b77490938dad15a923ff3 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Fri, 21 Aug 2026 20:33:02 -0700 Subject: [PATCH 10/21] Better formatting. --- .../task-spec/dynamic_graph/pass_expansion.cc | 34 ++++++++++--------- 1 file changed, 18 insertions(+), 16 deletions(-) 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 e536e3ebaa..8a3411d0d4 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 @@ -337,22 +337,24 @@ DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation( static std::map choose_subgradient_ids(std::set const &invocations) { - const std::vector 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}; + const std::vector 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 result = map_from_pairs( zip_with(vector_of(invocations), slot_names, From e308be07d6aa7da8fdd3ad26f28211985a3bb4a4 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Tue, 25 Aug 2026 13:32:24 -0700 Subject: [PATCH 11/21] Checkpoint. --- .../dynamic_graph/dynamic_node_attrs.h | 14 ++ .../dynamic_graph/dynamic_node_mapping.h | 5 + .../task-spec/dynamic_graph/copy_insertion.cc | 64 ++++++++- .../dynamic_graph/dynamic_node_attrs.cc | 16 +++ .../dynamic_graph/dynamic_node_mapping.cc | 19 +++ .../task-spec/dynamic_graph/pass_expansion.cc | 130 +++++++++++------- 6 files changed, 192 insertions(+), 56 deletions(-) create mode 100644 lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.h create mode 100644 lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_attrs.cc diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.h b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.h new file mode 100644 index 0000000000..0f31afc73d --- /dev/null +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.h @@ -0,0 +1,14 @@ +#ifndef _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_NODE_ATTRS_H +#define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_NODE_ATTRS_H + +#include "task-spec/dynamic_graph/dynamic_node_attrs.dtg.h" +#include "task-spec/dynamic_graph/dynamic_node_mapping.dtg.h" + +namespace FlexFlow { + +DynamicValueAttrs decide_dynamic_node_attrs_mapping(DynamicNodeAttrs const &, + DynamicNodeMapping const &); + +} // namespace FlexFlow + +#endif diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h index 3b28ea5149..e9ab82eeb2 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h @@ -2,10 +2,15 @@ #define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_NODE_MAPPING_H #include "task-spec/dynamic_graph/dynamic_node_mapping.dtg.h" +#include "task-spec/dynamic_graph/dynamic_tensor_slot.dtg.h" +#include "task-spec/dynamic_graph/parallel_tensor_mapping.dtg.h" #include "task-spec/global_device_id_t.dtg.h" namespace FlexFlow { +DynamicNodeMapping dynamic_node_mapping_from_value_mapping( + std::map const &); + bidict dynamic_node_mapping_get_shard_bindings(DynamicNodeMapping const &); diff --git a/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc b/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc index ef5695792f..76cfde36d9 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc @@ -15,6 +15,7 @@ #include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h" #include "task-spec/dynamic_graph/dynamic_value_attrs.h" #include "task-spec/dynamic_graph/parallel_tensor_mapping.dtg.h" +#include "task-spec/dynamic_graph/training_only_op_type.dtg.h" #include "task-spec/dynamic_graph/training_operation_attrs.h" #include "utils/bidict/algorithms/bidict_from_unstructured_relation.h" #include "utils/bidict/algorithms/unstructured_relation_from_bidict.h" @@ -42,6 +43,10 @@ bool node_is_copy(DynamicNodeAttrs const &n) { return n.op_attrs.has_value() && n.op_attrs.value().is_copy(); } +bool node_is_gradient_reduction(DynamicNodeAttrs const &n) { + return n.op_attrs.has_value() && n.op_attrs.value().is_gradient_reduction(); +} + bool value_is_mapped(DynamicValueAttrs const &n) { return n.mapping.has_value(); } @@ -49,7 +54,8 @@ bool value_is_mapped(DynamicValueAttrs const &n) { void require_node_is_ready_for_copy_insertion( DynamicNodeAttrs const &node_attrs) { ASSERT(node_attrs.op_attrs.has_value()); - ASSERT(node_attrs.mapping.has_value()); + ASSERT(node_attrs.mapping.has_value() || + node_is_gradient_reduction(node_attrs)); } void require_value_is_ready_for_copy_insertion(DynamicValueAttrs const &v) { @@ -148,6 +154,31 @@ DynamicNodeInvocation apply_mappings_for_invocation( return result; } +DynamicNodeInvocation apply_mappings_for_invocation_node_only( + dynamic_invocation_id_t const &id, + DynamicNodeInvocation const &i, + std::map const + &all_mappings) { + + require_invocation_is_ready_for_copy_insertion(i); + + std::map i_mappings = + get_mappings_for_invocation_id(id, all_mappings); + + DynamicNodeMapping mapping = + dynamic_node_mapping_from_value_mapping(i_mappings); + + DynamicNodeInvocation result = DynamicNodeInvocation{ + /*inputs=*/i.inputs, + /*node_attrs=*/ + decide_dynamic_node_attrs_mapping(i.node_attrs, mapping), + /*outputs=*/ + i.outputs, + }; + + return result; +} + DynamicNodeInvocation make_copy_invocation(DynamicValueCopyInfo const ©_info) { DynamicNodeInvocation result = DynamicNodeInvocation{ @@ -269,7 +300,6 @@ std::map std::map resolve_partial_tensor_mappings_from_node_mappings( DynamicOpenDataflowGraph const &g) { - require_graph_is_ready_for_copy_insertion(g); auto slots_to_map_for_replicate = @@ -290,7 +320,7 @@ std::map { InternalDynamicSlotSite slot_site = get_only(slot_sites); ASSERT(slot_site.slot_name.slot_name == TensorSlotName::OUTPUT); - }; + } return slot_sites; }; @@ -303,7 +333,9 @@ std::map dynamic_graph_get_id_for_invocation(g, invocation); std::set slot_sites_to_resolve = [&] { - if (op_type == TrainingOpType{OperatorType::REPLICATE}) { + if (op_type == TrainingOpType{TrainingOnlyOpType::GRADIENT_REDUCTION}) { + return std::set{}; + } else if (op_type == TrainingOpType{OperatorType::REPLICATE}) { return slots_to_map_for_replicate(invocation_id, invocation); } else { return get_dynamic_slot_sites_for_invocation(invocation_id, invocation); @@ -349,9 +381,31 @@ std::map TrainingOpType op_type = dynamic_node_invocation_get_op_type(invocation); std::optional task_type = invocation.node_attrs.task_type; + TrainingOpType gradient_reduction_op_type = + TrainingOpType{TrainingOnlyOpType::GRADIENT_REDUCTION}; TrainingOpType replicate_op_type = TrainingOpType{OperatorType::REPLICATE}; - if (op_type == replicate_op_type && task_type == DynamicTaskType::BWD) { + if (op_type == gradient_reduction_op_type) { + if (slot_site.direction == TensorDirection::OUTPUT) { + InternalDynamicSlotSite slot_site_sink = + get_only(dynamic_graph_find_sinks_of_slot_site(g, slot_site)); + + ASSERT(contains_key(resolved_mappings, slot_site_sink)); + + return resolved_mappings.at(slot_site_sink); + } else { + ASSERT(slot_site.direction == TensorDirection::INCOMING); + + InternalDynamicSlotSite slot_site_src = + dynamic_graph_find_source_of_slot_site(g, slot_site) + .require_internal(); + + ASSERT(contains_key(resolved_mappings, slot_site_src)); + + return resolved_mappings.at(slot_site_src); + } + } else if (op_type == replicate_op_type && + task_type == DynamicTaskType::BWD) { ASSERT(slot_site.direction == TensorDirection::OUTPUT); InternalDynamicSlotSite slot_site_sink = diff --git a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_attrs.cc b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_attrs.cc new file mode 100644 index 0000000000..0271188af1 --- /dev/null +++ b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_attrs.cc @@ -0,0 +1,16 @@ +#include "task-spec/dynamic_graph/dynamic_node_attrs.h" + +namespace FlexFlow { + +DynamicNodeAttrs + decide_dynamic_node_attrs_mapping(DynamicNodeAttrs const &attrs, + DynamicNodeMapping const &mapping) { + ASSERT(!attrs.mapping.has_value()); + + DynamicNodeAttrs result = attrs; + result.mapping = mapping; + + return result; +} + +} // namespace FlexFlow diff --git a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc index e66b61c6d4..a8f9d2db3c 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc @@ -1,10 +1,29 @@ #include "task-spec/dynamic_graph/dynamic_node_mapping.h" +#include "task-spec/dynamic_graph/parallel_tensor_mapping.h" #include "utils/bidict/algorithms/bidict_transform_keys.h" #include "utils/bidict/algorithms/bidict_transform_values.h" +#include "utils/containers/flatmap.h" +#include "utils/containers/require_all_same.h" #include "utils/containers/transform.h" +#include "utils/optional.h" namespace FlexFlow { +DynamicNodeMapping dynamic_node_mapping_from_value_mapping( + std::map const &mapping) { + MappedOperatorTaskGroup op_task_group; + DeviceType device_type = assert_unwrap(require_all_same( + transform(flatmap(values(mapping), + [](ParallelTensorMapping const &m) { + return pt_mapping_get_device_set(m); + }), + [](global_device_id_t const &id) { return id.device_type; }))); + return DynamicNodeMapping{ + /*op_task_group=*/op_task_group, + /*device_type=*/device_type, + }; +} + bidict dynamic_node_mapping_get_shard_bindings(DynamicNodeMapping const &m) { return bidict_transform_keys( 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 8a3411d0d4..dee31474d9 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,4 +1,5 @@ #include "task-spec/dynamic_graph/pass_expansion.h" +#include "task-spec/dynamic_graph/copy_insertion.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" @@ -365,22 +366,30 @@ static std::map return result; } -static std::vector reduce_gradients_for_invocations( +std::set multiply_defined_outputs( std::vector const &invocations) { std::multiset outputs = multiset_of(flatmap(invocations, [](DynamicNodeInvocation const &i) { return vector_of(values(i.outputs)); })); - std::set unique_outputs = set_of(outputs); + return filter(set_of(outputs), [&](DynamicValueAttrs const &value) { + return outputs.count(value) > 1; + }); +} + +static std::vector reduce_gradients_for_invocations( + std::vector const &invocations) { + std::set multiple_outputs = + multiply_defined_outputs(invocations); std::map invocation_index = map_from_keys_and_values(invocations, range(invocations.size())); std::map> output_invocation_set; for (int idx : range(invocations.size())) { - DynamicNodeInvocation const &i = invocations[idx]; + DynamicNodeInvocation const &i = invocations.at(idx); for (DynamicValueAttrs const &v : values(i.outputs)) { - if (outputs.count(v) > 1) { + if (multiple_outputs.count(v) != 0) { output_invocation_set[v].insert(idx); } } @@ -389,52 +398,52 @@ static std::vector reduce_gradients_for_invocations( output_subgradient_ids = map_values(output_invocation_set, choose_subgradient_ids); + auto compute_gradient_reduction = [&](DynamicValueAttrs const &output, + std::set 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{ + /*slot_name=*/TensorSlotName::OUTPUT, + /*slot_tensor_role=*/ + DynamicTensorRole{FwbTensorType::GRADIENT}, + /*task_shard=*/std::nullopt, + }, + output, + }, + }, + }; + }; + std::map gradient_reductions = - map_values2(output_invocation_set, - [&](DynamicValueAttrs const &output, - std::set 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{ - /*slot_name=*/TensorSlotName::OUTPUT, - /*slot_tensor_role=*/ - DynamicTensorRole{FwbTensorType::GRADIENT}, - /*task_shard=*/std::nullopt, - }, - output, - }, - }, - }; - }); + map_values2(output_invocation_set, compute_gradient_reduction); return concat_vectors( transform(range(invocations.size()), @@ -471,8 +480,27 @@ static DynamicOpenDataflowGraph new_invocation_set = reduce_gradients_for_invocations(new_invocation_set); - return dynamic_open_dataflow_graph_from_invocation_set( - set_of(new_invocation_set)); + DynamicOpenDataflowGraph unmapped_graph = + dynamic_open_dataflow_graph_from_invocation_set( + set_of(new_invocation_set)); + + std::map + fully_resolved_tensor_mappings = resolve_tensor_mappings(unmapped_graph); + + std::set mapped_invocations = + transform(get_dynamic_invocation_set(unmapped_graph), + [&](DynamicNodeInvocation const &i) -> DynamicNodeInvocation { + dynamic_invocation_id_t id = + dynamic_graph_get_id_for_invocation(unmapped_graph, i); + + return apply_mappings_for_invocation( + id, i, fully_resolved_tensor_mappings); + }); + + DynamicOpenDataflowGraph result = + dynamic_open_dataflow_graph_from_invocation_set(mapped_invocations); + + return result; } DynamicOpenDataflowGraph From 66b17a1a0d72ec7de0a971213af570bf98a0beb0 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Tue, 25 Aug 2026 15:33:40 -0700 Subject: [PATCH 12/21] Restructure pass expansion implementation. --- .../dynamic_graph/dynamic_node_attrs.h | 4 +- .../task-spec/dynamic_graph/pass_expansion.h | 11 +- .../task-spec/dynamic_graph/copy_insertion.cc | 25 -- .../dynamic_graph/dynamic_node_mapping.cc | 2 +- .../task-spec/dynamic_graph/pass_expansion.cc | 371 +++++++++--------- 5 files changed, 197 insertions(+), 216 deletions(-) diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.h b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.h index 0f31afc73d..a8e9c39dcc 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.h +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.h @@ -6,8 +6,8 @@ namespace FlexFlow { -DynamicValueAttrs decide_dynamic_node_attrs_mapping(DynamicNodeAttrs const &, - DynamicNodeMapping const &); +DynamicNodeAttrs decide_dynamic_node_attrs_mapping(DynamicNodeAttrs const &, + DynamicNodeMapping const &); } // namespace FlexFlow diff --git a/lib/task-spec/include/task-spec/dynamic_graph/pass_expansion.h b/lib/task-spec/include/task-spec/dynamic_graph/pass_expansion.h index 3adf66d117..cd01cf26c5 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/pass_expansion.h +++ b/lib/task-spec/include/task-spec/dynamic_graph/pass_expansion.h @@ -41,8 +41,15 @@ DynamicNodeAttrs pass_expand_node(DynamicNodeAttrs const &, DynamicTaskType); DynamicNodeInvocation perform_fwd_pass_expansion_for_invocation(DynamicNodeInvocation const &); -DynamicNodeInvocation - perform_bwd_pass_expansion_for_invocation(DynamicNodeInvocation const &); +DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation( + dynamic_invocation_id_t const &, + DynamicNodeInvocation const &, + std::map> const &); +DynamicNodeInvocation create_gradient_reduction_for_value( + DynamicOpenDataflowGraph const &, + DynamicValueAttrs const &, + std::map const &); DynamicOpenDataflowGraph perform_pass_expansion(DynamicOpenDataflowGraph const &); diff --git a/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc b/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc index 76cfde36d9..5c49c53232 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc @@ -154,31 +154,6 @@ DynamicNodeInvocation apply_mappings_for_invocation( return result; } -DynamicNodeInvocation apply_mappings_for_invocation_node_only( - dynamic_invocation_id_t const &id, - DynamicNodeInvocation const &i, - std::map const - &all_mappings) { - - require_invocation_is_ready_for_copy_insertion(i); - - std::map i_mappings = - get_mappings_for_invocation_id(id, all_mappings); - - DynamicNodeMapping mapping = - dynamic_node_mapping_from_value_mapping(i_mappings); - - DynamicNodeInvocation result = DynamicNodeInvocation{ - /*inputs=*/i.inputs, - /*node_attrs=*/ - decide_dynamic_node_attrs_mapping(i.node_attrs, mapping), - /*outputs=*/ - i.outputs, - }; - - return result; -} - DynamicNodeInvocation make_copy_invocation(DynamicValueCopyInfo const ©_info) { DynamicNodeInvocation result = DynamicNodeInvocation{ diff --git a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc index a8f9d2db3c..c9e6dbbefb 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc @@ -11,7 +11,7 @@ namespace FlexFlow { DynamicNodeMapping dynamic_node_mapping_from_value_mapping( std::map const &mapping) { - MappedOperatorTaskGroup op_task_group; + MappedOperatorTaskGroup op_task_group{{}}; DeviceType device_type = assert_unwrap(require_all_same( transform(flatmap(values(mapping), [](ParallelTensorMapping const &m) { 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 dee31474d9..e7d8a0972b 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,14 +1,19 @@ #include "task-spec/dynamic_graph/pass_expansion.h" #include "task-spec/dynamic_graph/copy_insertion.h" +#include "task-spec/dynamic_graph/dynamic_invocation_id_t.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" @@ -17,9 +22,12 @@ #include "utils/containers/multiset_of.h" #include "utils/containers/range.h" #include "utils/containers/repeat_until_converged.h" +#include "utils/containers/set_intersection.h" #include "utils/containers/set_of.h" #include "utils/containers/transform.h" +#include "utils/containers/try_at.h" #include "utils/containers/zip_with.h" +#include namespace FlexFlow { @@ -187,6 +195,65 @@ std::set return flatmap(required_values, get_sink_invocations_for_value); } +std::set determine_values_requiring_gradient_reduction( + DynamicOpenDataflowGraph const &g, + std::set const &in_bwd_pass) { + std::set internal_values = + dynamic_graph_get_internal_values(g); + return filter(internal_values, [&](DynamicValueAttrs const &v) { + dynamic_invocation_id_t source = dynamic_graph_find_source_of_value(g, v) + .require_internal() + .invocation_id; + return in_bwd_pass.count(source) != 0 && + dynamic_graph_find_sinks_of_value(g, v).size() > 1; + }); +} + +static std::map + assign_subgradient_ids( + std::set const &invocations) { + const std::vector 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 result = map_from_pairs( + zip_with(vector_of(invocations), + slot_names, + [](dynamic_invocation_id_t id, TensorSlotName slot_name) { + return std::pair{id, subgradient_id_t{slot_name}}; + })); + ASSERT(result.size() == invocations.size()); + return result; +} + +std::map> + compute_subgradient_ids(DynamicOpenDataflowGraph const &g, + std::set const + &values_requiring_gradient_reduction) { + return generate_map( + values_requiring_gradient_reduction, [&](DynamicValueAttrs const &v) { + return assign_subgradient_ids( + transform(dynamic_graph_find_sinks_of_value(g, v), + [](InternalDynamicSlotSite const &sink_site) { + return sink_site.invocation_id; + })); + }); +} + DynamicTensorSlot pass_expand_slot(DynamicTensorSlot const &s, FwbTensorType tensor_type) { require_slot_is_not_pass_expanded(s); @@ -197,12 +264,15 @@ DynamicTensorSlot pass_expand_slot(DynamicTensorSlot const &s, return result; } -DynamicValueAttrs pass_expand_value(DynamicValueAttrs const &v, - FwbTensorType tensor_type) { +DynamicValueAttrs + pass_expand_value(DynamicValueAttrs const &v, + FwbTensorType tensor_type, + std::optional const &subgradient_id) { require_value_is_not_pass_expanded(v); DynamicValueAttrs result = v; result.role = DynamicTensorRole{tensor_type}; + result.subgradient_id = subgradient_id; return result; }; @@ -235,13 +305,13 @@ DynamicNodeInvocation perform_fwd_pass_expansion_for_invocation( assert_unwrap(invocation.node_attrs.op_attrs); auto to_fwd_value = [](DynamicValueAttrs const &v) -> DynamicValueAttrs { - return pass_expand_value(v, FwbTensorType::FORWARD); + return pass_expand_value(v, FwbTensorType::FORWARD, std::nullopt); }; auto to_fwd = [](DynamicTensorSlot const &k, DynamicValueAttrs const &v) { return std::pair{ pass_expand_slot(k, FwbTensorType::FORWARD), - pass_expand_value(v, FwbTensorType::FORWARD), + pass_expand_value(v, FwbTensorType::FORWARD, std::nullopt), }; }; @@ -270,28 +340,54 @@ DynamicNodeInvocation perform_fwd_pass_expansion_for_invocation( } DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation( - DynamicNodeInvocation const &invocation) { + dynamic_invocation_id_t const &invocation_id, + DynamicNodeInvocation const &invocation, + std::map> const + &subgradient_ids) { require_invocation_is_ready_for_pass_expansion(invocation); TrainingOperationAttrs op_attrs = assert_unwrap(invocation.node_attrs.op_attrs); + auto subgradient_id_for_value = [&](DynamicValueAttrs const &v) { + return transform( + try_at(subgradient_ids, v), + [&](std::map const &s) { + return s.at(invocation_id); + }); + }; + auto to_fwd = [](DynamicTensorSlot const &k, DynamicValueAttrs const &v) { return std::pair{ pass_expand_slot(k, FwbTensorType::FORWARD), - pass_expand_value(v, FwbTensorType::FORWARD), + pass_expand_value(v, FwbTensorType::FORWARD, std::nullopt), }; }; auto to_grad_value = [](DynamicValueAttrs const &v) { - return pass_expand_value(v, FwbTensorType::GRADIENT); + return pass_expand_value(v, FwbTensorType::GRADIENT, std::nullopt); + }; + + auto to_grad_value_output = [&](DynamicValueAttrs const &v) { + return pass_expand_value( + v, FwbTensorType::GRADIENT, subgradient_id_for_value(v)); }; auto to_grad = [](DynamicTensorSlot const &k, DynamicValueAttrs const &v) { return std::pair{ pass_expand_slot(k, FwbTensorType::GRADIENT), - pass_expand_value(v, FwbTensorType::GRADIENT), + pass_expand_value(v, FwbTensorType::GRADIENT, std::nullopt), + }; + }; + + auto to_grad_output = [&](DynamicTensorSlot const &k, + DynamicValueAttrs const &v) { + return std::pair{ + pass_expand_slot(k, FwbTensorType::GRADIENT), + pass_expand_value( + v, FwbTensorType::GRADIENT, subgradient_id_for_value(v)), }; }; @@ -300,7 +396,7 @@ DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation( return DynamicNodeInvocation{ /*inputs=*/map_values(invocation.outputs, to_grad_value), /*node_attrs=*/invocation.node_attrs, - /*outputs=*/map_values(invocation.inputs, to_grad_value), + /*outputs=*/map_values(invocation.inputs, to_grad_value_output), }; } else if (training_op_attrs_has_op_type(op_attrs, OperatorType::REPLICATE)) { @@ -312,7 +408,7 @@ DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation( pass_expand_node(invocation.node_attrs, DynamicTaskType::BWD), /*outputs=*/ { - transform(invocation.inputs, to_grad), + transform(invocation.inputs, to_grad_output), }, }; } else { @@ -326,7 +422,7 @@ DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation( /*node_attrs=*/ pass_expand_node(invocation.node_attrs, DynamicTaskType::BWD), /*outputs=*/ - transform(invocation.inputs, to_grad), + transform(invocation.inputs, to_grad_output), }; }; }(); @@ -336,171 +432,56 @@ DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation( return result; } -static std::map - choose_subgradient_ids(std::set const &invocations) { - const std::vector 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 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; -} - -std::set multiply_defined_outputs( - std::vector const &invocations) { - std::multiset outputs = - multiset_of(flatmap(invocations, [](DynamicNodeInvocation const &i) { - return vector_of(values(i.outputs)); - })); - return filter(set_of(outputs), [&](DynamicValueAttrs const &value) { - return outputs.count(value) > 1; - }); -} +DynamicNodeInvocation create_gradient_reduction_for_value( + DynamicOpenDataflowGraph const &g, + DynamicValueAttrs const &value, + std::map const + &subgradient_inputs) { + DynamicSlotSite source_site = dynamic_graph_find_source_of_value(g, value); + std::set sink_sites = + dynamic_graph_find_sinks_of_value(g, value); -static std::vector reduce_gradients_for_invocations( - std::vector const &invocations) { - std::set multiple_outputs = - multiply_defined_outputs(invocations); - - std::map invocation_index = - map_from_keys_and_values(invocations, range(invocations.size())); - - std::map> output_invocation_set; - for (int idx : range(invocations.size())) { - DynamicNodeInvocation const &i = invocations.at(idx); - for (DynamicValueAttrs const &v : values(i.outputs)) { - if (multiple_outputs.count(v) != 0) { - output_invocation_set[v].insert(idx); - } - } - } - std::map> - output_subgradient_ids = - map_values(output_invocation_set, choose_subgradient_ids); - - auto compute_gradient_reduction = [&](DynamicValueAttrs const &output, - std::set 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{ - /*slot_name=*/TensorSlotName::OUTPUT, - /*slot_tensor_role=*/ - DynamicTensorRole{FwbTensorType::GRADIENT}, - /*task_shard=*/std::nullopt, - }, - output, - }, - }, + auto to_grad = [](DynamicTensorSlot const &k, DynamicValueAttrs const &v) { + return std::pair{ + pass_expand_slot(k, FwbTensorType::GRADIENT), + pass_expand_value(v, FwbTensorType::GRADIENT, std::nullopt), }; }; - std::map gradient_reductions = - map_values2(output_invocation_set, compute_gradient_reduction); - - 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( - DynamicNodeInvocation const &)> const &f) { - std::set current_invocation_set = - get_dynamic_invocation_set(g); - std::vector new_invocation_set = - flatmap(vector_of(current_invocation_set), f); - - ASSERT(!contains_duplicates(new_invocation_set)); - - new_invocation_set = reduce_gradients_for_invocations(new_invocation_set); - - DynamicOpenDataflowGraph unmapped_graph = - dynamic_open_dataflow_graph_from_invocation_set( - set_of(new_invocation_set)); - - std::map - fully_resolved_tensor_mappings = resolve_tensor_mappings(unmapped_graph); - - std::set mapped_invocations = - transform(get_dynamic_invocation_set(unmapped_graph), - [&](DynamicNodeInvocation const &i) -> DynamicNodeInvocation { - dynamic_invocation_id_t id = - dynamic_graph_get_id_for_invocation(unmapped_graph, i); - - return apply_mappings_for_invocation( - id, i, fully_resolved_tensor_mappings); - }); + auto to_grad_subgradient = [](DynamicTensorSlot const &k, + DynamicValueAttrs const &v, + subgradient_id_t const &s) { + return std::pair{ + pass_expand_slot(k, FwbTensorType::GRADIENT), + pass_expand_value(v, FwbTensorType::GRADIENT, s), + }; + }; - DynamicOpenDataflowGraph result = - dynamic_open_dataflow_graph_from_invocation_set(mapped_invocations); + DynamicNodeAttrs gradient_reduction_attrs{ + /*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, + }; - return result; + return DynamicNodeInvocation{ + /*inputs=*/map_from_pairs(transform( + values(subgradient_inputs), + [&](subgradient_id_t const &subgradient_id) { + return to_grad_subgradient(source_site.require_internal().slot_name, + value, + subgradient_id); + })), + /*node_attrs=*/gradient_reduction_attrs, + /*outputs=*/ + { + to_grad(source_site.require_internal().slot_name, value), + }, + }; } DynamicOpenDataflowGraph @@ -511,23 +492,41 @@ DynamicOpenDataflowGraph std::set needed_in_bwd_pass = determine_invocations_needed_in_backward_pass_for_gradient_computation(g); - 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{ + std::set values_requiring_gradient_reduction = + determine_values_requiring_gradient_reduction(g, needed_in_bwd_pass); + + std::map> + subgradient_ids = + compute_subgradient_ids(g, values_requiring_gradient_reduction); + + 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)) { + std::set gradient_reductions = + transform(set_intersection(keys(subgradient_ids), + set_of(values(invocation.outputs))), + [&](DynamicValueAttrs const &v) { + return create_gradient_reduction_for_value( + g, v, subgradient_ids.at(v)); + }); + + return set_union( + 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), - }; - } - }); + perform_bwd_pass_expansion_for_invocation( + invocation_id, invocation, subgradient_ids), + }, + gradient_reductions); + } else { + return std::set{ + perform_fwd_pass_expansion_for_invocation(invocation), + }; + } + }); require_graph_is_fully_pass_expanded(result); From 94f3eaa2ccb53246ead63416314c756761ec551c Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Tue, 25 Aug 2026 17:29:28 -0700 Subject: [PATCH 13/21] Cleanup. --- .../dynamic_graph/dynamic_value_attrs.h | 5 ++ .../dynamic_graph/dynamic_value_attrs.cc | 10 +++ .../task-spec/dynamic_graph/pass_expansion.cc | 80 +++++++++---------- 3 files changed, 53 insertions(+), 42 deletions(-) diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.h b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.h index 851facc2fd..6374ff7388 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.h +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.h @@ -3,12 +3,17 @@ #include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h" #include "task-spec/dynamic_graph/parallel_tensor_mapping.dtg.h" +#include "task-spec/dynamic_graph/subgradient_id_t.dtg.h" namespace FlexFlow { DynamicValueAttrs decide_dynamic_value_attrs_role(DynamicValueAttrs const &, DynamicTensorRole); +DynamicValueAttrs + decide_dynamic_value_attrs_subgradient_id(DynamicValueAttrs const &, + subgradient_id_t const &); + DynamicValueAttrs decide_dynamic_value_attrs_mapping(DynamicValueAttrs const &, ParallelTensorMapping const &); diff --git a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_value_attrs.cc b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_value_attrs.cc index a847c0ba75..7bfc123882 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_value_attrs.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_value_attrs.cc @@ -13,6 +13,16 @@ DynamicValueAttrs return result; } +DynamicValueAttrs decide_dynamic_value_attrs_subgradient_id( + DynamicValueAttrs const &attrs, subgradient_id_t const &subgradient_id) { + ASSERT(!attrs.subgradient_id.has_value()); + + DynamicValueAttrs result = attrs; + result.subgradient_id = subgradient_id; + + return result; +} + DynamicValueAttrs decide_dynamic_value_attrs_mapping(DynamicValueAttrs const &attrs, ParallelTensorMapping const &mapping) { 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 e7d8a0972b..beb04159ea 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,4 +1,5 @@ #include "task-spec/dynamic_graph/pass_expansion.h" +#include "op-attrs/tensor_slot_name.h" #include "task-spec/dynamic_graph/copy_insertion.h" #include "task-spec/dynamic_graph/dynamic_invocation_id_t.dtg.h" #include "task-spec/dynamic_graph/dynamic_node_invocation.dtg.h" @@ -6,6 +7,7 @@ #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/dynamic_value_attrs.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" @@ -15,6 +17,7 @@ #include "utils/containers/flatmap.h" #include "utils/containers/generate_map.h" #include "utils/containers/get_only.h" +#include "utils/containers/map_from_keys_and_values.h" #include "utils/containers/map_from_pairs.h" #include "utils/containers/map_values.h" #include "utils/containers/map_values2.h" @@ -24,9 +27,11 @@ #include "utils/containers/repeat_until_converged.h" #include "utils/containers/set_intersection.h" #include "utils/containers/set_of.h" +#include "utils/containers/slice.h" #include "utils/containers/transform.h" #include "utils/containers/try_at.h" #include "utils/containers/zip_with.h" +#include "utils/optional.h" #include namespace FlexFlow { @@ -212,32 +217,15 @@ std::set determine_values_requiring_gradient_reduction( static std::map assign_subgradient_ids( std::set const &invocations) { - const std::vector 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 result = map_from_pairs( - zip_with(vector_of(invocations), - slot_names, - [](dynamic_invocation_id_t id, TensorSlotName slot_name) { - return std::pair{id, subgradient_id_t{slot_name}}; - })); - ASSERT(result.size() == invocations.size()); - return result; + const std::vector slot_names = + get_variadic_inputs_slot_name_sequence(); + ASSERT(invocations.size() <= slot_names.size()); + return map_from_keys_and_values( + vector_of(invocations), + transform(slice(slot_names, 0, invocations.size()), + [](TensorSlotName const &slot_name) { + return subgradient_id_t{slot_name}; + })); } std::map> @@ -264,15 +252,12 @@ DynamicTensorSlot pass_expand_slot(DynamicTensorSlot const &s, return result; } -DynamicValueAttrs - pass_expand_value(DynamicValueAttrs const &v, - FwbTensorType tensor_type, - std::optional const &subgradient_id) { +DynamicValueAttrs pass_expand_value(DynamicValueAttrs const &v, + FwbTensorType tensor_type) { require_value_is_not_pass_expanded(v); DynamicValueAttrs result = v; result.role = DynamicTensorRole{tensor_type}; - result.subgradient_id = subgradient_id; return result; }; @@ -305,13 +290,13 @@ DynamicNodeInvocation perform_fwd_pass_expansion_for_invocation( assert_unwrap(invocation.node_attrs.op_attrs); auto to_fwd_value = [](DynamicValueAttrs const &v) -> DynamicValueAttrs { - return pass_expand_value(v, FwbTensorType::FORWARD, std::nullopt); + return pass_expand_value(v, FwbTensorType::FORWARD); }; auto to_fwd = [](DynamicTensorSlot const &k, DynamicValueAttrs const &v) { return std::pair{ pass_expand_slot(k, FwbTensorType::FORWARD), - pass_expand_value(v, FwbTensorType::FORWARD, std::nullopt), + pass_expand_value(v, FwbTensorType::FORWARD), }; }; @@ -359,26 +344,37 @@ DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation( }); }; + auto pass_expand_value_with_subgradient = [&](DynamicValueAttrs const &v, + FwbTensorType tensor_type) { + std::optional subgradient_id = + subgradient_id_for_value(v); + DynamicValueAttrs result = pass_expand_value(v, tensor_type); + if (subgradient_id.has_value()) { + result = decide_dynamic_value_attrs_subgradient_id( + result, assert_unwrap(subgradient_id)); + } + return result; + }; + auto to_fwd = [](DynamicTensorSlot const &k, DynamicValueAttrs const &v) { return std::pair{ pass_expand_slot(k, FwbTensorType::FORWARD), - pass_expand_value(v, FwbTensorType::FORWARD, std::nullopt), + pass_expand_value(v, FwbTensorType::FORWARD), }; }; auto to_grad_value = [](DynamicValueAttrs const &v) { - return pass_expand_value(v, FwbTensorType::GRADIENT, std::nullopt); + return pass_expand_value(v, FwbTensorType::GRADIENT); }; auto to_grad_value_output = [&](DynamicValueAttrs const &v) { - return pass_expand_value( - v, FwbTensorType::GRADIENT, subgradient_id_for_value(v)); + return pass_expand_value_with_subgradient(v, FwbTensorType::GRADIENT); }; auto to_grad = [](DynamicTensorSlot const &k, DynamicValueAttrs const &v) { return std::pair{ pass_expand_slot(k, FwbTensorType::GRADIENT), - pass_expand_value(v, FwbTensorType::GRADIENT, std::nullopt), + pass_expand_value(v, FwbTensorType::GRADIENT), }; }; @@ -386,8 +382,7 @@ DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation( DynamicValueAttrs const &v) { return std::pair{ pass_expand_slot(k, FwbTensorType::GRADIENT), - pass_expand_value( - v, FwbTensorType::GRADIENT, subgradient_id_for_value(v)), + pass_expand_value_with_subgradient(v, FwbTensorType::GRADIENT), }; }; @@ -444,7 +439,7 @@ DynamicNodeInvocation create_gradient_reduction_for_value( auto to_grad = [](DynamicTensorSlot const &k, DynamicValueAttrs const &v) { return std::pair{ pass_expand_slot(k, FwbTensorType::GRADIENT), - pass_expand_value(v, FwbTensorType::GRADIENT, std::nullopt), + pass_expand_value(v, FwbTensorType::GRADIENT), }; }; @@ -453,7 +448,8 @@ DynamicNodeInvocation create_gradient_reduction_for_value( subgradient_id_t const &s) { return std::pair{ pass_expand_slot(k, FwbTensorType::GRADIENT), - pass_expand_value(v, FwbTensorType::GRADIENT, s), + decide_dynamic_value_attrs_subgradient_id( + pass_expand_value(v, FwbTensorType::GRADIENT), s), }; }; From cef9b1b95203eca763639c09cd8cc48876655a49 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Tue, 25 Aug 2026 19:09:13 -0700 Subject: [PATCH 14/21] Work on mapping. --- .../dynamic_graph/dynamic_node_attrs.h | 14 --------- .../dynamic_graph/dynamic_node_attrs.cc | 16 ---------- .../task-spec/dynamic_graph/pass_expansion.cc | 30 +++++++++++++++++++ 3 files changed, 30 insertions(+), 30 deletions(-) delete mode 100644 lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.h delete mode 100644 lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_attrs.cc diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.h b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.h deleted file mode 100644 index a8e9c39dcc..0000000000 --- a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_attrs.h +++ /dev/null @@ -1,14 +0,0 @@ -#ifndef _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_NODE_ATTRS_H -#define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_NODE_ATTRS_H - -#include "task-spec/dynamic_graph/dynamic_node_attrs.dtg.h" -#include "task-spec/dynamic_graph/dynamic_node_mapping.dtg.h" - -namespace FlexFlow { - -DynamicNodeAttrs decide_dynamic_node_attrs_mapping(DynamicNodeAttrs const &, - DynamicNodeMapping const &); - -} // namespace FlexFlow - -#endif diff --git a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_attrs.cc b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_attrs.cc deleted file mode 100644 index 0271188af1..0000000000 --- a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_attrs.cc +++ /dev/null @@ -1,16 +0,0 @@ -#include "task-spec/dynamic_graph/dynamic_node_attrs.h" - -namespace FlexFlow { - -DynamicNodeAttrs - decide_dynamic_node_attrs_mapping(DynamicNodeAttrs const &attrs, - DynamicNodeMapping const &mapping) { - ASSERT(!attrs.mapping.has_value()); - - DynamicNodeAttrs result = attrs; - result.mapping = mapping; - - return result; -} - -} // namespace FlexFlow 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 beb04159ea..1bb588c08c 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 @@ -17,6 +17,7 @@ #include "utils/containers/flatmap.h" #include "utils/containers/generate_map.h" #include "utils/containers/get_only.h" +#include "utils/containers/invert_map.h" #include "utils/containers/map_from_keys_and_values.h" #include "utils/containers/map_from_pairs.h" #include "utils/containers/map_values.h" @@ -427,6 +428,35 @@ DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation( return result; } +DynamicNodeMapping + mapping_for_gradient_reduction(DynamicOpenDataflowGraph const &g, + DynamicValueAttrs const &value) { + DynamicSlotSite source_site = dynamic_graph_find_source_of_value(g, value); + std::set sink_sites = + dynamic_graph_find_sinks_of_value(g, value); + + DynamicNodeInvocation source = dynamic_graph_get_invocation_for_id( + g, source_site.require_internal().invocation_id); + DynamicNodeMapping source_mapping = assert_unwrap(source.node_attrs.mapping); + DynamicTensorSlot source_slot = + get_only(invert_map(source.outputs).at(value)); + + bidict + source_tensor_binding = get_tensor_bindings_for_slot_name( + source_mapping.op_task_group, source_slot.slot_name); + + DynamicNodeMapping result{ + /*op_task_group=*/MappedOperatorTaskGroup{ + bidict{ + {get_only(source_tensor_binding).second, + OperatorAtomicTaskShardBinding{ + std::map{ + {source_slot.slot_name, + get_only(source_tensor_binding).first}}}}}}, + /*device_type=*/source_mapping.device_type}; + return result; +} + DynamicNodeInvocation create_gradient_reduction_for_value( DynamicOpenDataflowGraph const &g, DynamicValueAttrs const &value, From b12e317c1206dbed6fa7edc107d8f76f7c46ef6f Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Wed, 26 Aug 2026 09:37:25 -0700 Subject: [PATCH 15/21] Cleanup. --- .../dynamic_graph/dynamic_node_mapping.h | 3 - .../task-spec/dynamic_graph/copy_insertion.cc | 35 +---------- .../dynamic_graph/dynamic_node_mapping.cc | 19 ------ .../task-spec/dynamic_graph/pass_expansion.cc | 59 ++++++++++++++----- 4 files changed, 47 insertions(+), 69 deletions(-) diff --git a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h index e9ab82eeb2..976153b558 100644 --- a/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h +++ b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h @@ -8,9 +8,6 @@ namespace FlexFlow { -DynamicNodeMapping dynamic_node_mapping_from_value_mapping( - std::map const &); - bidict dynamic_node_mapping_get_shard_bindings(DynamicNodeMapping const &); diff --git a/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc b/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc index 5c49c53232..2748e00ada 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc @@ -43,10 +43,6 @@ bool node_is_copy(DynamicNodeAttrs const &n) { return n.op_attrs.has_value() && n.op_attrs.value().is_copy(); } -bool node_is_gradient_reduction(DynamicNodeAttrs const &n) { - return n.op_attrs.has_value() && n.op_attrs.value().is_gradient_reduction(); -} - bool value_is_mapped(DynamicValueAttrs const &n) { return n.mapping.has_value(); } @@ -54,8 +50,7 @@ bool value_is_mapped(DynamicValueAttrs const &n) { void require_node_is_ready_for_copy_insertion( DynamicNodeAttrs const &node_attrs) { ASSERT(node_attrs.op_attrs.has_value()); - ASSERT(node_attrs.mapping.has_value() || - node_is_gradient_reduction(node_attrs)); + ASSERT(node_attrs.mapping.has_value()); } void require_value_is_ready_for_copy_insertion(DynamicValueAttrs const &v) { @@ -308,9 +303,7 @@ std::map dynamic_graph_get_id_for_invocation(g, invocation); std::set slot_sites_to_resolve = [&] { - if (op_type == TrainingOpType{TrainingOnlyOpType::GRADIENT_REDUCTION}) { - return std::set{}; - } else if (op_type == TrainingOpType{OperatorType::REPLICATE}) { + if (op_type == TrainingOpType{OperatorType::REPLICATE}) { return slots_to_map_for_replicate(invocation_id, invocation); } else { return get_dynamic_slot_sites_for_invocation(invocation_id, invocation); @@ -356,31 +349,9 @@ std::map TrainingOpType op_type = dynamic_node_invocation_get_op_type(invocation); std::optional task_type = invocation.node_attrs.task_type; - TrainingOpType gradient_reduction_op_type = - TrainingOpType{TrainingOnlyOpType::GRADIENT_REDUCTION}; TrainingOpType replicate_op_type = TrainingOpType{OperatorType::REPLICATE}; - if (op_type == gradient_reduction_op_type) { - if (slot_site.direction == TensorDirection::OUTPUT) { - InternalDynamicSlotSite slot_site_sink = - get_only(dynamic_graph_find_sinks_of_slot_site(g, slot_site)); - - ASSERT(contains_key(resolved_mappings, slot_site_sink)); - - return resolved_mappings.at(slot_site_sink); - } else { - ASSERT(slot_site.direction == TensorDirection::INCOMING); - - InternalDynamicSlotSite slot_site_src = - dynamic_graph_find_source_of_slot_site(g, slot_site) - .require_internal(); - - ASSERT(contains_key(resolved_mappings, slot_site_src)); - - return resolved_mappings.at(slot_site_src); - } - } else if (op_type == replicate_op_type && - task_type == DynamicTaskType::BWD) { + if (op_type == replicate_op_type && task_type == DynamicTaskType::BWD) { ASSERT(slot_site.direction == TensorDirection::OUTPUT); InternalDynamicSlotSite slot_site_sink = diff --git a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc index c9e6dbbefb..e66b61c6d4 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/dynamic_node_mapping.cc @@ -1,29 +1,10 @@ #include "task-spec/dynamic_graph/dynamic_node_mapping.h" -#include "task-spec/dynamic_graph/parallel_tensor_mapping.h" #include "utils/bidict/algorithms/bidict_transform_keys.h" #include "utils/bidict/algorithms/bidict_transform_values.h" -#include "utils/containers/flatmap.h" -#include "utils/containers/require_all_same.h" #include "utils/containers/transform.h" -#include "utils/optional.h" namespace FlexFlow { -DynamicNodeMapping dynamic_node_mapping_from_value_mapping( - std::map const &mapping) { - MappedOperatorTaskGroup op_task_group{{}}; - DeviceType device_type = assert_unwrap(require_all_same( - transform(flatmap(values(mapping), - [](ParallelTensorMapping const &m) { - return pt_mapping_get_device_set(m); - }), - [](global_device_id_t const &id) { return id.device_type; }))); - return DynamicNodeMapping{ - /*op_task_group=*/op_task_group, - /*device_type=*/device_type, - }; -} - bidict dynamic_node_mapping_get_shard_bindings(DynamicNodeMapping const &m) { return bidict_transform_keys( 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 1bb588c08c..a76820a6b6 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,4 +1,5 @@ #include "task-spec/dynamic_graph/pass_expansion.h" +#include "op-attrs/tensor_slot_name.dtg.h" #include "op-attrs/tensor_slot_name.h" #include "task-spec/dynamic_graph/copy_insertion.h" #include "task-spec/dynamic_graph/dynamic_invocation_id_t.dtg.h" @@ -10,7 +11,9 @@ #include "task-spec/dynamic_graph/dynamic_value_attrs.h" #include "task-spec/dynamic_graph/subgradient_id_t.dtg.h" #include "task-spec/dynamic_graph/training_operation_attrs.h" +#include "task-spec/fwb_tensor_type.dtg.h" #include "utils/containers/are_all_same.h" +#include "utils/containers/binary_merge_disjoint_maps.h" #include "utils/containers/concat_vectors.h" #include "utils/containers/contains_duplicates.h" #include "utils/containers/filter.h" @@ -428,9 +431,11 @@ DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation( return result; } -DynamicNodeMapping - mapping_for_gradient_reduction(DynamicOpenDataflowGraph const &g, - DynamicValueAttrs const &value) { +DynamicNodeMapping mapping_for_gradient_reduction( + DynamicOpenDataflowGraph const &g, + DynamicValueAttrs const &value, + std::map const + &subgradient_inputs) { DynamicSlotSite source_site = dynamic_graph_find_source_of_value(g, value); std::set sink_sites = dynamic_graph_find_sinks_of_value(g, value); @@ -441,18 +446,33 @@ DynamicNodeMapping DynamicTensorSlot source_slot = get_only(invert_map(source.outputs).at(value)); - bidict - source_tensor_binding = get_tensor_bindings_for_slot_name( - source_mapping.op_task_group, source_slot.slot_name); + std::pair + source_tensor_binding = get_only(get_tensor_bindings_for_slot_name( + source_mapping.op_task_group, source_slot.slot_name)); + + // For now we're going to map everything like the original source, both inputs + // and outputs. This means that in the presence of model parallelism we'll be + // potentially inducing extra copies that would otherwise be unnecessary. + std::set subgradient_ids = + set_of(values(subgradient_inputs)); + std::map input_bindings = + map_from_pairs(transform( + subgradient_ids, [&](subgradient_id_t const &subgradient_id) { + return std::pair{subgradient_id.gradient_reduction_slot, + source_tensor_binding.first}; + })); + std::map output_bindings{ + {TensorSlotName::OUTPUT, source_tensor_binding.first}, + }; + std::map bindings = + binary_merge_disjoint_maps(input_bindings, output_bindings); DynamicNodeMapping result{ /*op_task_group=*/MappedOperatorTaskGroup{ bidict{ - {get_only(source_tensor_binding).second, - OperatorAtomicTaskShardBinding{ - std::map{ - {source_slot.slot_name, - get_only(source_tensor_binding).first}}}}}}, + {source_tensor_binding.second, + OperatorAtomicTaskShardBinding{bindings}}, + }}, /*device_type=*/source_mapping.device_type}; return result; } @@ -494,18 +514,27 @@ DynamicNodeInvocation create_gradient_reduction_for_value( /*per_device_op_state=*/std::nullopt, }; + DynamicTensorSlot output_slot{ + /*slot_name=*/TensorSlotName::OUTPUT, + /*slot_tensor_role=*/std::nullopt, + /*task_shard=*/std::nullopt, + }; + return DynamicNodeInvocation{ /*inputs=*/map_from_pairs(transform( values(subgradient_inputs), [&](subgradient_id_t const &subgradient_id) { - return to_grad_subgradient(source_site.require_internal().slot_name, - value, - subgradient_id); + DynamicTensorSlot input_slot{ + /*slot_name=*/subgradient_id.gradient_reduction_slot, + /*slot_tensor_role=*/std::nullopt, + /*task_shard=*/std::nullopt, + }; + return to_grad_subgradient(input_slot, value, subgradient_id); })), /*node_attrs=*/gradient_reduction_attrs, /*outputs=*/ { - to_grad(source_site.require_internal().slot_name, value), + to_grad(output_slot, value), }, }; } From b43b6588ce40acbfe5f337a6516d00b8f3b59d5f Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Wed, 26 Aug 2026 09:53:40 -0700 Subject: [PATCH 16/21] Update test case. --- .../task-spec/dynamic_graph/pass_expansion.cc | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 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 c66777b7e6..ce424a261f 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,7 @@ #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_invocation_id_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/serializable_dynamic_node_invocation.h" @@ -386,6 +387,8 @@ TEST_SUITE(FF_TEST_SUITE) { }, }; + dynamic_invocation_id_t invocation_id{0_n}; + DynamicNodeInvocation invocation = [&]() -> DynamicNodeInvocation { return DynamicNodeInvocation{ /*inputs=*/{ @@ -409,8 +412,8 @@ TEST_SUITE(FF_TEST_SUITE) { }; }(); - DynamicNodeInvocation result = - perform_bwd_pass_expansion_for_invocation(invocation); + DynamicNodeInvocation result = perform_bwd_pass_expansion_for_invocation( + invocation_id, invocation, {}); DynamicNodeInvocation correct = [&]() -> DynamicNodeInvocation { return DynamicNodeInvocation{ @@ -452,6 +455,8 @@ TEST_SUITE(FF_TEST_SUITE) { }, }; + dynamic_invocation_id_t invocation_id{0_n}; + DynamicNodeInvocation invocation = [&]() -> DynamicNodeInvocation { return DynamicNodeInvocation{ /*inputs=*/{ @@ -473,8 +478,8 @@ TEST_SUITE(FF_TEST_SUITE) { }; }(); - DynamicNodeInvocation result = - perform_bwd_pass_expansion_for_invocation(invocation); + DynamicNodeInvocation result = perform_bwd_pass_expansion_for_invocation( + invocation_id, invocation, {}); DynamicNodeInvocation correct = [&]() -> DynamicNodeInvocation { DynamicTensorRole fwd_role = DynamicTensorRole{FwbTensorType::FORWARD}; @@ -508,6 +513,8 @@ TEST_SUITE(FF_TEST_SUITE) { SUBCASE("copy operator") { TrainingOperationAttrs op_attrs = TrainingOperationAttrs{CopyAttrs{}}; + dynamic_invocation_id_t invocation_id{0_n}; + DynamicNodeInvocation invocation = [&]() -> DynamicNodeInvocation { return DynamicNodeInvocation{ /*inputs=*/{ @@ -529,8 +536,8 @@ TEST_SUITE(FF_TEST_SUITE) { }; }(); - DynamicNodeInvocation result = - perform_bwd_pass_expansion_for_invocation(invocation); + DynamicNodeInvocation result = perform_bwd_pass_expansion_for_invocation( + invocation_id, invocation, {}); DynamicNodeInvocation correct = [&]() -> DynamicNodeInvocation { DynamicTensorRole fwd_role = DynamicTensorRole{FwbTensorType::FORWARD}; From 767d74f5ce5b5ad7e929acb18833ecec702e09c4 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Wed, 26 Aug 2026 11:53:44 -0700 Subject: [PATCH 17/21] Fixes for pass expansion. --- .../task-spec/dynamic_graph/pass_expansion.cc | 46 +++++++++++-------- 1 file changed, 26 insertions(+), 20 deletions(-) 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 a76820a6b6..fa88ac94ac 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 @@ -12,6 +12,7 @@ #include "task-spec/dynamic_graph/subgradient_id_t.dtg.h" #include "task-spec/dynamic_graph/training_operation_attrs.h" #include "task-spec/fwb_tensor_type.dtg.h" +#include "utils/containers/any_of.h" #include "utils/containers/are_all_same.h" #include "utils/containers/binary_merge_disjoint_maps.h" #include "utils/containers/concat_vectors.h" @@ -210,11 +211,14 @@ std::set determine_values_requiring_gradient_reduction( std::set internal_values = dynamic_graph_get_internal_values(g); return filter(internal_values, [&](DynamicValueAttrs const &v) { - dynamic_invocation_id_t source = dynamic_graph_find_source_of_value(g, v) - .require_internal() - .invocation_id; - return in_bwd_pass.count(source) != 0 && - dynamic_graph_find_sinks_of_value(g, v).size() > 1; + std::set sinks = + dynamic_graph_find_sinks_of_value(g, v); + if (sinks.size() <= 1) { + return false; + } + return any_of(sinks, [&](InternalDynamicSlotSite const &sink_site) { + return in_bwd_pass.count(sink_site.invocation_id); + }); }); } @@ -333,7 +337,7 @@ DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation( DynamicNodeInvocation const &invocation, std::map> const - &subgradient_ids) { + &subgradient_ids_by_value) { require_invocation_is_ready_for_pass_expansion(invocation); @@ -342,7 +346,7 @@ DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation( auto subgradient_id_for_value = [&](DynamicValueAttrs const &v) { return transform( - try_at(subgradient_ids, v), + try_at(subgradient_ids_by_value, v), [&](std::map const &s) { return s.at(invocation_id); }); @@ -552,7 +556,7 @@ DynamicOpenDataflowGraph std::map> - subgradient_ids = + subgradient_ids_by_value = compute_subgradient_ids(g, values_requiring_gradient_reduction); DynamicOpenDataflowGraph result = flatmap_dynamic_invocation_set( @@ -560,26 +564,28 @@ DynamicOpenDataflowGraph dynamic_invocation_id_t invocation_id = dynamic_graph_get_id_for_invocation(g, invocation); - if (contains(needed_in_bwd_pass, invocation_id)) { - std::set gradient_reductions = - transform(set_intersection(keys(subgradient_ids), - set_of(values(invocation.outputs))), - [&](DynamicValueAttrs const &v) { - return create_gradient_reduction_for_value( - g, v, subgradient_ids.at(v)); - }); + std::set gradient_reductions = + transform(set_intersection(keys(subgradient_ids_by_value), + set_of(values(invocation.outputs))), + [&](DynamicValueAttrs const &v) { + return create_gradient_reduction_for_value( + g, v, subgradient_ids_by_value.at(v)); + }); + if (contains(needed_in_bwd_pass, invocation_id)) { return set_union( std::set{ perform_fwd_pass_expansion_for_invocation(invocation), perform_bwd_pass_expansion_for_invocation( - invocation_id, invocation, subgradient_ids), + invocation_id, invocation, subgradient_ids_by_value), }, gradient_reductions); } else { - return std::set{ - perform_fwd_pass_expansion_for_invocation(invocation), - }; + return set_union( + std::set{ + perform_fwd_pass_expansion_for_invocation(invocation), + }, + gradient_reductions); } }); From 893f6acf73b6edcec8f0d9c5af0778efb757d033 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Wed, 26 Aug 2026 11:55:49 -0700 Subject: [PATCH 18/21] Cleanup. --- lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc | 1 - 1 file changed, 1 deletion(-) diff --git a/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc b/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc index 2748e00ada..ca6d4a05fb 100644 --- a/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc +++ b/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc @@ -15,7 +15,6 @@ #include "task-spec/dynamic_graph/dynamic_value_attrs.dtg.h" #include "task-spec/dynamic_graph/dynamic_value_attrs.h" #include "task-spec/dynamic_graph/parallel_tensor_mapping.dtg.h" -#include "task-spec/dynamic_graph/training_only_op_type.dtg.h" #include "task-spec/dynamic_graph/training_operation_attrs.h" #include "utils/bidict/algorithms/bidict_from_unstructured_relation.h" #include "utils/bidict/algorithms/unstructured_relation_from_bidict.h" From ac746b085100161c054d7a54491920a6f899fbf4 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Wed, 26 Aug 2026 12:01:51 -0700 Subject: [PATCH 19/21] Actually invoke mapping. --- lib/task-spec/src/task-spec/dynamic_graph/pass_expansion.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 fa88ac94ac..f153e5dedd 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 @@ -510,7 +510,7 @@ DynamicNodeInvocation create_gradient_reduction_for_value( DynamicNodeAttrs gradient_reduction_attrs{ /*task_type=*/DynamicTaskType::BWD, /*device_ids=*/std::nullopt, - /*mapping=*/std::nullopt, + /*mapping=*/mapping_for_gradient_reduction(g, value, subgradient_inputs), /*op_attrs=*/ TrainingOperationAttrs{GradientReductionAttrs{}}, /*layer_guid=*/ From 0176b6d8779088f03e5a1b7988ce78978bc44a43 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Wed, 26 Aug 2026 12:04:55 -0700 Subject: [PATCH 20/21] Don't compute mapping if not present. --- .../src/task-spec/dynamic_graph/pass_expansion.cc | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) 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 f153e5dedd..3c46e82569 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 @@ -435,7 +435,7 @@ DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation( return result; } -DynamicNodeMapping mapping_for_gradient_reduction( +std::optional mapping_for_gradient_reduction( DynamicOpenDataflowGraph const &g, DynamicValueAttrs const &value, std::map const @@ -446,7 +446,12 @@ DynamicNodeMapping mapping_for_gradient_reduction( DynamicNodeInvocation source = dynamic_graph_get_invocation_for_id( g, source_site.require_internal().invocation_id); + + if (!source.node_attrs.mapping.has_value()) { + return std::nullopt; + } DynamicNodeMapping source_mapping = assert_unwrap(source.node_attrs.mapping); + DynamicTensorSlot source_slot = get_only(invert_map(source.outputs).at(value)); From 20f7f89aba235997afad619a241d1c3a7971495e Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Wed, 26 Aug 2026 12:10:28 -0700 Subject: [PATCH 21/21] Header cleanup. --- .../src/task-spec/dynamic_graph/pass_expansion.cc | 15 --------------- 1 file changed, 15 deletions(-) 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 3c46e82569..5c52f97327 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,22 +1,12 @@ #include "task-spec/dynamic_graph/pass_expansion.h" -#include "op-attrs/tensor_slot_name.dtg.h" #include "op-attrs/tensor_slot_name.h" -#include "task-spec/dynamic_graph/copy_insertion.h" -#include "task-spec/dynamic_graph/dynamic_invocation_id_t.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/dynamic_value_attrs.h" -#include "task-spec/dynamic_graph/subgradient_id_t.dtg.h" #include "task-spec/dynamic_graph/training_operation_attrs.h" -#include "task-spec/fwb_tensor_type.dtg.h" #include "utils/containers/any_of.h" -#include "utils/containers/are_all_same.h" #include "utils/containers/binary_merge_disjoint_maps.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" @@ -25,9 +15,7 @@ #include "utils/containers/map_from_keys_and_values.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_intersection.h" @@ -35,9 +23,6 @@ #include "utils/containers/slice.h" #include "utils/containers/transform.h" #include "utils/containers/try_at.h" -#include "utils/containers/zip_with.h" -#include "utils/optional.h" -#include namespace FlexFlow {