From b851973f227863acce798f1b0e597a8b3c06bc38 Mon Sep 17 00:00:00 2001 From: Weixing Zhang Date: Tue, 27 Oct 2020 11:57:46 -0700 Subject: [PATCH] pipeline_worker_pool_.JoinAll() should be called in pipeline code path (#5604) Co-authored-by: Weixing Zhang --- .../orttraining/models/runner/training_runner.cc | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/orttraining/orttraining/models/runner/training_runner.cc b/orttraining/orttraining/models/runner/training_runner.cc index 18b76ad118..f42776e37f 100644 --- a/orttraining/orttraining/models/runner/training_runner.cc +++ b/orttraining/orttraining/models/runner/training_runner.cc @@ -628,13 +628,6 @@ 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); - } - // TODO: move this to an operator in graph. onnxruntime::contrib::OrtEventPool::GetInstance().ResetAllEvents(); @@ -905,7 +898,9 @@ Status TrainingRunner::TrainingLoop(IDataLoader& training_data_loader, IDataLoad } } // end of one file/shard - pipeline_worker_pool_.JoinAll(); + if (params_.pipeline_parallel_size > 1) { + pipeline_worker_pool_.JoinAll(); + } if (step_ < params_.num_train_steps) { training_data_loader.MoveToNextDataSet(); }