diff --git a/lib/realm-execution/src/realm-execution/pcg_instance.cc b/lib/realm-execution/src/realm-execution/pcg_instance.cc index 1d102f8cca..a037ca2654 100644 --- a/lib/realm-execution/src/realm-execution/pcg_instance.cc +++ b/lib/realm-execution/src/realm-execution/pcg_instance.cc @@ -341,6 +341,7 @@ static Realm::Event spawn_dynamic_node_invocation( }, [&](LossAttrs const &) { return spawn_task(); }, [&](CopyAttrs const &) { return issue_copy(); }, + [&](GradientReductionAttrs const &) { return issue_reduction(); }, }); } 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 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/dynamic_value_attrs.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_value_attrs.dtg.toml index 41638fd4ef..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 @@ -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/subgradient_id_t.dtg.h", "op-attrs/parallel_tensor_space_coordinate.dtg.h", "pcg/machine_space_coordinate.dtg.h", "utils/bidict/bidict.h", @@ -48,6 +49,15 @@ For a \ref DynamicOpenDataflowGraph originating form a \ref MappedParallelComput name = "create_grad" type = "std::optional" +[[fields]] +name = "subgradient_id" +type = "std::optional<::FlexFlow::subgradient_id_t>" +docstring = ''' +\brief The unique identity of the subgradient represented by this value. + +This field is filled in by \ref pass_expansion.h when a value is consumed multiple times, requiring a \ref GradientReductionAttrs in the backward pass. +''' + [[fields]] name = "shard_coord" type = "std::optional<::FlexFlow::ParallelTensorSpaceCoordinate>" diff --git a/lib/task-spec/include/task-spec/dynamic_graph/gradient_reduction_attrs.dtg.toml b/lib/task-spec/include/task-spec/dynamic_graph/gradient_reduction_attrs.dtg.toml new file mode 100644 index 0000000000..82f759fd19 --- /dev/null +++ b/lib/task-spec/include/task-spec/dynamic_graph/gradient_reduction_attrs.dtg.toml @@ -0,0 +1,13 @@ +namespace = "FlexFlow" +name = "GradientReductionAttrs" +type = "struct" +features = [ + "eq", + "ord", + "hash", + "json", + "fmt", + "rapidcheck", +] + +fields = [] 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..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,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/subgradient_id_t.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::subgradient_id_t>" + [[fields]] name = "shard_coord" type = "std::optional<::FlexFlow::ParallelTensorSpaceCoordinate>" 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..61fd935996 --- /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 TensorSlotName as it is used as an input to \ref GradientReduction. +""" + +includes = [ + "op-attrs/tensor_slot_name.dtg.h", +] + +src_includes = [] + +[[fields]] +name = "gradient_reduction_slot" +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..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,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_attrs.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/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..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 @@ -4,12 +4,21 @@ #include "task-spec/dynamic_graph/dynamic_tensor_role.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/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" namespace FlexFlow { @@ -27,6 +36,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) { @@ -39,6 +49,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( @@ -324,6 +335,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_for_invocations( + 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_for_invocations(new_invocation_set); + + return dynamic_open_dataflow_graph_from_invocation_set( + set_of(new_invocation_set)); +} + DynamicOpenDataflowGraph perform_pass_expansion(DynamicOpenDataflowGraph const &g) { @@ -332,22 +481,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/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/src/task-spec/dynamic_graph/training_operation_attrs.cc b/lib/task-spec/src/task-spec/dynamic_graph/training_operation_attrs.cc index 3bdfe0d30a..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 @@ -12,6 +12,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 +30,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/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 d3a915fb5c..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 @@ -56,6 +56,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, @@ -181,6 +182,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, @@ -337,6 +339,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, @@ -587,6 +590,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, @@ -894,4 +898,339 @@ 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_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, + std::optional const &subgradient_id) + -> 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, + /*subgradient_id=*/subgradient_id, + /*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, 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 { + 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); + + 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(11, relu_op_attrs, DynamicTaskType::BWD); + DynamicNodeAttrs relu2_node_bwd = + mk_node_attrs(12, 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(), 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(), std::nullopt); + DynamicValueAttrs relu1_output_tensor_activation = + 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(), std::nullopt); + DynamicValueAttrs relu2_output_tensor_activation = + 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(), std::nullopt); + + 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_grad_slot(TensorSlotName::OUTPUT), + relu1_output_tensor_gradient, + }, + }, + /*node_attrs=*/relu1_node_bwd, + /*outputs=*/ + std::map{ + std::pair{ + mk_grad_slot(TensorSlotName::INPUT), + input_tensor_subgradient1, + }, + }, + }, + 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_grad_slot(TensorSlotName::OUTPUT), + relu2_output_tensor_gradient, + }, + }, + /*node_attrs=*/relu2_node_bwd, + /*outputs=*/ + std::map{ + std::pair{ + mk_grad_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, + }, + }, + }, + }; + + 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)); + } } 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,