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_node_mapping.h b/lib/task-spec/include/task-spec/dynamic_graph/dynamic_node_mapping.h index 3b28ea5149..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 @@ -2,6 +2,8 @@ #define _FLEXFLOW_LIB_TASK_SPEC_INCLUDE_TASK_SPEC_DYNAMIC_GRAPH_DYNAMIC_NODE_MAPPING_H #include "task-spec/dynamic_graph/dynamic_node_mapping.dtg.h" +#include "task-spec/dynamic_graph/dynamic_tensor_slot.dtg.h" +#include "task-spec/dynamic_graph/parallel_tensor_mapping.dtg.h" #include "task-spec/global_device_id_t.dtg.h" namespace FlexFlow { 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/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/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/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/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/copy_insertion.cc b/lib/task-spec/src/task-spec/dynamic_graph/copy_insertion.cc index ef5695792f..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 @@ -269,7 +269,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 +289,7 @@ std::map { InternalDynamicSlotSite slot_site = get_only(slot_sites); ASSERT(slot_site.slot_name.slot_name == TensorSlotName::OUTPUT); - }; + } return slot_sites; }; 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/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..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,15 +1,28 @@ #include "task-spec/dynamic_graph/pass_expansion.h" +#include "op-attrs/tensor_slot_name.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.h" #include "task-spec/dynamic_graph/training_operation_attrs.h" -#include "utils/containers/are_all_same.h" +#include "utils/containers/any_of.h" +#include "utils/containers/binary_merge_disjoint_maps.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/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" #include "utils/containers/merge_disjoint_maps.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/slice.h" #include "utils/containers/transform.h" +#include "utils/containers/try_at.h" namespace FlexFlow { @@ -27,6 +40,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 +53,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( @@ -175,6 +190,51 @@ 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) { + 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); + }); + }); +} + +static std::map + assign_subgradient_ids( + std::set const &invocations) { + 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> + 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); @@ -258,13 +318,37 @@ 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_by_value) { 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_by_value, v), + [&](std::map const &s) { + return s.at(invocation_id); + }); + }; + + 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), @@ -276,6 +360,10 @@ DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation( return pass_expand_value(v, FwbTensorType::GRADIENT); }; + auto to_grad_value_output = [&](DynamicValueAttrs const &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), @@ -283,12 +371,20 @@ DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation( }; }; + auto to_grad_output = [&](DynamicTensorSlot const &k, + DynamicValueAttrs const &v) { + return std::pair{ + pass_expand_slot(k, FwbTensorType::GRADIENT), + pass_expand_value_with_subgradient(v, FwbTensorType::GRADIENT), + }; + }; + DynamicNodeInvocation result = [&]() -> DynamicNodeInvocation { if (op_attrs.is_copy()) { 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)) { @@ -300,7 +396,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 { @@ -314,7 +410,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), }; }; }(); @@ -324,6 +420,119 @@ DynamicNodeInvocation perform_bwd_pass_expansion_for_invocation( return result; } +std::optional 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); + + 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)); + + 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{ + {source_tensor_binding.second, + OperatorAtomicTaskShardBinding{bindings}}, + }}, + /*device_type=*/source_mapping.device_type}; + return result; +} + +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); + + auto to_grad = [](DynamicTensorSlot const &k, DynamicValueAttrs const &v) { + return std::pair{ + pass_expand_slot(k, FwbTensorType::GRADIENT), + pass_expand_value(v, FwbTensorType::GRADIENT), + }; + }; + + auto to_grad_subgradient = [](DynamicTensorSlot const &k, + DynamicValueAttrs const &v, + subgradient_id_t const &s) { + return std::pair{ + pass_expand_slot(k, FwbTensorType::GRADIENT), + decide_dynamic_value_attrs_subgradient_id( + pass_expand_value(v, FwbTensorType::GRADIENT), s), + }; + }; + + DynamicNodeAttrs gradient_reduction_attrs{ + /*task_type=*/DynamicTaskType::BWD, + /*device_ids=*/std::nullopt, + /*mapping=*/mapping_for_gradient_reduction(g, value, subgradient_inputs), + /*op_attrs=*/ + TrainingOperationAttrs{GradientReductionAttrs{}}, + /*layer_guid=*/ + dynamic_layer_guid_t{dynamic_gradient_reduction_layer_guid_t{}}, + /*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) { + 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(output_slot, value), + }, + }; +} + DynamicOpenDataflowGraph perform_pass_expansion(DynamicOpenDataflowGraph const &g) { @@ -332,20 +541,41 @@ DynamicOpenDataflowGraph std::set needed_in_bwd_pass = determine_invocations_needed_in_backward_pass_for_gradient_computation(g); + std::set values_requiring_gradient_reduction = + determine_values_requiring_gradient_reduction(g, needed_in_bwd_pass); + + std::map> + subgradient_ids_by_value = + 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); + 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 std::set{ - perform_fwd_pass_expansion_for_invocation(invocation), - perform_bwd_pass_expansion_for_invocation(invocation), - }; + return set_union( + std::set{ + perform_fwd_pass_expansion_for_invocation(invocation), + perform_bwd_pass_expansion_for_invocation( + 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); } }); 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..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" @@ -56,6 +57,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 +183,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 +340,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, @@ -383,6 +387,8 @@ TEST_SUITE(FF_TEST_SUITE) { }, }; + dynamic_invocation_id_t invocation_id{0_n}; + DynamicNodeInvocation invocation = [&]() -> DynamicNodeInvocation { return DynamicNodeInvocation{ /*inputs=*/{ @@ -406,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{ @@ -449,6 +455,8 @@ TEST_SUITE(FF_TEST_SUITE) { }, }; + dynamic_invocation_id_t invocation_id{0_n}; + DynamicNodeInvocation invocation = [&]() -> DynamicNodeInvocation { return DynamicNodeInvocation{ /*inputs=*/{ @@ -470,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}; @@ -505,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=*/{ @@ -526,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}; @@ -587,6 +597,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 +905,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,