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

add pipeline to distributed context config (#3789)

Browse source 
* add pipeline to distributed context

* white space
This commit is contained in:
Xueyun Zhu 2020-05-01 13:49:51 -07:00 committed by GitHub
parent 517bff9675
commit e8e95110d3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 28 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
orttraining/orttraining/core/framework/distributed_run_context.cc
View file

@ -7,8 +7,13 @@
namespace onnxruntime {
namespace training {
DistributedRunContext::DistributedRunContext(int32_t world_rank, int32_t world_size, int32_t local_rank,
int32_t local_size, int32_t data_parallel_size, int32_t horizontal_parallel_size) {
DistributedRunContext::DistributedRunContext(int32_t world_rank,
int32_t world_size,
int32_t local_rank,
int32_t local_size,
int32_t data_parallel_size,
int32_t horizontal_parallel_size,
int32_t pipeline_stage_size) {
// We only check world_size and world_rank since local_size and local_rank might not be set if using NCCL.
// TODO tix, refactor the mpi related code to populate all fields correctly by default.
ORT_ENFORCE(world_rank >= 0 && world_size > 0,
@ -22,10 +27,10 @@ DistributedRunContext::DistributedRunContext(int32_t world_rank, int32_t world_s
ORT_ENFORCE(world_size % horizontal_parallel_size == 0, "world size is not divisible by horizontal model parallel size.");
ORT_ENFORCE(world_size % data_parallel_size == 0, "world size is not divisible by data parallel size.");
// Be noted: this check and subsequent logic should be udpated when we introduce pipeline group
// Be noted: this check and subsequent logic should be updated when we introduce pipeline group
// depending how to split the pipeline groups.
ORT_ENFORCE(data_parallel_size * horizontal_parallel_size == world_size,
"total worker number != data_parallel_size * horizontal_parallel_size");
ORT_ENFORCE(data_parallel_size * horizontal_parallel_size * pipeline_stage_size == world_size,
"total worker number != data_parallel_size * horizontal_parallel_size * pipeline_stage_size");
params_.world_rank = world_rank;
params_.world_size = world_size;
@ -33,6 +38,7 @@ DistributedRunContext::DistributedRunContext(int32_t world_rank, int32_t world_s
params_.local_size = local_size;
params_.data_parallel_size = data_parallel_size;
params_.horizontal_parallel_size = horizontal_parallel_size;
params_.pipeline_stage_size = pipeline_stage_size;
groups_.resize(2);
// Initialize Data Parallel Group

11
orttraining/orttraining/core/framework/distributed_run_context.h
View file

@ -31,6 +31,7 @@ struct DistributedRunConfig {
int32_t local_size{1}; // Get local size of one physical node.
int32_t data_parallel_size{1};
int32_t horizontal_parallel_size{1};
int32_t pipeline_stage_size{1};
};
// Context managing global distribute run config, also responsible for splitting workers into groups
@ -45,7 +46,8 @@ class DistributedRunContext {
return DistributedRunContext::GetOrCreateInstance(config.world_rank, config.world_size,
config.local_rank, config.local_size,
config.data_parallel_size,
config.horizontal_parallel_size);
config.horizontal_parallel_size,
config.pipeline_stage_size);
}
static DistributedRunContext& GetInstance() {
@ -84,14 +86,15 @@ class DistributedRunContext {
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(DistributedRunContext);
DistributedRunContext(int32_t world_rank, int32_t world_size, int32_t local_rank, int32_t local_size,
int32_t data_parallel_size, int32_t horizontal_parallel_size);
int32_t data_parallel_size, int32_t horizontal_parallel_size, int32_t pipeline_stage_size = 1);
static DistributedRunContext& GetOrCreateInstance(int32_t world_rank = 0, int32_t world_size = 1,
int32_t local_rank = 0, int32_t local_size = 1,
int32_t data_parallel_size = 1,
int32_t horizontal_parallel_size = 1) {
int32_t horizontal_parallel_size = 1,
int32_t pipeline_stage_size = 1) {
static DistributedRunContext instance(world_rank, world_size, local_rank, local_size,
data_parallel_size, horizontal_parallel_size);
data_parallel_size, horizontal_parallel_size, pipeline_stage_size);
return instance;
}

3
orttraining/orttraining/core/session/training_session.cc
View file

@ -124,7 +124,8 @@ Status TrainingSession::ConfigureForTraining(
config.distributed_config.local_rank,
config.distributed_config.local_size,
config.distributed_config.data_parallel_size,
config.distributed_config.horizontal_parallel_size});
config.distributed_config.horizontal_parallel_size,
config.distributed_config.pipeline_stage_size});
ORT_RETURN_IF_ERROR(ApplyTransformationsToMainGraph());

3
orttraining/orttraining/core/session/training_session.h
View file

@ -61,6 +61,8 @@ class TrainingSession : public InferenceSession {
int data_parallel_size{1};
// The number of ranks for horizontal model parallel group
int horizontal_parallel_size{1};
// The number of stages for pipeline model parallel group
int pipeline_stage_size{1};
};
// The distributed training configuration.
DistributedConfiguration distributed_config{};
@ -225,7 +227,6 @@ class TrainingSession : public InferenceSession {
* @return The list of feed names.
*/
std::unordered_set<std::string> GetDropoutEvalFeeds() const { return dropout_eval_feeds_; }
/** Override Run function in InferenceSession to inject some training-specific logics **/
using InferenceSession::Run; // For overload resolution.
common::Status Run(const RunOptions& run_options, IOBinding& io_binding) override;

5
orttraining/orttraining/models/bert/main.cc
View file

@ -151,6 +151,7 @@ Status ParseArguments(int argc, char* argv[], BertParameters& params, OrtParamet
("cuda_mem_limit_in_gb", "Max cuda memory ort can use, in GB", cxxopts::value<float>()->default_value("-1.0"))
("data_parallel_size", "Data parallel group size.", cxxopts::value<int>()->default_value("1"))
("horizontal_parallel_size", "Horizontal model parallel group size.", cxxopts::value<int>()->default_value("1"))
("pipeline_stage_size", "pipeline model parallel group size.", cxxopts::value<int>()->default_value("1"))
("enable_grad_norm_clip", "Specify whether to enable gradient clipping for optimizers.",
cxxopts::value<bool>()->default_value("true"));
options
@ -358,8 +359,10 @@ Status ParseArguments(int argc, char* argv[], BertParameters& params, OrtParamet
params.data_parallel_size = flags["data_parallel_size"].as<int>();
params.horizontal_parallel_size = flags["horizontal_parallel_size"].as<int>();
params.pipeline_stage_size = flags["pipeline_stage_size"].as<int>();
ORT_RETURN_IF_NOT(params.data_parallel_size > 0, "data_parallel_size must > 0");
ORT_RETURN_IF_NOT(params.horizontal_parallel_size > 0, "horizontal_parallel_size must > 0");
ORT_RETURN_IF_NOT(params.pipeline_stage_size > 0, "pipeline_stage_size must > 0");
int64_t seed = flags["seed"].as<int64_t>();
if (params.horizontal_parallel_size > 1 && seed <= 0) {
@ -429,7 +432,7 @@ void setup_training_params(BertParameters& params) {
return;
}
auto data_group_size = params.mpi_context.world_size / params.horizontal_parallel_size;
auto data_group_size = params.mpi_context.world_size / (params.horizontal_parallel_size * params.pipeline_stage_size);
if (data_group_size != params.data_parallel_size) {
LOGS_DEFAULT(WARNING) << "WARNING: data_parallel_size is not correct, tuned automatically to "
<< data_group_size << std::endl;

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

@ -82,6 +82,7 @@ Status TrainingRunner::Initialize() {
config.distributed_config.local_rank = params_.mpi_context.local_rank;
config.distributed_config.data_parallel_size = params_.data_parallel_size;
config.distributed_config.horizontal_parallel_size = params_.horizontal_parallel_size;
config.distributed_config.pipeline_stage_size = params_.pipeline_stage_size;
if (params_.use_mixed_precision) {
TrainingSession::TrainingConfiguration::MixedPrecisionConfiguration mp{};

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

@ -151,6 +151,7 @@ class TrainingRunner {
int data_parallel_size = 1;
int horizontal_parallel_size = 1;
int pipeline_stage_size = 1;
// Enable gradient clipping.
bool enable_grad_norm_clip=true;
};