diff --git a/orttraining/orttraining/core/framework/distributed_run_context.cc b/orttraining/orttraining/core/framework/distributed_run_context.cc index 4ac9304990..c7c1201bae 100644 --- a/orttraining/orttraining/core/framework/distributed_run_context.cc +++ b/orttraining/orttraining/core/framework/distributed_run_context.cc @@ -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 diff --git a/orttraining/orttraining/core/framework/distributed_run_context.h b/orttraining/orttraining/core/framework/distributed_run_context.h index 6734fbd307..6d3dc27cdd 100644 --- a/orttraining/orttraining/core/framework/distributed_run_context.h +++ b/orttraining/orttraining/core/framework/distributed_run_context.h @@ -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; } diff --git a/orttraining/orttraining/core/session/training_session.cc b/orttraining/orttraining/core/session/training_session.cc index 00cd81f95f..1c661b1ccd 100644 --- a/orttraining/orttraining/core/session/training_session.cc +++ b/orttraining/orttraining/core/session/training_session.cc @@ -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()); diff --git a/orttraining/orttraining/core/session/training_session.h b/orttraining/orttraining/core/session/training_session.h index 1731c696cd..3550a2389e 100644 --- a/orttraining/orttraining/core/session/training_session.h +++ b/orttraining/orttraining/core/session/training_session.h @@ -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 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; diff --git a/orttraining/orttraining/models/bert/main.cc b/orttraining/orttraining/models/bert/main.cc index fe2a7b9d14..faee11b2e2 100644 --- a/orttraining/orttraining/models/bert/main.cc +++ b/orttraining/orttraining/models/bert/main.cc @@ -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()->default_value("-1.0")) ("data_parallel_size", "Data parallel group size.", cxxopts::value()->default_value("1")) ("horizontal_parallel_size", "Horizontal model parallel group size.", cxxopts::value()->default_value("1")) + ("pipeline_stage_size", "pipeline model parallel group size.", cxxopts::value()->default_value("1")) ("enable_grad_norm_clip", "Specify whether to enable gradient clipping for optimizers.", cxxopts::value()->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(); params.horizontal_parallel_size = flags["horizontal_parallel_size"].as(); + params.pipeline_stage_size = flags["pipeline_stage_size"].as(); 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(); 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; diff --git a/orttraining/orttraining/models/runner/training_runner.cc b/orttraining/orttraining/models/runner/training_runner.cc index 54c20c6a46..a30b4e9d1f 100644 --- a/orttraining/orttraining/models/runner/training_runner.cc +++ b/orttraining/orttraining/models/runner/training_runner.cc @@ -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{}; diff --git a/orttraining/orttraining/models/runner/training_runner.h b/orttraining/orttraining/models/runner/training_runner.h index 94d972edca..97e0c3cc2f 100644 --- a/orttraining/orttraining/models/runner/training_runner.h +++ b/orttraining/orttraining/models/runner/training_runner.h @@ -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; };