// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "orttraining/models/runner/training_runner.h"
#include <algorithm>
#include <memory>
#include <sstream>
#include <thread>

#include "core/common/logging/logging.h"
#include "core/common/logging/sinks/clog_sink.h"
#include "core/framework/tensorprotoutils.h"
#include "core/platform/env.h"
#include "core/platform/path_lib.h"
#ifdef ENABLE_NVTX_PROFILE
#include "core/providers/cuda/nvtx_profile_context.h"
#endif
#include "core/session/environment.h"
#include "orttraining/core/framework/checkpointing.h"
#include "orttraining/core/framework/distributed_run_context.h"
#include "orttraining/core/graph/optimizer_graph_builder.h"
#include "orttraining/models/runner/training_util.h"
#include "orttraining/training_ops/cpu/controlflow/event_pool.h"
#if defined(USE_CUDA) && defined(ORT_USE_NCCL) && defined(USE_NCCL_P2P)
#include "orttraining/training_ops/cuda/communication/nccl_service.h"
#endif
#include "nlohmann/json.hpp"
#include "test/perftest/utils.h"

using json = nlohmann::json;

namespace onnxruntime {
namespace training {

static std::vector<FreeDimensionOverride> overrides = {};
static SessionOptions SESSION_OPTION = {
    ExecutionMode::ORT_SEQUENTIAL,     // execution_mode
    ExecutionOrder::PRIORITY_BASED,    // execution_order
    false,                             // enable_profiling
    ORT_TSTR(""),                      // optimized_model_filepath
    true,                              // enable_mem_pattern
    true,                              // enable_mem_reuse
    true,                              // enable_cpu_mem_arena
    ORT_TSTR("onnxruntime_profile_"),  // profile_file_prefix
    "",                                // session_logid
    -1,                                // session_log_severity_level
    0,                                 // session_log_verbosity_level
    5,                                 // max_num_graph_transformation_steps
    TransformerLevel::Level1,          // graph_optimization_level
    {},                                // intra_op_param
    {},                                // inter_op_param
    overrides,                         // free_dimension_overrides
    true,                              // use_per_session_threads
    true,                              // thread_pool_allow_spinning
    false,                             // use_deterministic_compute
    {},                                // config_options
    {},                                // initializers_to_share_map
    {},                                // external_initializers
};

TrainingRunner::TrainingRunner(Parameters params, const Environment& env)
    : TrainingRunner(params, env, SESSION_OPTION) {
}

TrainingRunner::TrainingRunner(Parameters params, const Environment& env, SessionOptions session_options)
    : step_(0),
      round_(0),
      weight_update_step_count_(0),
      training_data_set_index_(0),
      params_(params),
      session_options_(session_options),
      session_(session_options, env),
      input_allocator_(params.input_allocator ? params.input_allocator : TrainingUtil::GetCpuAllocator()) {
  ORT_ENFORCE(!params_.model_path.empty());
  if (!params.weights_to_train.empty())
    ORT_ENFORCE(params.weights_not_to_train.empty());
  ORT_ENFORCE(!params_.training_optimizer_name.empty());
  if (params.deepspeed_zero.stage != 0)
    ORT_ENFORCE(params.use_nccl,
                "DeepSpeed ZeRO partitioning is only supported with NCCL distributed training.");
  ORT_ENFORCE(params.num_train_steps % params.gradient_accumulation_steps == 0,
              "Number of training steps must be a multiple of number of gradient accumulation step.");
}

Status TrainingRunner::Initialize() {
  if (params_.pipeline_parallel_size > 1 && !params_.pipeline_stage_paths.empty()) {
    // Pipeline partition happens outside ORT. We just load the result of partitioning forward graph.
    // Backward graph will be generated using ORT's graph transformers.
    ORT_ENFORCE(static_cast<size_t>(MPIContext::GetInstance().GetWorldSize()) == params_.pipeline_stage_paths.size());
    ORT_RETURN_IF_ERROR(session_.Load(params_.pipeline_stage_paths[MPIContext::GetInstance().GetWorldRank()]));
  } else {
    ORT_RETURN_IF_ERROR(session_.Load(params_.model_path));
  }

  TrainingSession::TrainingConfiguration config{};
  config.model_with_loss_function_path = params_.model_with_loss_func_path;
  config.model_with_training_graph_path = params_.model_with_training_graph_path;

  config.weight_names_to_train = params_.weights_to_train;
  config.weight_names_to_not_train = params_.weights_not_to_train;
  config.immutable_weights = params_.immutable_weights;

  config.gradient_graph_config.use_memory_efficient_gradient = params_.use_memory_efficient_gradient;
  config.gradient_graph_config.set_gradients_as_graph_outputs = false;

  config.gradient_accumulation_steps = params_.gradient_accumulation_steps;

  config.distributed_config.world_rank = MPIContext::GetInstance().GetWorldRank();
  config.distributed_config.world_size = MPIContext::GetInstance().GetWorldSize();
  config.distributed_config.local_size = MPIContext::GetInstance().GetLocalSize();
  config.distributed_config.local_rank = MPIContext::GetInstance().GetLocalRank();
  config.distributed_config.data_parallel_size = params_.data_parallel_size;
  config.distributed_config.horizontal_parallel_size = params_.horizontal_parallel_size;
  config.distributed_config.pipeline_parallel_size = params_.pipeline_parallel_size;

  if (params_.use_mixed_precision) {
    TrainingSession::TrainingConfiguration::MixedPrecisionConfiguration mp{};
    mp.use_mixed_precision_initializers = params_.use_mixed_precision_initializer;
    if (params_.use_bfloat16) {
      mp.mixed_precision_type = MixedPrecisionDataType::BF16;
    }
    mp.layernorm_stash_as_fp32 = params_.layernorm_stash_as_fp32;
    config.mixed_precision_config = mp;
  }

  // configure the loss function if no pipeline is used or it's the last stage of pipeline
  auto pipeline_stage_id = GetPipelineStageId(MPIContext::GetInstance().GetWorldRank(),
                                              params_.horizontal_parallel_size,
                                              params_.data_parallel_size);
  if (params_.pipeline_parallel_size == 1 ||
      (pipeline_stage_id + 1) == params_.pipeline_parallel_size) {
    TrainingSession::TrainingConfiguration::LossFunctionConfiguration lf{};
    lf.loss_function_info = params_.loss_func_info;

    config.loss_function_config = lf;
  }

  // always configure the optimizer
  {
    TrainingSession::TrainingConfiguration::OptimizerConfiguration opt{};
    opt.name = params_.training_optimizer_name;
    opt.learning_rate_input_name = params_.lr_params.feed_name;
    opt.weight_attributes_generator = params_.optimizer_attributes;
    opt.weight_int_attributes_generator = params_.optimizer_int_attributes;
    opt.use_mixed_precision_moments = params_.use_mixed_precision_moments;
    opt.do_all_reduce_in_mixed_precision_type = params_.allreduce_in_mixed_precision_type;
    opt.use_nccl = params_.use_nccl;
    opt.deepspeed_zero = params_.deepspeed_zero;
    opt.adasum_reduction_type = params_.GetAdasumReductionType();
    opt.enable_grad_norm_clip = params_.enable_grad_norm_clip;
    config.optimizer_config = opt;
  }

  if (params_.EnableTensorboard()) {
    TrainingSession::TrainingConfiguration::TensorboardConfiguration tb{};
    tb.summary_name = params_.summary_name;
    tb.scalar_node_names = params_.scalar_names;
    tb.histogram_node_names = params_.histogram_names;
    tb.norm_node_names = params_.norm_names;
    tb.dump_convergence_metrics = params_.dump_convergence_metrics;

    config.tensorboard_config = tb;
  }

  if (params_.use_gist) {
    TrainingSession::TrainingConfiguration::GistConfiguration gist{};
    gist.op_type = params_.gist_config.op_type;
    gist.compr_type = params_.gist_config.compr_type;

    config.gist_config = gist;
  }

  // Prepare pipeline information to do configuration.
  if (params_.pipeline_parallel_size > 1) {
    TrainingSession::TrainingConfiguration::PipelineConfiguration pipe{};
    // If partition is done outside and the paths to partitioned model are provided,
    // the session already loads a pipeline stage.
    pipe.do_partition = params_.pipeline_stage_paths.empty() ? true : false;
    pipe.fetch_names = params_.fetch_names;
    pipe.cut_list = params_.pipeline_partition_cut_list;
    pipe.op_id_to_stage = params_.op_id_to_stage;
    pipe.partitioned_model_path = params_.pipeline_partitioned_model_path;
    // Do not assign value to config.pipeline_config if pipeline is not used.
    config.pipeline_config = pipe;
  }

  // always configure the graph transformer
  {
    TrainingGraphTransformerConfiguration gt_config{};
    gt_config.enable_gelu_approximation = params_.enable_gelu_approximation;
    gt_config.attn_dropout_recompute = params_.attn_dropout_recompute;
    gt_config.gelu_recompute = params_.gelu_recompute;
    gt_config.transformer_layer_recompute = params_.transformer_layer_recompute;
    gt_config.number_recompute_layers = params_.number_recompute_layers;

    config.graph_transformer_config = gt_config;
  }

  TrainingSession::TrainingConfigurationResult config_result{};

  ORT_RETURN_IF_ERROR(session_.ConfigureForTraining(config, config_result));

  if (config_result.mixed_precision_config_result.has_value()) {
    const std::string& loss_scale_input_name =
        config_result.mixed_precision_config_result.value().loss_scale_input_name;
    if (params_.loss_scale == 0.0f) {
      // use dynamic loss_scale
      loss_scaler_ = std::make_unique<LossScaler>(loss_scale_input_name, true, static_cast<float>(1 << 16));
    } else {
      // use static loss_scale
      loss_scaler_ = std::make_unique<LossScaler>(loss_scale_input_name, false, params_.loss_scale);
    }
  }

  opt_graph_outputs_ = config_result.opt_config_result.value().output_key_to_graph_output_name;

  // Retrieve pipeline information from configuration result.
  VectorString fetch_names;
  if (params_.pipeline_parallel_size > 1) {
    // Instead of inside constructor, we initialize pipeline_schedule_ and pipeline_worker_pool_ here
    // because they dependent on the result of TraningSession::ConfigureForTraining(...).
    pipeline_schedule_ = pipeline::PipelineScheduler(
        params_.gradient_accumulation_steps,
        params_.pipeline_parallel_size,
        DistributedRunContext::GetRanks(WorkerGroupType::PipelineParallel));
    pipeline_worker_pool_ = pipeline::PipelineWorkerPool(params_.pipeline_parallel_size);

    fetch_names = config_result.pipeline_config_result.value().fetch_names;

    // Set tensor names for event IDs and outputs of event ops.
    pipeline_context_.pipeline_tensor_names = config_result.pipeline_config_result.value().pipeline_tensor_names;

    // Create a local function to append non-empty name to fetch_names list.
    auto append_non_empty_name = [&](const std::string& name) {
      if (!name.empty()) {
        fetch_names.push_back(name);
      }
    };

    // Append first output of each event operator.
    pipeline_context_.pipeline_tensor_names.ForEachOutputName(append_non_empty_name);

    // Names of allowed inputs after pipeline partition.
    pipeline_context_.feed_names = config_result.pipeline_config_result.value().feed_names;
    // Names of allowed outputs after pipeline partition.
    // Configure dimension of this pipeline.
    pipeline_context_.fetch_names = fetch_names;

    pipeline_context_.pipeline_stage_id = config_result.pipeline_config_result.value().pipeline_stage_id;
    pipeline_context_.num_pipeline_micro_batches = params_.gradient_accumulation_steps;
  } else {
    fetch_names = params_.fetch_names;
    pipeline_context_.pipeline_stage_id = 0;
  }

  // Expose all optimizer outputs as graph outputs.
  for (const auto& it : opt_graph_outputs_) {
    fetch_names.push_back(it.second);
  }

  // Expose all optimizer outputs and pipeline outputs and as graph outputs.
  ORT_RETURN_IF_ERROR(session_.OverrideGraphOutputs(fetch_names));

  for (const auto& factory : params_.providers) {
    auto provider = factory.second->CreateProvider();
    ORT_ENFORCE(factory.first == provider->Type());
    ORT_RETURN_IF_ERROR(session_.RegisterExecutionProvider(std::move(provider)));
  }

  if (params_.use_profiler && !session_options_.enable_profiling) {
    // Profiling has not already been enabled, so override from command line options.
    session_.StartProfiling(session_options_.profile_file_prefix);
  }

  ORT_RETURN_IF_ERROR(session_.Initialize());

  // Checkpointing initialization
  // session_.Initialize() must be called prior to LoadCheckpoint()
  if (!params_.checkpoints_dir.empty()) {
    checkpoint_registry_ = std::make_unique<CheckpointRegistry>(
        params_.checkpoints_dir, params_.max_num_checkpoints);

    // Load checkpoint, if any
    PathString checkpoint_to_load_path = params_.checkpoint_to_load_path;
    if (!checkpoint_to_load_path.empty() ||
        checkpoint_registry_->TryGetLatestCheckpoint(checkpoint_to_load_path)) {
      ORT_RETURN_IF_ERROR(LoadCheckpoint(checkpoint_to_load_path));
    }
  }

  return Status::OK();
}

Status TrainingRunner::Run(IDataLoader* training_data_loader, IDataLoader* test_data_loader,
                           const MapStringToString& mapped_dimensions) {
  if (MPIContext::GetInstance().GetWorldRank() == 0 && !params_.model_actual_running_graph_path.empty()) {
    ORT_RETURN_IF_ERROR(session_.Save(params_.model_actual_running_graph_path, TrainingSession::SaveOption::NO_RELOAD));
  }

  // maybe in the future we can support an evaluation-only run
  if (!training_data_loader) {
    LOGS_DEFAULT(WARNING) << "training data loader not provided, nothing to do";
    return Status::OK();
  }

  ORT_RETURN_IF_ERROR(TrainingLoop(*training_data_loader, test_data_loader, mapped_dimensions));

  // after successful Run(), update counters
  ++round_;
  step_ = 0;

  return Status::OK();
}

// Prepare feeds for a call to one session run.
Status TrainingRunner::PrepareFeedNamesAndFeeds(const SessionMode mode,
                                                IDataLoader& training_data_loader,
                                                DataSet& training_data,
                                                LearningRateScheduler* lr_scheduler,
                                                const size_t batch_index,
                                                std::vector<std::string>& feed_names,
                                                std::vector<OrtValue>& feeds) {
  // Initialize outputs of this function.
  feed_names = std::vector<std::string>();
  feeds = std::vector<OrtValue>();

  auto allowed_feed_begin = pipeline_context_.feed_names.begin();
  auto allowed_feed_end = pipeline_context_.feed_names.end();

  // Pick up feeds from data loader
  {
    std::vector<std::string> data_feed_names = training_data_loader.DataSetTensorNames();
    std::vector<OrtValue> data_feeds = training_data.GetKthBatch(params_.batch_size, batch_index, input_allocator_);
    for (size_t i = 0; i < data_feed_names.size(); ++i) {
      const auto name = data_feed_names[i];
      if (params_.pipeline_parallel_size == 1 || std::find(allowed_feed_begin, allowed_feed_end, name) != allowed_feed_end) {
        feed_names.push_back(name);
        feeds.push_back(data_feeds[i]);
      }
    }
  }

  // Pick up feed from loss scaling.
  if (loss_scaler_) {
    const auto name = loss_scaler_->GetLossScaleInputName();
    if (params_.pipeline_parallel_size == 1 || std::find(allowed_feed_begin, allowed_feed_end, name) != allowed_feed_end) {
      feed_names.push_back(name);
      const float loss_scale = (mode == EvaluateStep) ? 1.0f : loss_scaler_->GetLossScale();
      OrtValue loss_scale_val;
      TrainingUtil::CreateCpuMLValue(std::array<int64_t, 1>{1}, std::vector<float>{loss_scale}, &loss_scale_val, input_allocator_);
      feeds.push_back(loss_scale_val);
    }
  }

  // Pick up feed from learning rate schedule.
  {
    const auto name = params_.lr_params.feed_name;
    if (params_.pipeline_parallel_size == 1 || std::find(allowed_feed_begin, allowed_feed_end, name) != allowed_feed_end) {
      feed_names.push_back(name);
      // learning rate is 0 if there is no learning-rate scheduler. Otherwise, learning rate is obtained from the scheduler.
      const float learning_rate = lr_scheduler ? lr_scheduler->GetLearningRate(step_ + 1) : 0.0f;
      OrtValue lr_val;
      TrainingUtil::CreateCpuMLValue(std::array<int64_t, 1>{1}, std::vector<float>{learning_rate}, &lr_val, input_allocator_);
      feeds.push_back(lr_val);
    }
  }

  // Define a helper function to append scalar inputs to feeds for event operators.
  auto append_event_to_feeds = [&](const std::string event_name, const int64_t event_value) -> void {
    if (event_name.empty()) {
      return;
    }
    feed_names.push_back(event_name);
    OrtValue event_id;
    TrainingUtil::CreateCpuMLScalar(
        event_value,
        &event_id,
        input_allocator_);
    feeds.push_back(event_id);
  };

  // Add event IDs to feeds.
  if (params_.pipeline_parallel_size > 1) {
    const auto batch_id = static_cast<int>(step_) % pipeline_context_.num_pipeline_micro_batches;
    const auto stage_id = pipeline_context_.pipeline_stage_id;

    int64_t id = -1;

    // Forward Recv
    id = (mode == EvaluateStep) ? -1 : pipeline_schedule_.GetForwardRecvWaitedEvent(batch_id, stage_id);
    append_event_to_feeds(pipeline_context_.pipeline_tensor_names.forward_recv_waited_event_name, id);
    id = (mode == EvaluateStep) ? -1 : pipeline_schedule_.GetForwardRecvRecordedEvent(batch_id, stage_id);
    append_event_to_feeds(pipeline_context_.pipeline_tensor_names.forward_recv_recorded_event_name, id);

    // Forward Send
    id = (mode == EvaluateStep) ? -1 : pipeline_schedule_.GetForwardSendWaitedEvent(batch_id, stage_id);
    append_event_to_feeds(pipeline_context_.pipeline_tensor_names.forward_send_waited_event_name, id);
    id = (mode == EvaluateStep) ? -1 : pipeline_schedule_.GetForwardSendRecordedEvent(batch_id, stage_id);
    append_event_to_feeds(pipeline_context_.pipeline_tensor_names.forward_send_recorded_event_name, id);

    // Backward Recv
    id = (mode == EvaluateStep) ? -1 : pipeline_schedule_.GetBackwardRecvWaitedEvent(batch_id, stage_id);
    append_event_to_feeds(pipeline_context_.pipeline_tensor_names.backward_recv_waited_event_name, id);
    id = (mode == EvaluateStep) ? -1 : pipeline_schedule_.GetBackwardRecvRecordedEvent(batch_id, stage_id);
    append_event_to_feeds(pipeline_context_.pipeline_tensor_names.backward_recv_recorded_event_name, id);

    // Backward Send
    id = (mode == EvaluateStep) ? -1 : pipeline_schedule_.GetBackwardSendWaitedEvent(batch_id, stage_id);
    append_event_to_feeds(pipeline_context_.pipeline_tensor_names.backward_send_waited_event_name, id);
    id = (mode == EvaluateStep) ? -1 : pipeline_schedule_.GetBackwardSendRecordedEvent(batch_id, stage_id);
    append_event_to_feeds(pipeline_context_.pipeline_tensor_names.backward_send_recorded_event_name, id);

    // Forward Compute
    id = (mode == EvaluateStep) ? -1 : pipeline_schedule_.GetForwardComputeWaitedEvent(batch_id, stage_id);
    append_event_to_feeds(pipeline_context_.pipeline_tensor_names.forward_compute_waited_event_name, id);
    id = (mode == EvaluateStep) ? -1 : pipeline_schedule_.GetForwardComputeRecordedEvent(batch_id, stage_id);
    append_event_to_feeds(pipeline_context_.pipeline_tensor_names.forward_compute_recorded_event_name, id);

    // Backward Compute
    id = (mode == EvaluateStep) ? -1 : pipeline_schedule_.GetBackwardComputeWaitedEvent(batch_id, stage_id);
    append_event_to_feeds(pipeline_context_.pipeline_tensor_names.backward_compute_waited_event_name, id);
    id = (mode == EvaluateStep) ? -1 : pipeline_schedule_.GetBackwardComputeRecordedEvent(batch_id, stage_id);
    append_event_to_feeds(pipeline_context_.pipeline_tensor_names.backward_compute_recorded_event_name, id);
  }

  for (auto name : feed_names) {
    ORT_ENFORCE(!name.empty(), "feed name cannot be empty string.");
  }
  return Status::OK();
}

Status TrainingRunner::PrepareFetchNamesAndFetches(const SessionMode mode,
                                                   std::vector<std::string>& fetch_names,
                                                   std::vector<OrtValue>& fetches) {
  // Initialize outputs of this function.
  fetch_names = std::vector<std::string>();
  fetches = std::vector<OrtValue>();

  const auto& allowed_fetch_names = pipeline_context_.fetch_names;

  if (mode == ModelUpdateStep) {
    // Set up tensor to be fetched when doing model update.

    if (params_.pipeline_parallel_size > 1) {
      // If pipeline is used, we need to filter out fetches which are not in this pipeline stage.

      for (size_t i = 0; i < params_.fetch_names.size(); ++i) {
        const auto name = params_.fetch_names[i];
        auto it = std::find(allowed_fetch_names.begin(), allowed_fetch_names.end(), name);
        if (it == allowed_fetch_names.end()) {
          continue;
        }
        fetch_names.push_back(name);
      }
    } 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) {
        if (!params_.use_bfloat16) {
          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_.enable_adasum) {
          auto 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) {
    // Set up tensor to be fetched when doing gradient accumulation.

    if (params_.gradient_accumulation_steps > 1) {
      auto it = opt_graph_outputs_.find(OptimizerOutputKey::GradientAccumulation);
      ORT_RETURN_IF(it == opt_graph_outputs_.end(), "Gradient accumulation output is missing in the optimizer output");
      fetch_names.push_back(it->second);
    }

    // Always execute event operators to avoid deadlock if pipeline is used.
    // TODO: create a list of must-to-fetch tensors and pass it to all graph transformer.
    if (params_.pipeline_parallel_size > 1) {
      // Create a local function to append non-empty name to fetch_names list.
      auto append_non_empty_name = [&](const std::string& name) {
        if (!name.empty()) {
          fetch_names.push_back(name);
        }
      };

      // Append first output of each event operator to fetch_names list to make sure all event ops will
      // be computed.
      pipeline_context_.pipeline_tensor_names.ForEachOutputName(append_non_empty_name);
    }
  } else if (mode == EvaluateStep) {
    // Set up tensor to be fetched when doing model evaluation.
    // Ideally, this path should not fetch optimizer and gradient accumulation.
    // This path may fetch predicted scores, loss value, and so on.

    if (params_.pipeline_parallel_size > 1) {
      // If pipeline is used, we need to filter out fetches which are not in this pipeline stage.

      for (size_t i = 0; i < params_.fetch_names.size(); ++i) {
        const auto name = params_.fetch_names[i];
        auto it = std::find(allowed_fetch_names.begin(), allowed_fetch_names.end(), name);
        if (it == allowed_fetch_names.end()) {
          continue;
        }
        fetch_names.push_back(name);
      }
    } else {
      // No pipeline. All fetched names should appear in the graph handled by this process.
      fetch_names = params_.fetch_names;
    }
  }

  // We need to fetch at least one variable.
  // If there is nothing to fetch, we fetch all model outputs.
  if (fetch_names.empty()) {
    fetch_names = allowed_fetch_names;
  }

  for (auto name : fetch_names) {
    ORT_ENFORCE(!name.empty(), "fetch name cannot be empty string.");
  }
  return Status::OK();
}

// If any exceptions happen during worker execution, it means there is an error
// during training. The worker thread propagates the exception to the main thread so
// we can properly cleanup and exit the execution.
void TrainingRunner::CheckWorkerException(const std::exception_ptr& p) {
  try {
    if (p) {
      std::rethrow_exception(p);
    }
  } catch (const std::exception& e) {
    ORT_THROW("Error in worker thread: ", e.what());
  }
}

// Launch synced session.Run on the main thread.
void TrainingRunner::RunWithUpdate(VectorString& feed_names,
                                   VectorString& fetch_names,
                                   std::vector<OrtValue>& feeds,
                                   std::vector<OrtValue>& fetches) {
  if (params_.pipeline_parallel_size > 1) {
    // Cyclically pick up a worker ID.
    const size_t worker_id = step_ % params_.pipeline_parallel_size;

    // Wait for the previous work to finish its job.
    // Its resource cannot be overrided when it's still working.
    pipeline_worker_pool_.Join(worker_id);
    CheckWorkerException(pipeline_worker_pool_.worker_states[worker_id].execution_exception);

    // Copy thread-used variable to thread-specific buffer to maintain their life.
    pipeline_worker_pool_.worker_states[worker_id].feed_names = feed_names;
    pipeline_worker_pool_.worker_states[worker_id].feeds = feeds;
    pipeline_worker_pool_.worker_states[worker_id].fetch_names = fetch_names;
    pipeline_worker_pool_.worker_states[worker_id].fetches = std::vector<OrtValue>();

    pipeline_worker_pool_.workers[worker_id] = std::thread([&](const size_t worker_id, const size_t step) {
      try {
#ifdef ENABLE_NVTX_PROFILE
        // Store the tag for the thread which runs session_.Run(...).
        // It will be used to name range in Nvidia's visual profiler.
        auto& profile_context = profile::Context::GetInstance();
        profile_context.SetThreadTag(
            std::this_thread::get_id(), std::to_string(step));
#else
        ORT_UNUSED_PARAMETER(step);
#endif
        RunOptions run_options;
        auto status = session_.Run(
            run_options,
            pipeline_worker_pool_.worker_states[worker_id].feed_names,
            pipeline_worker_pool_.worker_states[worker_id].feeds,
            pipeline_worker_pool_.worker_states[worker_id].fetch_names,
            &(pipeline_worker_pool_.worker_states[worker_id].fetches));

        ORT_THROW_IF_ERROR(status);
      } catch (std::exception&) {
        // If exception happens during worker execution, propogate the exception to main thread.
        pipeline_worker_pool_.worker_states[worker_id].execution_exception = std::current_exception();
      }
    },
                                                           worker_id, step_);

    // 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.
    pipeline_worker_pool_.JoinAll();
    for (auto& status : pipeline_worker_pool_.worker_states) {
      CheckWorkerException(status.execution_exception);
    }

    fetches = pipeline_worker_pool_.worker_states[worker_id].fetches;
  } else {
// Entering this branch means we will run graph without multi-threading.
// This branch is only hit in pipeline parallel.
#ifdef ENABLE_NVTX_PROFILE
    // Store the tag for the thread which runs session_.Run(...).
    // It will be used to name range in Nvidia's visual profiler.
    auto& profile_context = profile::Context::GetInstance();
    profile_context.SetThreadTag(
        std::this_thread::get_id(), std::to_string(step_));
#endif
    auto status = session_.Run(
        RunOptions(),
        feed_names,
        feeds,
        fetch_names,
        &fetches);
    ORT_ENFORCE(status.IsOK(), status.ErrorMessage());
  }

  if (loss_scaler_) {
    auto it = std::find(fetch_names.begin(), fetch_names.end(), opt_graph_outputs_[OptimizerOutputKey::GradientAllIsFinite]);
    if (it != fetch_names.end()) {
      const size_t index = static_cast<size_t>(std::distance(fetch_names.begin(), it));
      const Tensor& all_is_finite_t = fetches[index].Get<Tensor>();
      const bool is_all_finite = *(all_is_finite_t.template Data<bool>());
      loss_scaler_->UpdateLossScale(is_all_finite);
    }
  }

  // Assume that only the last pipeline stage can see loss, predicted value, and so on.
  // Thus, the error function should only be called when we are at the last stage.
  const bool session_can_see_loss = params_.pipeline_parallel_size == 1 ||
                                    pipeline_context_.pipeline_stage_id == params_.pipeline_parallel_size - 1;
  if (session_can_see_loss &&
      !params_.is_perf_test &&
      weight_update_step_count_ % params_.display_loss_steps == 0) {
    if (params_.error_function) {
      params_.error_function(feed_names, feeds, fetch_names, fetches, weight_update_step_count_);
    }
    if (params_.post_evaluation_callback) {
      params_.post_evaluation_callback(params_.batch_size, weight_update_step_count_, "train");
    }
  }

  // TODO: move this to an operator in graph.
  onnxruntime::contrib::OrtEventPool::GetInstance().ResetAllEvents();

  // Add one after process one batch.
  ++step_;
  // Add one after update the model once.
  ++weight_update_step_count_;
}

// Launch async session.Run on non-main thread.
void TrainingRunner::RunWithoutUpdate(VectorString& feed_names,
                                      VectorString& fetch_names,
                                      std::vector<OrtValue>& feeds,
                                      size_t& gradient_accumulation_step_count) {
  if (params_.pipeline_parallel_size > 1) {
    // Launch the graph using other threads when pipeline parallel is enabled.

    // Cyclically pick up a worker ID.
    const size_t worker_id = step_ % params_.pipeline_parallel_size;

    // Wait for the previous work to finish its job.
    // Its resource cannot be overrided when it's still working.
    pipeline_worker_pool_.Join(worker_id);
    CheckWorkerException(pipeline_worker_pool_.worker_states[worker_id].execution_exception);

    // Prepare async launch of session.
    // All used variables have to be copied to a buffer object to maintain their lifetime.
    pipeline_worker_pool_.worker_states[worker_id].feeds = feeds;
    pipeline_worker_pool_.worker_states[worker_id].feed_names = feed_names;
    pipeline_worker_pool_.worker_states[worker_id].fetch_names = fetch_names;
    pipeline_worker_pool_.worker_states[worker_id].fetches = std::vector<OrtValue>();

    // Async launch of a session.
    pipeline_worker_pool_.workers[worker_id] = std::thread([&](const size_t worker_id, const size_t step) {
      try {
#ifdef ENABLE_NVTX_PROFILE
        // Store the tag for the thread which runs session_.Run(...).
        // It will be used to name range in Nvidia's visual profiler.
        auto& profile_context = profile::Context::GetInstance();
        profile_context.SetThreadTag(
            std::this_thread::get_id(), std::to_string(step));
#else
        ORT_UNUSED_PARAMETER(step);
#endif
        RunOptions run_options;
        run_options.only_execute_path_to_fetches = true;
        run_options.training_mode = true;
        auto status = session_.Run(
            run_options,
            pipeline_worker_pool_.worker_states[worker_id].feed_names,
            pipeline_worker_pool_.worker_states[worker_id].feeds,
            pipeline_worker_pool_.worker_states[worker_id].fetch_names,
            &(pipeline_worker_pool_.worker_states[worker_id].fetches));
        ORT_THROW_IF_ERROR(status);
      } catch (std::exception&) {
        pipeline_worker_pool_.worker_states[worker_id].execution_exception = std::current_exception();
      }
    },
                                                           worker_id, step_);
  } else {
// Pipeline is not enabled, so we run session using the main thread.
#ifdef ENABLE_NVTX_PROFILE
    // Store the tag for the thread which runs session_.Run(...).
    // It will be used to name range in Nvidia's visual profiler.
    auto& profile_context = profile::Context::GetInstance();
    profile_context.SetThreadTag(
        std::this_thread::get_id(), std::to_string(step_));
#endif

    RunOptions run_options;
    run_options.only_execute_path_to_fetches = true;
    std::vector<OrtValue> fetches;
    auto status = session_.Run(
        run_options,
        feed_names,
        feeds,
        fetch_names,
        &fetches);
    ORT_ENFORCE(status.IsOK(), status.ErrorMessage());
  }

  // The following variables are only read and changed by the main thread, so
  // no thread-sync is needed.

  // Add one after process one batch.
  ++step_;
  // Add one after comuting one forward-backward path without applying optimizer.
  ++gradient_accumulation_step_count;
}

Status TrainingRunner::TrainingLoop(IDataLoader& training_data_loader, IDataLoader* test_data_loader,
                                    const MapStringToString& mapped_dimensions) {
  const bool enable_checkpoint_saving =
      MPIContext::GetInstance().GetWorldRank() == 0 &&
      checkpoint_registry_ && params_.checkpoint_period > 0;

  std::unique_ptr<perftest::utils::ICPUUsage> cpu_usage_calculator;
  if (!params_.perf_output_dir.empty()) {
    cpu_usage_calculator = perftest::utils::CreateICPUUsage();
  }

  if (test_data_loader) {
    ORT_RETURN_IF_ERROR(test_data_loader->InitializeDataSetIndex(0));
  }
  ORT_RETURN_IF_ERROR(training_data_loader.InitializeDataSetIndex(training_data_set_index_));

  const size_t num_shards_to_visit = training_data_loader.NumShards();
  const auto lr_scheduler = LearningRateScheduler::Create(params_.lr_params, params_.num_train_steps);

  double total_time{0};
  size_t epoch = 0;  // Note: epoch is not set properly when loaded from a checkpoint, but it's only for display.
  size_t gradient_accumulation_step_count = 0;
  const auto step_start = step_;
  const auto weight_update_step_count_start = weight_update_step_count_;

  // how many steps at last we used for stabilized perf benchmarking.
  const size_t stabilized_perf_total_step_count = std::min(static_cast<size_t>(128), params_.num_train_steps);
  const size_t stabilized_perf_start_step = params_.num_train_steps - stabilized_perf_total_step_count;
  double stabilized_total_time{0};
  constexpr size_t end_to_end_perf_start_step = 128;
  auto end_to_end_start = std::chrono::high_resolution_clock::now();
  bool end_to_end_measurement_started = false;

#if defined(USE_CUDA) && defined(ORT_USE_NCCL) && defined(USE_NCCL_P2P)
  // Create communication plan.
  auto& nccl_service = cuda::INcclService::GetInstance();

  nccl_service.PlanStart();
  for (auto& slot : pipeline_schedule_.GetSchedule(pipeline_context_.pipeline_stage_id)) {
    if (!slot.HasCommute()) {
      continue;
    }
    nccl_service.PlanNewGroupStart();
    for (auto& task : slot.GetTasks()) {
      if (task.type == pipeline::PipelineTask::Type::Send) {
        nccl_service.PlanSend(task.peer_rank);
      } else if (task.type == pipeline::PipelineTask::Type::Recv) {
        nccl_service.PlanRecv(task.peer_rank);
      }
    }
    nccl_service.PlanNewGroupEnd();
  }
  nccl_service.PlanEnd();

  // Launch NCCL service to execute the plan.
  nccl_service.Launch();
#endif

  auto all_steps_time_start = std::chrono::high_resolution_clock::now();
  while (step_ < params_.num_train_steps) {
    for (size_t shard_it = 0; shard_it < num_shards_to_visit; ++shard_it) {
      auto training_data = training_data_loader.CurrentDataSet();
      training_data_set_index_ = training_data_loader.CurrentDataSetIndex();
      if (training_data == nullptr) {
        printf("Skipping shard at index %d, which failed to load.\n",
               static_cast<int>(training_data_loader.CurrentDataSetIndex()));
        training_data_loader.MoveToNextDataSet();
        continue;
      }

      // Shuffle the data for each epoch
      if (params_.shuffle_data) {
        printf("Randomly shuffle training data.\n");
        training_data->RandomShuffle();
      }

      // loop through the data
      size_t batch_num_cur_shard = training_data->TotalBatch(params_.batch_size);
      for (size_t batch = 0; batch < batch_num_cur_shard && step_ < params_.num_train_steps; ++batch) {
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
        GetSession().GetMemoryProfiler().GetMemoryInfo().SetIteration(step_);
#endif
        const bool is_weight_update_step = (step_ + 1) % params_.gradient_accumulation_steps == 0;

        const bool stablized_perf_measurement_started = step_ >= stabilized_perf_start_step;
        if (!end_to_end_measurement_started && step_ >= end_to_end_perf_start_step) {
          end_to_end_start = std::chrono::high_resolution_clock::now();
          end_to_end_measurement_started = true;
        }

        VectorString feed_names;
        VectorString fetch_names;
        std::vector<OrtValue> feeds;
        std::vector<OrtValue> fetches;

        auto start = std::chrono::high_resolution_clock::now();

        if (is_weight_update_step) {
          ORT_RETURN_IF_ERROR(PrepareFeedNamesAndFeeds(ModelUpdateStep,
                                                       training_data_loader,
                                                       *training_data,
                                                       lr_scheduler.get(),
                                                       batch,
                                                       feed_names,
                                                       feeds));
          ORT_RETURN_IF_ERROR(
              PrepareFetchNamesAndFetches(ModelUpdateStep,
                                          fetch_names,
                                          fetches));
          RunWithUpdate(feed_names, fetch_names, feeds, fetches);
#if defined(USE_CUDA) && defined(ORT_USE_NCCL) && defined(USE_NCCL_P2P)
          nccl_service.Reset();
#endif
        } else {
          ORT_RETURN_IF_ERROR(PrepareFeedNamesAndFeeds(GradientAccumulateStep,
                                                       training_data_loader,
                                                       *training_data,
                                                       lr_scheduler.get(),
                                                       batch,
                                                       feed_names,
                                                       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.
        auto end = std::chrono::high_resolution_clock::now();
        std::chrono::duration<double> duration_seconds = end - start;
        total_time += duration_seconds.count();
        if (stablized_perf_measurement_started) {
          stabilized_total_time += duration_seconds.count();
        }

        printf("Stage %d, Round %d, Step: %d, epoch: %d, batch: %d/%d, shard_iteration: %d/%d, time: %.2f ms, throughput: %.2f ex/sec \n",
               pipeline_context_.pipeline_stage_id,
               static_cast<int>(round_),
               static_cast<int>(step_),
               static_cast<int>(epoch),
               static_cast<int>(batch),
               static_cast<int>(batch_num_cur_shard),
               static_cast<int>(shard_it + 1),
               static_cast<int>(num_shards_to_visit),
               duration_seconds.count() * 1000,
               params_.batch_size * (step_ - step_start) / total_time);
        printf("Training data range: [%d - %d)\n",
               static_cast<int>(batch * params_.batch_size),
               static_cast<int>((batch + 1) * params_.batch_size - 1));

        if (test_data_loader &&
            params_.do_eval && step_ % params_.evaluation_period == 0) {
          ORT_RETURN_IF_ERROR(Evaluate(session_, *test_data_loader));
        }

        if (enable_checkpoint_saving && is_weight_update_step &&
            weight_update_step_count_ % params_.checkpoint_period == 0) {
          PathString new_checkpoint_path, old_checkpoint_path;
          bool should_remove_old_checkpoint;

          ORT_RETURN_IF_ERROR(checkpoint_registry_->AddCheckpoint(
              weight_update_step_count_, new_checkpoint_path,
              should_remove_old_checkpoint, old_checkpoint_path));

          // ensure checkpoint directory exists
          if (!Env::Default().FolderExists(params_.checkpoints_dir)) {
            ORT_RETURN_IF_ERROR(Env::Default().CreateFolder(params_.checkpoints_dir));
          }

          if (should_remove_old_checkpoint) {
            const auto status = Env::Default().DeleteFolder(old_checkpoint_path);
            LOGS_DEFAULT_IF(!status.IsOK(), WARNING)
                << "Failed to delete old checkpoint. "
                << "Path: " << ToUTF8String(old_checkpoint_path)
                << ", error: " << status.ErrorMessage();
          }

          ORT_RETURN_IF_ERROR(SaveCheckpoint(new_checkpoint_path));
        }
      }  // end of one file/shard

      if (params_.pipeline_parallel_size > 1) {
        pipeline_worker_pool_.JoinAll();
      }
      if (step_ < params_.num_train_steps) {
        training_data_loader.MoveToNextDataSet();
      }
    }  // end of one epoch

    ++epoch;
  }
  auto all_steps_time_end = std::chrono::high_resolution_clock::now();
  std::chrono::duration<double> all_steps_duration_seconds = all_steps_time_end - all_steps_time_start;

  const double e2e_throughput = [&]() {
    if (end_to_end_perf_start_step >= params_.num_train_steps) return 0.0;
    auto end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> duration_seconds = end - end_to_end_start;
    const double total_e2e_time = duration_seconds.count();
    const size_t end_to_end_step_count = params_.num_train_steps - std::max(step_start, end_to_end_perf_start_step);
    return params_.batch_size * end_to_end_step_count / total_e2e_time;
  }();

  const size_t number_of_batches = step_ - step_start;
  const size_t weight_update_steps = weight_update_step_count_ - weight_update_step_count_start;
  const double avg_time_per_batch = total_time / (step_ - step_start) * 1000;
  const double throughput = params_.batch_size * (step_ - step_start) / total_time;
  const double stabilized_throughput = params_.batch_size / (stabilized_total_time / stabilized_perf_total_step_count);

  if (params_.perf_output_dir.empty()) {
    printf("No perf output directory specified, skipping save of trained perf metrics.\n");
  } else {
    const short average_cpu_usage = cpu_usage_calculator->GetUsage();
    const size_t peak_workingset_size = perftest::utils::GetPeakWorkingSetSize();
    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,
                                        e2e_throughput, mapped_dimensions,
                                        average_cpu_usage, peak_workingset_size));
  }

  std::cout << "Round: " << round_ << "\n"
            << "Batch size: " << params_.batch_size << "\n"
            << "Number of Batches: " << number_of_batches << "\n"
            << "Gradient Accumulation Steps: " << gradient_accumulation_step_count << "\n"
            << "Weight Update Steps: " << weight_update_steps << "\n"
            << "Total Running Time: " << total_time << " Seconds \n"
            << "Average Running Time Per Batch: " << avg_time_per_batch << " ms\n"
            << "Throughput: " << throughput << " Examples / Second\n"
            << "Stabilized Throughput: " << stabilized_throughput << " Examples / Second\n"
            << "EndToEnd Throughput: " << e2e_throughput << " Examples / Second\n"
            << "Average Step Time: " << all_steps_duration_seconds.count() / (step_ - step_start) << " Second\n"
            << "Average Step Throughput: " << params_.batch_size * (step_ - step_start) / (all_steps_duration_seconds.count()) << " Examples / Second\n";

#if defined(USE_CUDA) && defined(ORT_USE_NCCL) && defined(USE_NCCL_P2P)
  nccl_service.Terminate();
#endif
  return Status::OK();
}

Status TrainingRunner::SavePerfMetrics(const size_t number_of_batches, const size_t gradient_accumulation_steps,
                                       const size_t weight_update_steps, const double total_time,
                                       const double avg_time_per_batch, const double throughput, const double stabilized_throughput,
                                       const double e2e_throughput, const MapStringToString& mapped_dimensions,
                                       const short average_cpu_usage, const size_t peak_workingset_size) {
  // populate metrics for reporting
  json perf_metrics;
  perf_metrics["Model"] = params_.model_type;

  // loop thru the mapped_dimensions and put it in json sub-structure
  std::string seq_len;
  for (auto const& it : mapped_dimensions) {
    if (it.first == "SeqLen") {
      seq_len = it.second;
    }
    perf_metrics["DerivedProperties"][it.first] = it.second;
  }

  perf_metrics["Round"] = round_;
  perf_metrics["BatchSize"] = params_.batch_size;
  perf_metrics["NumOfBatches"] = number_of_batches;
  perf_metrics["GradAccSteps"] = gradient_accumulation_steps;
  perf_metrics["WeightUpdateSteps"] = weight_update_steps;
  perf_metrics["TotalTime"] = total_time;
  perf_metrics["AvgTimePerBatch"] = avg_time_per_batch;
  perf_metrics["Throughput"] = throughput;
  perf_metrics["StabilizedThroughput"] = stabilized_throughput;
  perf_metrics["EndToEndThroughput"] = e2e_throughput;
  perf_metrics["UseMixedPrecision"] = params_.use_mixed_precision;

  std::string optimizer = params_.training_optimizer_name;
  std::size_t pos = optimizer.find("Optimizer");
  if (pos != std::string::npos)
    optimizer = optimizer.substr(0, pos);
  perf_metrics["Optimizer"] = optimizer;

  Path model_path{};
  ORT_RETURN_IF_ERROR(Path::Parse(params_.model_path, model_path));
  PathString leaf = model_path.GetComponents().back();
  std::string model_name = ToUTF8String(leaf.c_str());
  perf_metrics["ModelName"] = model_name;

  std::string display_name = model_name + "_" + params_.model_type + "_" + (params_.use_mixed_precision ? "fp16" : "fp32") +
                             (seq_len.empty() ? "" : "_" + seq_len) + "_" + optimizer;
  perf_metrics["DisplayName"] = display_name;

  perf_metrics["Memory"] = peak_workingset_size >> 20;  // mb
  perf_metrics["AvgCPU"] = average_cpu_usage;

  //
  // we will get date/time and commitId in post-run pipeline
  //

  // populate other basic params for bookkeeping - add more as needed
  json bookkeeping_params;
  bookkeeping_params["LearningRate"] = params_.lr_params.initial_lr;
  bookkeeping_params["WarmupRatio"] = params_.lr_params.warmup_ratio;
  bookkeeping_params["WarmupMode"] = params_.lr_params.warmup_mode;
  bookkeeping_params["TrainSteps"] = params_.num_train_steps;
  bookkeeping_params["ModelPath"] = ToUTF8String(params_.model_path.c_str());
  bookkeeping_params["TrainDataDir"] = ToUTF8String(params_.train_data_dir.c_str());
  bookkeeping_params["TestDataDir"] = ToUTF8String(params_.test_data_dir.c_str());

  perf_metrics["RunConfig"] = bookkeeping_params.dump();  // serialize the params as json string

  std::string json_string = perf_metrics.dump();

  // write to a file - the next task in CI will pick up all files with the same prefix
  const PathString perf_metrics_path =
      params_.perf_output_dir + GetPathSep<PathChar>() + ORT_TSTR("onnxruntime_perf_metrics_") +
      ToPathString(display_name) + ORT_TSTR(".json");

  std::ofstream perf_metrics_stream;
  perf_metrics_stream.open(perf_metrics_path, std::ios::out | std::ios::trunc);
  ORT_RETURN_IF_NOT(perf_metrics_stream << json_string << "\n", "Failed to write to output file.");

  std::cout << "\n\nSaved perf metrics file: " << ToUTF8String(perf_metrics_path) << "\n\n";

  return Status::OK();
}

Status TrainingRunner::EndTraining(IDataLoader* data_loader) {
  if (params_.use_profiler) {
    // Write profiler data to disk.
    // We do this first in case there are any problems saving the trained model.
    std::string profile_file = session_.EndProfiling();
    std::cout << "Profiler data written to file " << profile_file << "\n";
  }

  if (data_loader) {
    // Test the in-memory model before saving.
    printf("\nEvaluating the final model on the test set.\n");
    ORT_RETURN_IF_ERROR(Evaluate(session_, *data_loader));
  }

  if (params_.output_dir.empty()) {
    printf("No output directory specified, skipping save of trained model.\n");
    return Status::OK();
  }

  // Create output directory if needed.
  if (!params_.output_dir.empty()) {
    ORT_RETURN_IF_ERROR(Env::Default().CreateFolder(params_.output_dir));
  }

  // We only want to save with WITH_UPDATED_WEIGHTS or WITH_UPDATED_WEIGHTS_AND_LOSS_FUNC
  // if the Loss function is in the graph, or otherwise saving will fail.
  // Assuming here that the Loss function will be assigned to the last device of the pipeline.
  if (MPIContext::GetInstance().GetWorldRank() == params_.pipeline_parallel_size - 1) {
    printf("\nSaving the trained model.\n");
    const PathString model_base_name = GetLastComponent(params_.model_path);

    const PathString trained_model_path =
        params_.output_dir + GetPathSep<PathChar>() + model_base_name + ORT_TSTR("_trained.onnx");
    ORT_RETURN_IF_ERROR(session_.Save(
        trained_model_path, TrainingSession::SaveOption::WITH_UPDATED_WEIGHTS));

    const PathString trained_model_with_loss_func_path =
        params_.output_dir + GetPathSep<PathChar>() + model_base_name + ORT_TSTR("_with_cost_trained.onnx");
    ORT_RETURN_IF_ERROR(session_.Save(
        trained_model_with_loss_func_path, TrainingSession::SaveOption::WITH_UPDATED_WEIGHTS_AND_LOSS_FUNC));
  }

  return Status::OK();
}

Status TrainingRunner::Evaluate(TrainingSession& session, IDataLoader& data_loader) {
  if (params_.skip_evaluation) {
    printf("Skipping evaluation...\n");
    return Status::OK();
  }

  // A static batch index representing current test batch
  static size_t current_batch = 0;
  auto test_data = data_loader.CurrentDataSet();
  if (params_.shuffle_data && current_batch == 0) {
    printf("Randomly shuffle test data.\n");
    test_data->RandomShuffle();
  }

  const size_t evaluation_batch_size = params_.eval_batch_size;

  printf("Test data range: [%d - %d)\n",
         static_cast<int>(current_batch * evaluation_batch_size),
         static_cast<int>((current_batch + 1) * evaluation_batch_size - 1));

  const size_t num_batches = size_t(ceil((float)evaluation_batch_size / (float)params_.batch_size));
  if (evaluation_batch_size % params_.batch_size != 0) {
    printf(
        "WARNING: evaluation_batch_size %zu is not an integer multiple of batch_size %zu. "
        "Using evaluation_batch_size %zu\n",
        evaluation_batch_size,
        params_.batch_size,
        num_batches * params_.batch_size);
  }

  RunOptions run_options;
  for (size_t batch_idx = 0; batch_idx < num_batches; ++batch_idx) {
    std::vector<std::string> feed_names;
    std::vector<OrtValue> feeds;
    std::vector<std::string> fetch_names;
    std::vector<OrtValue> fetches;

    ORT_RETURN_IF_ERROR(PrepareFeedNamesAndFeeds(EvaluateStep,
                                                 data_loader,
                                                 *test_data,
                                                 nullptr,
                                                 batch_idx,
                                                 feed_names,
                                                 feeds));
    if (!session.GetDropoutEvalFeeds().empty()) {
      float eval_ratio = 0.0f;
      for (auto& dropout_ratio : session.GetDropoutEvalFeeds()) {
        feed_names.push_back(dropout_ratio);
        OrtValue ratio_val;
        TrainingUtil::CreateCpuMLScalar(eval_ratio, &ratio_val, input_allocator_);
        feeds.push_back(ratio_val);
      }
    }
    const std::string training_mode_string = "training_mode";
    auto input_list = session.GetOverridableInitializers().second;
    for (auto input : *input_list) {
      if (input->Name().compare(training_mode_string) == 0) {
        feed_names.push_back("training_mode");
        OrtValue mode_val;
        TrainingUtil::CreateCpuMLScalar(false, &mode_val, input_allocator_);
        feeds.push_back(mode_val);
        break;
      }
    }

    ORT_RETURN_IF_ERROR(PrepareFetchNamesAndFetches(EvaluateStep,
                                                    fetch_names,
                                                    fetches));

    run_options.only_execute_path_to_fetches = true;
    run_options.training_mode = false;
    // Training threads are fully used by pipeline stages.
    // Pipeline cannot reuse training threads to do evaluation.
    // Otherwise, deadlock may happens.
    ORT_RETURN_IF_ERROR(session.Run(run_options,
                                    feed_names,
                                    feeds,
                                    fetch_names,
                                    &fetches));

    // Assume that user-specified fetches are avaliable only on the last pipeline stage.
    // When there is no pipeline, all pipeline_context_.pipeline_stage_id should be 0 and
    // params_.pipeline_parallel_size is 1. Thus, the following condition is always true if there
    // is no pipeline.
    const bool session_can_see_loss = pipeline_context_.pipeline_stage_id == params_.pipeline_parallel_size - 1;

    // Call error function
    if (session_can_see_loss && params_.error_function) {
      params_.error_function(feed_names, feeds, params_.fetch_names, fetches, step_);
    }

    // Set to next batch
    if (++current_batch >= test_data->TotalBatch(params_.batch_size)) {
      // Move to next shard
      test_data = data_loader.MoveToNextDataSet();
      current_batch = 0;
    }
  }

  // Call after a test batch.
  if (params_.post_evaluation_callback) {
    params_.post_evaluation_callback(evaluation_batch_size, step_, "test");
  }

  return Status::OK();
}

Status TrainingRunner::SaveCheckpoint(const PathString& checkpoint_path) {
  NameMLValMap checkpointed_tensors{};
  ORT_RETURN_IF_ERROR(session_.GetStateTensors(checkpointed_tensors));

  std::unordered_map<std::string, std::string> checkpointed_properties{};
  ORT_RETURN_IF_ERROR(SaveCheckpointProperties(checkpointed_properties));

  ORT_RETURN_IF_ERROR(SaveModelCheckpoint(
      checkpoint_path, session_.GetDataTransferManager(),
      checkpointed_tensors, checkpointed_properties));

  return Status::OK();
}

namespace {
Status WithOrtValuesFromTensorProtos(
    const PathString& model_location,
    const std::vector<ONNX_NAMESPACE::TensorProto>& tensor_protos,
    std::function<Status(const NameMLValMap&)> use_name_to_ort_value_fn) {
  static const OrtMemoryInfo cpu_alloc_info{onnxruntime::CPU, OrtDeviceAllocator};

  NameMLValMap name_to_ort_value{};
  std::vector<std::vector<char>> tensor_buffers{};

  for (const auto& tensor_proto : tensor_protos) {
    const auto* tensor_type = DataTypeImpl::TensorTypeFromONNXEnum(tensor_proto.data_type());
    const size_t element_size = tensor_type->GetElementType()->Size();
    const TensorShape shape{
        tensor_proto.dims().data(), static_cast<size_t>(tensor_proto.dims().size())};

    std::vector<char> tensor_buffer{};
    tensor_buffer.resize(element_size * shape.Size());

    const MemBuffer mem_buffer{tensor_buffer.data(), tensor_buffer.size(), cpu_alloc_info};

    OrtValue ort_value;

    ORT_RETURN_IF_ERROR(utils::TensorProtoToMLValue(
        Env::Default(), model_location.c_str(), tensor_proto, mem_buffer,
        ort_value));

    name_to_ort_value.emplace(tensor_proto.name(), ort_value);
    tensor_buffers.emplace_back(std::move(tensor_buffer));
  }

  ORT_RETURN_IF_ERROR(use_name_to_ort_value_fn(name_to_ort_value));

  return Status::OK();
}
}  // namespace

Status TrainingRunner::LoadCheckpoint(const PathString& checkpoint_path) {
  std::vector<ONNX_NAMESPACE::TensorProto> checkpointed_tensors{};
  std::unordered_map<std::string, std::string> checkpointed_properties{};
  ORT_RETURN_IF_ERROR(LoadModelCheckpoint(
      checkpoint_path, session_.GetModelLocation(),
      checkpointed_tensors, checkpointed_properties));

  ORT_RETURN_IF_ERROR(WithOrtValuesFromTensorProtos(
      session_.GetModelLocation(), checkpointed_tensors,
      [this](const NameMLValMap& name_to_ort_value) -> Status {
        ORT_RETURN_IF_ERROR(session_.SetStateTensors(name_to_ort_value, true));
        return Status::OK();
      }));

  ORT_RETURN_IF_ERROR(LoadCheckpointProperties(checkpointed_properties));

  return Status::OK();
}

namespace {
namespace property_names {
constexpr const char* k_step = "step";
constexpr const char* k_round = "round";
constexpr const char* k_weight_update_step = "weight_update_step";
constexpr const char* k_training_data_set_index = "training_data_set_index";
constexpr const char* k_loss_scaler_state = "loss_scaler_state";
}  // namespace property_names

template <typename T>
Status FromString(const std::string& s, T& t) {
  std::istringstream i{s};
  ORT_RETURN_IF_NOT(i >> t && i.eof(), "i >> t && i.eof() was false");
  return Status::OK();
}
}  // namespace

Status TrainingRunner::SaveCheckpointProperties(
    std::unordered_map<std::string, std::string>& properties) const {
  auto save_property = [&properties](const char* name, auto val) {
    properties[name] = std::to_string(val);
  };

  save_property(property_names::k_step, step_);
  save_property(property_names::k_round, round_);
  save_property(property_names::k_weight_update_step, weight_update_step_count_);
  save_property(property_names::k_training_data_set_index, training_data_set_index_);

  if (loss_scaler_) {
    properties[property_names::k_loss_scaler_state] = loss_scaler_->SaveToString();
  }

  return Status::OK();
}

Status TrainingRunner::LoadCheckpointProperties(
    const std::unordered_map<std::string, std::string>& properties) {
  auto load_property = [&properties](const char* name, auto& val) {
    auto prop_it = properties.find(name);
    ORT_RETURN_IF_NOT(prop_it != properties.end(), "prop_it == properties.end()");
    ORT_RETURN_IF_ERROR(FromString(prop_it->second, val));
    return Status::OK();
  };

  ORT_RETURN_IF_ERROR(load_property(property_names::k_step, step_));
  ORT_RETURN_IF_ERROR(load_property(property_names::k_round, round_));
  ORT_RETURN_IF_ERROR(load_property(
      property_names::k_weight_update_step, weight_update_step_count_));
  ORT_RETURN_IF_ERROR(load_property(
      property_names::k_training_data_set_index, training_data_set_index_));

  if (loss_scaler_) {
    auto prop_it = properties.find(property_names::k_loss_scaler_state);
    ORT_RETURN_IF_NOT(prop_it != properties.end(), "prop_it == properties.end()");
    ORT_RETURN_IF_ERROR(loss_scaler_->LoadFromString(prop_it->second));
  }

  return Status::OK();
}

Status TrainingRunner::UpdateParams(Parameters params) {
  params_.lr_params.initial_lr = params.lr_params.initial_lr;
  params_.lr_params.warmup_ratio = params.lr_params.warmup_ratio;
  params_.num_train_steps = params.num_train_steps;
  params_.batch_size = params.batch_size;
  params_.gradient_accumulation_steps = params.gradient_accumulation_steps;
  return Status::OK();
}

Status TrainingRunner::ResetLossScaler() {
  if (loss_scaler_) {
    loss_scaler_->Reset();
  }
  return Status::OK();
}
}  // namespace training
}  // namespace onnxruntime
