enable pipeline to run with mixed precision (#4113)

* enable pipeline to run with mixed precision

* address feedback

* address feedback

* test log

* pipe infomation if test fails

* ci failure

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.

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;
 

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.

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;i<ints_size;++i){
+        if(element_type->ints(i) == static_cast<int64_t>(ONNX_NAMESPACE::TensorProto_DataType_FLOAT)){
+          element_type->set_ints(i, static_cast<int64_t>(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));
 

orttraining/orttraining/core/graph/pipeline_transformer.cc

@@ -464,6 +464,7 @@ Status AddBackwardRecordBeforeSend(
 }
 
 Status SetInputsOutputsAndResolve(Graph& graph,
+                                  const std::unordered_set<std::string>& weights_to_train,
                                   const std::vector<std::string>& new_input_names,
                                   const std::vector<std::string>& 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<std::string>& 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<int64_t>(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<int64_t>(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();
 }
 

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<std::string>& initializer_names_to_preserve,
     std::string& forward_waited_event_name,
     std::string& forward_recorded_event_name,
     std::string& backward_waited_event_name,

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<std::string> loss_scale_input_name =
-        is_mixed_precision_enabled_ ? optional<std::string>{""} : optional<std::string>{};
+        enable_loss_scale ? optional<std::string>{""} : optional<std::string>{};
   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<std::string>& 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,

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<std::string>& 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,

orttraining/orttraining/models/bert/main.cc

@@ -504,13 +504,22 @@ const std::map<std::string, std::pair<std::string, size_t>> 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) {

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;

orttraining/orttraining/test/graph/gradient_graph_builder_test.cc

@@ -425,7 +425,7 @@ static void RunBertTrainingWithChecks(
 
   std::vector<OrtValue> 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<bool> 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> 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> 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";
+      }
+    }
   }
 }