saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-10 17:37:14 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add exception check in training_runner when worker runs into error, and misc check on nccl and mpi calls (#4380)

Browse source 
* error check

* fix build warning treated as error
This commit is contained in:
Xueyun Zhu 2020-07-22 14:32:19 -07:00 committed by GitHub
parent c2ec3b734b
commit e2acb165e9
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 112 additions and 74 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
orttraining/orttraining/core/framework/mpi_setup.h
View file

@ -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;

1
orttraining/orttraining/models/runner/pipeline.h
View file

@ -214,6 +214,7 @@ struct PipelineWorkerState {
std::vector<MLValue> feeds;
std::vector<std::string> fetch_names;
std::vector<MLValue> fetches;
std::exception_ptr execution_exception{nullptr};
};
struct PipelineWorkerPool {

100
orttraining/orttraining/models/runner/training_runner.cc
View file

@ -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<MLValue>();
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<MLValue>();
// 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_);

1
orttraining/orttraining/models/runner/training_runner.h
View file

@ -213,6 +213,7 @@ class TrainingRunner {
VectorString& fetch_names,
std::vector<MLValue>& 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);

16
orttraining/orttraining/training_ops/cuda/collective/nccl_common.cc
View file

@ -4,6 +4,8 @@
#include "nccl_common.h"
#include <mpi.h>
#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();
}

27
orttraining/orttraining/training_ops/cuda/communication/recv.cc
View file

@ -8,6 +8,8 @@
#include "core/profile/profile.h"
#include <mpi.h>
#include "orttraining/core/framework/mpi_setup.h"
namespace onnxruntime {
namespace cuda {
@ -67,17 +69,14 @@ Status Recv::ComputeInternal(OpKernelContext* ctx) const {
static_cast<int>(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.

28
orttraining/orttraining/training_ops/cuda/communication/send.cc
View file

@ -9,6 +9,8 @@
#include <limits>
#include <mpi.h>
#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<CommInfo_t*>(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<int>(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.