diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h index 048218744f..5f41d7f819 100644 --- a/include/onnxruntime/core/graph/graph.h +++ b/include/onnxruntime/core/graph/graph.h @@ -264,6 +264,9 @@ class Node { /** Gets the Node's attributes. */ const NodeAttributes& GetAttributes() const noexcept; + /** Gets the Node's mutable attributes. */ + NodeAttributes& GetMutableAttributes() noexcept; + /** Gets the Graph instance that is instantiated from a GraphProto attribute during Graph::Resolve. @param attr_name Attribute name for the GraphProto attribute. @returns nullptr if the Graph instance has not been instantiated or attribute does not contain a GraphProto. diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index cf1a90c051..0e11eb1ba3 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -654,6 +654,10 @@ const NodeAttributes& Node::GetAttributes() const noexcept { return attributes_; } +NodeAttributes& Node::GetMutableAttributes() noexcept { + return attributes_; +} + Graph* Node::GetMutableGraphAttribute(const std::string& attr_name) { Graph* subgraph = nullptr; diff --git a/orttraining/orttraining/core/graph/gradient_schema_defs.cc b/orttraining/orttraining/core/graph/gradient_schema_defs.cc index e3bc26cf6b..e49c858181 100644 --- a/orttraining/orttraining/core/graph/gradient_schema_defs.cc +++ b/orttraining/orttraining/core/graph/gradient_schema_defs.cc @@ -219,7 +219,7 @@ OpSchema& RegisterLambOpSchema(OpSchema&& op_schema) { } } - // Handle other tensors including new weight, new gradient (update direction), + // Handle other tensors including new weight, new gradient (update direction), // new momentums. for (size_t i = 0; i < ctx.getNumInputs() - 5; ++i) { const size_t input_index = 5 + i; // The first 5 inputs don't affect output shape. diff --git a/orttraining/orttraining/core/graph/mixed_precision_transformer.cc b/orttraining/orttraining/core/graph/mixed_precision_transformer.cc index e79db74a92..82b66eb7a4 100644 --- a/orttraining/orttraining/core/graph/mixed_precision_transformer.cc +++ b/orttraining/orttraining/core/graph/mixed_precision_transformer.cc @@ -415,6 +415,24 @@ Status TransformGraphForMixedPrecision(Graph& graph, graph.AddInitializedTensor(weight_tensor_proto); } + // Handle pipeline case + for (auto& node : graph.Nodes()) { + // For send and recv node, if the tensor being sent or received is FP32, update its + // attribute and change it to FP16. + if (!node.OpType().compare("Send") || !node.OpType().compare("Recv")) { + auto& attributes = node.GetMutableAttributes(); + auto* element_type = &(attributes.find("element_types")->second); + int ints_size = element_type->ints_size(); + for(int i=0;iints(i) == static_cast(ONNX_NAMESPACE::TensorProto_DataType_FLOAT)){ + element_type->set_ints(i, static_cast(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16)); + // Need to resolve and populate the new type through the graph. + graph.SetGraphResolveNeeded(); + } + } + } + } + // Handle implicit data type casting nodes such as Cast, ConstantOfShape ORT_RETURN_IF_ERROR(TransformConstants(graph)); diff --git a/orttraining/orttraining/core/graph/pipeline_transformer.cc b/orttraining/orttraining/core/graph/pipeline_transformer.cc index 4ca48178e9..6fde902670 100644 --- a/orttraining/orttraining/core/graph/pipeline_transformer.cc +++ b/orttraining/orttraining/core/graph/pipeline_transformer.cc @@ -464,6 +464,7 @@ Status AddBackwardRecordBeforeSend( } Status SetInputsOutputsAndResolve(Graph& graph, + const std::unordered_set& weights_to_train, const std::vector& new_input_names, const std::vector& new_output_names) { auto fill_node_args = [&](const Graph& graph, @@ -491,7 +492,14 @@ Status SetInputsOutputsAndResolve(Graph& graph, graph.SetGraphResolveNeeded(); graph.SetGraphProtoSyncNeeded(); - return graph.Resolve(); + Graph::ResolveOptions options; + // Reserve the training weights. In mixed precision case, without this field, + // the original fp32 initializers could be removed due to not being used + // at this point. But we still need to preserve them because later when optimizer is + // is constructed, the isolated fp32 initializers will be inputs for optimizer. + options.initializer_names_to_preserve = &weights_to_train; + + return graph.Resolve(options); } // This function inserts WaitEvent's and RecordEvent's to the input graph for @@ -542,6 +550,7 @@ Status SetInputsOutputsAndResolve(Graph& graph, // Record-3: Tell others that backward result has been passed to another stage. Status TransformGraphForPipeline( Graph& graph, + const std::unordered_set& weights_to_train, std::string& forward_waited_event_name, std::string& forward_recorded_event_name, std::string& backward_waited_event_name, @@ -727,7 +736,7 @@ Status TransformGraphForPipeline( )); } - ORT_RETURN_IF_ERROR(SetInputsOutputsAndResolve(graph, new_input_names, new_output_names)); + ORT_RETURN_IF_ERROR(SetInputsOutputsAndResolve(graph, weights_to_train, new_input_names, new_output_names)); return Status::OK(); } @@ -887,15 +896,8 @@ common::Status SplitGraph(Graph& graph, send_input_args.push_back(updated_node_arg); auto dtype = original_node_arg->TypeAsProto()->tensor_type().elem_type(); - switch (dtype) { - case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: - element_types.add_ints(static_cast(1)); - break; - default: - // Assume all tensors are of type float. - // TODO: update if graph supports other data type. - ORT_THROW("pipeline partition unsupported 'type' value: ", dtype); - } + + element_types.add_ints(static_cast(dtype)); auto& new_receive_output = CreateNodeArg(graph, updated_node_arg); const auto old_shape = *(updated_node_arg->Shape()); @@ -957,7 +959,7 @@ common::Status SplitGraph(Graph& graph, recv_nodes.push_back(&recv_node); } - ORT_RETURN_IF_ERROR(SetInputsOutputsAndResolve(graph, new_input_names, new_output_names)); + ORT_RETURN_IF_ERROR(SetInputsOutputsAndResolve(graph, {} /* weights_to_train */, new_input_names, new_output_names)); return Status::OK(); } diff --git a/orttraining/orttraining/core/graph/pipeline_transformer.h b/orttraining/orttraining/core/graph/pipeline_transformer.h index bbc424439b..be58c66682 100644 --- a/orttraining/orttraining/core/graph/pipeline_transformer.h +++ b/orttraining/orttraining/core/graph/pipeline_transformer.h @@ -13,6 +13,7 @@ void GetPipelineSendOutput(const Graph& graph, std::string& loss_name); Status TransformGraphForPipeline( Graph& graph, + const std::unordered_set& initializer_names_to_preserve, std::string& forward_waited_event_name, std::string& forward_recorded_event_name, std::string& backward_waited_event_name, diff --git a/orttraining/orttraining/core/session/training_session.cc b/orttraining/orttraining/core/session/training_session.cc index e200866088..d8a0eeceea 100644 --- a/orttraining/orttraining/core/session/training_session.cc +++ b/orttraining/orttraining/core/session/training_session.cc @@ -142,8 +142,15 @@ Status TrainingSession::ConfigureForTraining( is_mixed_precision_enabled_ = config.mixed_precision_config.has_value(); std::string loss_name{}; + // Enable loss scale if mixed precision is enabled AND at pipeline last stage if pipeline is used. + // We are currently making the assumption that no data parallelism is used together with model parallelism. + // So we can check the last stage by checking the world_rank and world_size. Once DP and MP combination is + // enabled, we need to devise another way to check MP stages. + bool enable_loss_scale = is_mixed_precision_enabled_ && + (!config.pipeline_config.has_value() || + (config.distributed_config.world_rank + 1 == config.distributed_config.world_size)); optional loss_scale_input_name = - is_mixed_precision_enabled_ ? optional{""} : optional{}; + enable_loss_scale ? optional{""} : optional{}; if (config.pipeline_config.has_value()) { // if use pipeline, first check if model contains send op. If it does, set the // send node's output as the start tensor to build gradient graph @@ -164,7 +171,7 @@ Status TrainingSession::ConfigureForTraining( !loss_scale_input_name.has_value() || !loss_scale_input_name.value().empty(), "loss_scale_input_name should not be set to an empty string."); - if (is_mixed_precision_enabled_) { + if (enable_loss_scale) { TrainingConfigurationResult::MixedPrecisionConfigurationResult mp_result{}; mp_result.loss_scale_input_name = loss_scale_input_name.value(); config_result.mixed_precision_config_result = mp_result; @@ -225,7 +232,8 @@ Status TrainingSession::ConfigureForTraining( if (config.pipeline_config.has_value()) { TrainingConfigurationResult::PipelineConfigurationResult pipeline_result{}; - ORT_RETURN_IF_ERROR(InsertPipelineOps(pipeline_result.forward_waited_event_name, + ORT_RETURN_IF_ERROR(InsertPipelineOps(weight_names_to_train, + pipeline_result.forward_waited_event_name, pipeline_result.forward_recorded_event_name, pipeline_result.backward_waited_event_name, pipeline_result.backward_recorded_event_name, @@ -259,7 +267,8 @@ Status TrainingSession::ConfigureForTraining( for (auto it = weights_to_train_.begin(); it != weights_to_train_.end();) { const auto* node_arg = model_->MainGraph().GetNodeArg(*it); ORT_RETURN_IF_NOT(node_arg, "Failed to get NodeArg with name ", *it); - if (node_arg->TypeAsProto()->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT) { + if (node_arg->TypeAsProto()->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT && + node_arg->TypeAsProto()->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) { it = weights_to_train_.erase(it); } else{ @@ -549,6 +558,7 @@ Status TrainingSession::AddTensorboard(const std::string& summary_name, } Status TrainingSession::InsertPipelineOps( + const std::unordered_set& initializer_names_to_preserve, std::string& forward_waited_event_name, std::string& forward_recorded_event_name, std::string& backward_waited_event_name, @@ -563,6 +573,7 @@ Status TrainingSession::InsertPipelineOps( std::string& backward_recorded_event_before_send_name) { ORT_RETURN_IF_ERROR(TransformGraphForPipeline( model_->MainGraph(), + initializer_names_to_preserve, forward_waited_event_name, forward_recorded_event_name, backward_waited_event_name, diff --git a/orttraining/orttraining/core/session/training_session.h b/orttraining/orttraining/core/session/training_session.h index b5c03c51a1..266c0ed996 100644 --- a/orttraining/orttraining/core/session/training_session.h +++ b/orttraining/orttraining/core/session/training_session.h @@ -370,7 +370,8 @@ class TrainingSession : public InferenceSession { // 3. Backward operators' descriptions are all "Backward pass". This assumption is used to // identify backward nodes. // 4. No event operator is inserted by other graph transform. - common::Status InsertPipelineOps(std::string& forward_waited_event_name, + common::Status InsertPipelineOps(const std::unordered_set& initializer_names_to_preserve, + std::string& forward_waited_event_name, std::string& forward_recorded_event_name, std::string& backward_waited_event_name, std::string& backward_recorded_event_name, diff --git a/orttraining/orttraining/models/bert/main.cc b/orttraining/orttraining/models/bert/main.cc index edf2bcdf8b..b4894156ee 100644 --- a/orttraining/orttraining/models/bert/main.cc +++ b/orttraining/orttraining/models/bert/main.cc @@ -504,13 +504,22 @@ const std::map> input_to_dimension_m MapStringToString mapped_dimensions; void setup_training_params(BertParameters& params) { - params.model_path = ToPathString(params.model_name) + ORT_TSTR(".onnx"); - params.model_with_loss_func_path = ToPathString(params.model_name) + ORT_TSTR("_with_cost.onnx"); - params.model_with_training_graph_path = ToPathString(params.model_name) + ORT_TSTR("_bw.onnx"); - params.model_actual_running_graph_path = ToPathString(params.model_name) + ORT_TSTR("_bw_running.onnx"); + auto model_name_base = ToPathString(params.model_name); + params.model_path = model_name_base + ORT_TSTR(".onnx"); + params.model_with_loss_func_path = model_name_base + ORT_TSTR("_with_cost.onnx"); + params.model_with_training_graph_path = model_name_base + ORT_TSTR("_bw.onnx"); + params.model_actual_running_graph_path = model_name_base + ORT_TSTR("_bw_running.onnx"); #ifdef USE_HOROVOD params.mpi_context = setup_horovod(); + + if (params.pipeline_parallel_size > 1) { + auto pipeline_model_name_base = model_name_base + ToPathString(std::to_string(params.mpi_context.world_rank)); + params.model_with_loss_func_path = pipeline_model_name_base + ORT_TSTR("_with_cost.onnx"); + params.model_with_training_graph_path = pipeline_model_name_base + ORT_TSTR("_bw.onnx"); + params.model_actual_running_graph_path = pipeline_model_name_base + ORT_TSTR("_bw_running.onnx"); + } + ORT_ENFORCE(params.horizontal_parallel_size <= params.mpi_context.world_size); ORT_ENFORCE(params.data_parallel_size <= params.mpi_context.world_size); if (params.mpi_context.world_size % params.horizontal_parallel_size != 0) { diff --git a/orttraining/orttraining/models/runner/training_runner.cc b/orttraining/orttraining/models/runner/training_runner.cc index e7c19d8008..37385ecc8b 100644 --- a/orttraining/orttraining/models/runner/training_runner.cc +++ b/orttraining/orttraining/models/runner/training_runner.cc @@ -561,16 +561,16 @@ Status TrainingRunner::PrepareFetchNamesAndFetches(const SessionMode mode, } else { // No pipeline. All fetched names should appear in the graph handled by this process. fetch_names = params_.fetch_names; - } - if (params_.use_mixed_precision) { - auto it = opt_graph_outputs_.find(OptimizerOutputKey::GradientAllIsFinite); - ORT_RETURN_IF(it == opt_graph_outputs_.end(), "Gradient norm's IsFinite output is missing in the optimizer output"); - fetch_names.push_back(it->second); - if (params_.use_adasum) { - it = opt_graph_outputs_.find(OptimizerOutputKey::DeltaAllIsFinite); - ORT_RETURN_IF(it == opt_graph_outputs_.end(), "Adasum delta's IsFinite output is missing in the optimizer output"); + if (params_.use_mixed_precision) { + auto it = opt_graph_outputs_.find(OptimizerOutputKey::GradientAllIsFinite); + ORT_RETURN_IF(it == opt_graph_outputs_.end(), "Gradient norm's IsFinite output is missing in the optimizer output"); fetch_names.push_back(it->second); + if (params_.use_adasum) { + it = opt_graph_outputs_.find(OptimizerOutputKey::DeltaAllIsFinite); + ORT_RETURN_IF(it == opt_graph_outputs_.end(), "Adasum delta's IsFinite output is missing in the optimizer output"); + fetch_names.push_back(it->second); + } } } } else if (mode == GradientAccumulateStep) { @@ -667,7 +667,7 @@ void TrainingRunner::RunWithUpdate(VectorString& feed_names, &(pipeline_worker_pool_.worker_states[worker_id].fetches)); }, worker_id, step_); - // Wait all workers to finish this round of pipeline parallelism. + // Wait all workers to finish this round of pipeline parallelism. // The last batch in a pipeline collects gradient and update the model. // We must join here because main thread needs to access thread-produced // fetches and those fetches must be ready. @@ -831,33 +831,32 @@ Status TrainingRunner::TrainingLoop(IDataLoader& training_data_loader, IDataLoad auto start = std::chrono::high_resolution_clock::now(); if (is_weight_update_step) { - PrepareFeedNamesAndFeeds(ModelUpdateStep, + ORT_RETURN_IF_ERROR(PrepareFeedNamesAndFeeds(ModelUpdateStep, training_data_loader, *training_data, lr_scheduler.get(), batch, feed_names, - feeds); + feeds)); ORT_RETURN_IF_ERROR( PrepareFetchNamesAndFetches(ModelUpdateStep, fetch_names, fetches)); RunWithUpdate(feed_names, fetch_names, feeds, fetches); } else { - PrepareFeedNamesAndFeeds(GradientAccumulateStep, + ORT_RETURN_IF_ERROR(PrepareFeedNamesAndFeeds(GradientAccumulateStep, training_data_loader, *training_data, lr_scheduler.get(), batch, feed_names, - feeds); + feeds)); ORT_RETURN_IF_ERROR( PrepareFetchNamesAndFetches(GradientAccumulateStep, fetch_names, fetches)); RunWithoutUpdate(feed_names, fetch_names, feeds, gradient_accumulation_step_count); - } // at this point, step_ already be increased by 1. @@ -946,7 +945,7 @@ Status TrainingRunner::TrainingLoop(IDataLoader& training_data_loader, IDataLoad ORT_RETURN_IF_ERROR(Env::Default().CreateFolder(params_.perf_output_dir)); // saving json file ORT_RETURN_IF_ERROR(SavePerfMetrics(number_of_batches, gradient_accumulation_step_count, weight_update_steps, - total_time, avg_time_per_batch, throughput, stabilized_throughput, + total_time, avg_time_per_batch, throughput, stabilized_throughput, e2e_throughput, mapped_dimensions, average_cpu_usage, peak_workingset_size)); } @@ -1138,7 +1137,7 @@ Status TrainingRunner::Evaluate(InferenceSession& session, IDataLoader& data_loa if (params_.pipeline_parallel_size == 1) { auto status = Status::OK(); // When there is no pipeline, we always use the first thread - // to launch session_.Run(...) to avoid multiple activation allocations. + // to launch session_.Run(...) to avoid multiple activation allocations. // Always use the first thread to evaluate. const size_t worker_id = 0; diff --git a/orttraining/orttraining/test/graph/gradient_graph_builder_test.cc b/orttraining/orttraining/test/graph/gradient_graph_builder_test.cc index b2f888557c..1e2fb4f9e3 100644 --- a/orttraining/orttraining/test/graph/gradient_graph_builder_test.cc +++ b/orttraining/orttraining/test/graph/gradient_graph_builder_test.cc @@ -425,7 +425,7 @@ static void RunBertTrainingWithChecks( std::vector fetches; - EXPECT_TRUE(training_session->Run(run_options, feed_names, feeds, fetch_names, &fetches).IsOK()); + ASSERT_STATUS_OK(training_session->Run(run_options, feed_names, feeds, fetch_names, &fetches)); for (size_t i = 0; i < fetch_names.size(); ++i) { if (!fetches[i].IsAllocated() || !!fetches[i].IsTensor()) @@ -1127,31 +1127,52 @@ TEST(GradientGraphBuilderTest, PipelineOnlinePartition) { pipe.cut_list.emplace_back(cut0); pipe.cut_list.emplace_back(cut1); - for (int i = 0; i < 3; ++i) { + TrainingSession::TrainingConfiguration::MixedPrecisionConfiguration mixed_precision_config{}; + mixed_precision_config.use_fp16_initializers = true; + + // 2 test variations - full precision and mixed precision + const std::vector test_with_fp32{true, false}; + for(auto is_fp32 : test_with_fp32) { + // graph is partitioned into 3 parts. + for (int i = 0; i < 3; ++i) { #ifdef _WIN32 - auto surfix = std::to_wstring(i); + auto surfix = std::to_wstring(i); #else - auto surfix = std::to_string(i); + auto surfix = std::to_string(i); #endif - PathString output_file = ORT_TSTR("pipeline_partition_") + surfix + ORT_TSTR("_back.onnx"); + PathString output_file = ORT_TSTR("pipeline_partition_") + surfix + ORT_TSTR("_back.onnx"); - auto config = MakeBasicTrainingConfig(); - config.pipeline_config = pipe; - config.distributed_config.world_rank = i; - config.distributed_config.world_size = 3; - config.distributed_config.local_rank = i; - config.distributed_config.local_size = 3; - config.distributed_config.data_parallel_size = 1; - config.distributed_config.horizontal_parallel_size = 1; - config.distributed_config.pipeline_parallel_size = 3; - config.model_with_training_graph_path = output_file; + auto config = MakeBasicTrainingConfig(); + config.pipeline_config = pipe; + config.distributed_config.world_rank = i; + config.distributed_config.world_size = 3; + config.distributed_config.local_rank = i; + config.distributed_config.local_size = 3; + config.distributed_config.data_parallel_size = 1; + config.distributed_config.horizontal_parallel_size = 1; + config.distributed_config.pipeline_parallel_size = 3; + config.model_with_training_graph_path = output_file; - PathString backprop_model_file; - ASSERT_STATUS_OK(BuildBackPropGraph(model_uri, config, backprop_model_file)); + if (!is_fp32) { + config.mixed_precision_config = mixed_precision_config; + } - std::shared_ptr model; - // Ensure the partitioned model load. - ASSERT_STATUS_OK(Model::Load(backprop_model_file, model, nullptr, DefaultLoggingManager().DefaultLogger())); + PathString backprop_model_file; + Status status = BuildBackPropGraph(model_uri, config, backprop_model_file); + ASSERT_TRUE(status.IsOK()) << status<<" (is_fp32 = " << is_fp32 << ", stage = " << i << ").\n"; + + // Skip the re-load for mixed-precision case. This model contains grad op that has function body, + // which takes a const tensor input. Const cast for input in function body won't be saved in the output + // model so reload will run into error. + // For the purpose of testing mixed-precision, BuildBackPropGraph above will be sufficient to verify the + // partition logic and validate the graph. + if (is_fp32) { + std::shared_ptr model; + // Ensure the partitioned model load. + status = Model::Load(backprop_model_file, model, nullptr, DefaultLoggingManager().DefaultLogger()); + ASSERT_TRUE(status.IsOK()) << status<<" (is_fp32 = " << is_fp32 << ", stage = " << i << ").\n"; + } + } } }