diff --git a/orttraining/orttraining/core/framework/mpi_setup.h b/orttraining/orttraining/core/framework/mpi_setup.h index 927f751fc1..1d9f6515d6 100644 --- a/orttraining/orttraining/core/framework/mpi_setup.h +++ b/orttraining/orttraining/core/framework/mpi_setup.h @@ -11,6 +11,19 @@ namespace onnxruntime { namespace training { +#define MPI_CHECK(condition) \ + do { \ + int error = (condition); \ + ORT_ENFORCE( \ + error == MPI_SUCCESS, \ + "MPI Error at: ", \ + __FILE__, \ + ":", \ + __LINE__, \ + ": ", \ + error); \ + } while (0) + struct MPIContext { MPIContext(int world_rank = 0, int local_rank = 0, int world_size = 1, int local_size = 1); int world_rank; diff --git a/orttraining/orttraining/models/runner/pipeline.h b/orttraining/orttraining/models/runner/pipeline.h index 9d40bc8d3f..a2dde6534a 100644 --- a/orttraining/orttraining/models/runner/pipeline.h +++ b/orttraining/orttraining/models/runner/pipeline.h @@ -214,6 +214,7 @@ struct PipelineWorkerState { std::vector feeds; std::vector fetch_names; std::vector fetches; + std::exception_ptr execution_exception{nullptr}; }; struct PipelineWorkerPool { diff --git a/orttraining/orttraining/models/runner/training_runner.cc b/orttraining/orttraining/models/runner/training_runner.cc index f97b72d0aa..55ab25a951 100644 --- a/orttraining/orttraining/models/runner/training_runner.cc +++ b/orttraining/orttraining/models/runner/training_runner.cc @@ -631,6 +631,19 @@ Status TrainingRunner::PrepareFetchNamesAndFetches(const SessionMode mode, 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, @@ -642,6 +655,7 @@ void TrainingRunner::RunWithUpdate(VectorString& feed_names, // 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; @@ -649,25 +663,30 @@ void TrainingRunner::RunWithUpdate(VectorString& feed_names, pipeline_worker_pool_.worker_states[worker_id].fetch_names = fetch_names; pipeline_worker_pool_.worker_states[worker_id].fetches = std::vector(); - Status status = Status::OK(); - pipeline_worker_pool_.workers[worker_id] = std::thread([&]( - const size_t worker_id, const size_t step) { + 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)); + // 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); + ORT_UNUSED_PARAMETER(step); #endif - RunOptions run_options; - 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)); + 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_); @@ -676,9 +695,9 @@ void TrainingRunner::RunWithUpdate(VectorString& feed_names, // We must join here because main thread needs to access thread-produced // fetches and those fetches must be ready. pipeline_worker_pool_.JoinAll(); - - // If the updating thread fails, we return with its error status. - ORT_THROW_IF_ERROR(status); + for(auto& status : pipeline_worker_pool_.worker_states){ + CheckWorkerException(status.execution_exception); + } // Copy back from thread-specific buffer to main thread's memory. fetches = pipeline_worker_pool_.worker_states[worker_id].fetches; @@ -711,6 +730,9 @@ void TrainingRunner::RunWithUpdate(VectorString& feed_names, // Wait all workers to finish this around of pipeline parallism. // The last batch in a pipeline collects gradient and update the model. pipeline_worker_pool_.JoinAll(); + for(auto& status : pipeline_worker_pool_.worker_states){ + CheckWorkerException(status.execution_exception); + } // Add one after process one batch. ++step_; @@ -729,6 +751,7 @@ void TrainingRunner::RunWithoutUpdate(VectorString& feed_names, // 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. @@ -738,27 +761,30 @@ void TrainingRunner::RunWithoutUpdate(VectorString& feed_names, pipeline_worker_pool_.worker_states[worker_id].fetches = std::vector(); // Async launch of a session. - pipeline_worker_pool_.workers[worker_id] = std::thread([&]( - const size_t worker_id, const size_t step) { + 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)); + // 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); + 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); + 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_); diff --git a/orttraining/orttraining/models/runner/training_runner.h b/orttraining/orttraining/models/runner/training_runner.h index 51ef943ba0..ebd25411dc 100644 --- a/orttraining/orttraining/models/runner/training_runner.h +++ b/orttraining/orttraining/models/runner/training_runner.h @@ -213,6 +213,7 @@ class TrainingRunner { VectorString& fetch_names, std::vector& feeds, size_t& gradient_accumulation_step_count); + void CheckWorkerException(const std::exception_ptr& p); Status TrainingLoop(IDataLoader& training_data_loader, IDataLoader* test_data_loader, const MapStringToString& mapped_dimensions); Status Evaluate(TrainingSession& session, IDataLoader& data_loader); diff --git a/orttraining/orttraining/training_ops/cuda/collective/nccl_common.cc b/orttraining/orttraining/training_ops/cuda/collective/nccl_common.cc index 4e5db008c8..c4e2a10f81 100644 --- a/orttraining/orttraining/training_ops/cuda/collective/nccl_common.cc +++ b/orttraining/orttraining/training_ops/cuda/collective/nccl_common.cc @@ -4,6 +4,8 @@ #include "nccl_common.h" #include +#include "orttraining/core/framework/mpi_setup.h" + namespace onnxruntime { namespace cuda { @@ -39,25 +41,25 @@ static Status CreateNcclCommunicator(MPI_Group* mpi_world_group, // Create new group MPI_Group mpi_group; - MPI_Group_incl(*mpi_world_group, worker_group.ranks.size(), worker_group.ranks.data(), &mpi_group); + MPI_CHECK(MPI_Group_incl(*mpi_world_group, worker_group.ranks.size(), worker_group.ranks.data(), &mpi_group)); // Create new MPI communicator MPI_Comm mpi_comm; static int32_t mpi_group_id = 0; - MPI_Comm_create_group(MPI_COMM_WORLD, mpi_group, ++mpi_group_id, &(mpi_comm)); + MPI_CHECK(MPI_Comm_create_group(MPI_COMM_WORLD, mpi_group, ++mpi_group_id, &(mpi_comm))); ORT_ENFORCE(mpi_comm != MPI_COMM_NULL, "MPI communicator creation failed."); // Create new NCCL communicator ncclUniqueId nccl_id; if (worker_group.rank_in_group == 0) { - ncclGetUniqueId(&nccl_id); + NCCL_RETURN_IF_ERROR(ncclGetUniqueId(&nccl_id)); } - MPI_Bcast(&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, mpi_comm); - ncclCommInitRank(group_comm, worker_group.ranks.size(), nccl_id, worker_group.rank_in_group); + MPI_CHECK(MPI_Bcast(&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, mpi_comm)); + NCCL_RETURN_IF_ERROR(ncclCommInitRank(group_comm, worker_group.ranks.size(), nccl_id, worker_group.rank_in_group)); // Clean up - MPI_Group_free(&mpi_group); - MPI_Comm_free(&mpi_comm); + MPI_CHECK(MPI_Group_free(&mpi_group)); + MPI_CHECK(MPI_Comm_free(&mpi_comm)); return Status::OK(); } diff --git a/orttraining/orttraining/training_ops/cuda/communication/recv.cc b/orttraining/orttraining/training_ops/cuda/communication/recv.cc index db48bbe4e5..9d43323b88 100644 --- a/orttraining/orttraining/training_ops/cuda/communication/recv.cc +++ b/orttraining/orttraining/training_ops/cuda/communication/recv.cc @@ -8,6 +8,8 @@ #include "core/profile/profile.h" #include +#include "orttraining/core/framework/mpi_setup.h" + namespace onnxruntime { namespace cuda { @@ -67,17 +69,14 @@ Status Recv::ComputeInternal(OpKernelContext* ctx) const { static_cast(tag_)}; - int mpi_code = 0; - // Directly use CPU to wait MPI_Recv. We cannot use GPU callback because // MPI_Recv may block the entire GPU until it returns. - mpi_code = MPI_Recv( + MPI_CHECK(MPI_Recv( info_shape_sizes.buffer, info_shape_sizes.size, MPI_CHAR, - info_shape_sizes.rank, info_shape_sizes.tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE); - ORT_ENFORCE(mpi_code == MPI_SUCCESS, "MPI Recv fails."); + info_shape_sizes.rank, info_shape_sizes.tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE)); #ifdef ENABLE_NVTX_PROFILE - // This range object includes the first MPI_Recv which receives a scalar. + // This range object includes the first MPI_Recv which receives a scalar. // It means we count the MPI's initialization in pre-recv stage. preRange.End(); #endif @@ -91,10 +90,9 @@ Status Recv::ComputeInternal(OpKernelContext* ctx) const { recvRange.Begin(); #endif - mpi_code = MPI_Recv( + MPI_CHECK(MPI_Recv( info_aggregated_size.buffer, info_aggregated_size.size, MPI_CHAR, - info_aggregated_size.rank, info_aggregated_size.tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE); - ORT_ENFORCE(mpi_code == MPI_SUCCESS, "MPI Recv fails."); + info_aggregated_size.rank, info_aggregated_size.tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE)); // Prepare receive shapes and data buffer aggregated_tensor_shapes.resize(prefix_tensor_shape_sizes[tensor_num - 1]); @@ -111,14 +109,13 @@ Status Recv::ComputeInternal(OpKernelContext* ctx) const { // Directly use CPU to wait MPI_Recv. We cannot use GPU callback because // MPI_Recv may block the entire GPU until it returns. - mpi_code = MPI_Recv( + MPI_CHECK(MPI_Recv( info_shapes.buffer, info_shapes.size, MPI_CHAR, - info_shapes.rank, info_shapes.tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE); - ORT_ENFORCE(mpi_code == MPI_SUCCESS, "MPI Recv fails."); - mpi_code = MPI_Recv( + info_shapes.rank, info_shapes.tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE)); + + MPI_CHECK(MPI_Recv( info_data.buffer, info_data.size, MPI_CHAR, - info_data.rank, info_data.tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE); - ORT_ENFORCE(mpi_code == MPI_SUCCESS, "MPI Recv fails."); + info_data.rank, info_data.tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE)); #ifdef ENABLE_NVTX_PROFILE // End of actual communication. diff --git a/orttraining/orttraining/training_ops/cuda/communication/send.cc b/orttraining/orttraining/training_ops/cuda/communication/send.cc index a3863e5d97..de03538aa8 100644 --- a/orttraining/orttraining/training_ops/cuda/communication/send.cc +++ b/orttraining/orttraining/training_ops/cuda/communication/send.cc @@ -9,6 +9,8 @@ #include #include +#include "orttraining/core/framework/mpi_setup.h" + namespace onnxruntime { namespace cuda { @@ -28,8 +30,7 @@ ONNX_OPERATOR_KERNEL_EX( void CUDART_CB HostSend(void* args) { CommInfo_t* info = reinterpret_cast(args); - int mpi_code = MPI_Send(info->buffer, info->size, MPI_CHAR, info->rank, info->tag, MPI_COMM_WORLD); - ORT_ENFORCE(mpi_code == MPI_SUCCESS, "MPI Send fails."); + MPI_CHECK(MPI_Send(info->buffer, info->size, MPI_CHAR, info->rank, info->tag, MPI_COMM_WORLD)); } Status Send::ComputeInternal(OpKernelContext* ctx) const { @@ -126,14 +127,12 @@ Status Send::ComputeInternal(OpKernelContext* ctx) const { dst, static_cast(tag_)}; - int mpi_code = 0; // Directly use CPU to wait MPI_Send. We cannot use GPU callback because // MPI_Send may block the entire GPU until it returns. - mpi_code = MPI_Send( + MPI_CHECK(MPI_Send( info_shape_sizes.buffer, info_shape_sizes.size, MPI_CHAR, - info_shape_sizes.rank, info_shape_sizes.tag, MPI_COMM_WORLD); - ORT_ENFORCE(mpi_code == MPI_SUCCESS, "MPI Send fails."); + info_shape_sizes.rank, info_shape_sizes.tag, MPI_COMM_WORLD)); #ifdef ENABLE_NVTX_PROFILE preRange.End(); @@ -148,18 +147,17 @@ Status Send::ComputeInternal(OpKernelContext* ctx) const { sendRange.Begin(); #endif - mpi_code = MPI_Send( + MPI_CHECK(MPI_Send( info_aggregated_size.buffer, info_aggregated_size.size, MPI_CHAR, - info_aggregated_size.rank, info_aggregated_size.tag, MPI_COMM_WORLD); - ORT_ENFORCE(mpi_code == MPI_SUCCESS, "MPI Send fails."); - mpi_code = MPI_Send( + info_aggregated_size.rank, info_aggregated_size.tag, MPI_COMM_WORLD)); + + MPI_CHECK(MPI_Send( info_shapes.buffer, info_shapes.size, MPI_CHAR, - info_shapes.rank, info_shapes.tag, MPI_COMM_WORLD); - ORT_ENFORCE(mpi_code == MPI_SUCCESS, "MPI Send fails."); - mpi_code = MPI_Send( + info_shapes.rank, info_shapes.tag, MPI_COMM_WORLD)); + + MPI_CHECK(MPI_Send( info_data.buffer, info_data.size, MPI_CHAR, - info_data.rank, info_data.tag, MPI_COMM_WORLD); - ORT_ENFORCE(mpi_code == MPI_SUCCESS, "MPI Send fails."); + info_data.rank, info_data.tag, MPI_COMM_WORLD)); #ifdef ENABLE_NVTX_PROFILE // End of major communication tasks.