saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

pipeline_worker_pool_.JoinAll() should be called in pipeline code path (#5604)

Browse source 
Co-authored-by: Weixing Zhang <wezhan@microsoft.com>
This commit is contained in:
Weixing Zhang 2020-10-27 11:57:46 -07:00 committed by GitHub
parent 6f824c25e5
commit b851973f22
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -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();
}