Make gradient clipping configurable. (#3243)

* Make gradient clipping configurable.
add control flag to c++ and python frontend
This commit is contained in:
Tixxx 2020-03-23 12:21:48 -07:00 committed by GitHub
parent 5c52332223
commit 7f610caca0
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 28 additions and 7 deletions

View file

@ -45,6 +45,7 @@ struct OptimizerGraphConfig {
int64_t horovod_reduce_op{1};
std::string loss_scale_input_name{}; // empty string means no loss scaling factor is applied
AdasumReductionType adasum_reduction_type{AdasumReductionType::None};
bool enable_grad_norm_clip{true};
};
} // namespace training

View file

@ -190,7 +190,7 @@ Status OptimizerGraphBuilder::BuildOptimizerNode(
global_gradient_norm_argdef, global_gradient_norm_finite_argdef,
opt_configs, graph_defs,
new_initializers,
output_weight_argdefs, output_gradient_argdefs));
output_weight_argdefs, output_gradient_argdefs, opt_graph_config_.enable_grad_norm_clip));
return Status::OK();
}

View file

@ -76,6 +76,7 @@ Status SetupOptimizerParams(
opt_graph_config.allreduce_in_fp16 = optimizer_config.do_all_reduce_in_fp16;
opt_graph_config.use_nccl = optimizer_config.use_nccl;
opt_graph_config.adasum_reduction_type = optimizer_config.adasum_reduction_type;
opt_graph_config.enable_grad_norm_clip = optimizer_config.enable_grad_norm_clip;
#if USE_HOROVOD
opt_graph_config.horovod_reduce_op =
opt_graph_config.adasum_reduction_type == AdasumReductionType::None

View file

@ -129,6 +129,8 @@ class TrainingSession : public InferenceSession {
bool partition_optimizer{};
// Selects the reduction algorithm for Adasum.
AdasumReductionType adasum_reduction_type{AdasumReductionType::None};
// Whether to enable gradient clipping.
bool enable_grad_norm_clip{true};
};
// The optimizer configuration.
// If not provided, no optimizer is added.

View file

@ -140,7 +140,9 @@ Status ParseArguments(int argc, char* argv[], BertParameters& params, OrtParamet
("ratio_max", "Lamb max ratio parameter", cxxopts::value<float>()->default_value("5.0"))
("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"));
("horizontal_parallel_size", "Horizontal 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
.add_options("ORT configuration")
("ort_log_severity", "ORT minimum logging severity (see onnxruntime::logging::Severity values)",
@ -305,6 +307,7 @@ Status ParseArguments(int argc, char* argv[], BertParameters& params, OrtParamet
}
params.partition_optimizer = flags["partition_optimizer"].as<bool>();
params.enable_grad_norm_clip = flags["enable_grad_norm_clip"].as<bool>();
float alpha = flags["alpha"].as<float>();
float beta = flags["beta"].as<float>();
float lambda = flags["lambda"].as<float>();

View file

@ -81,7 +81,9 @@ Status ParseArguments(int argc, char* argv[], GPT2Parameters& params, OrtParamet
("lambda", "Adam/Lamb lambda parameter", cxxopts::value<float>()->default_value("0.01"))
("epsilon", "Adam/Lamb epsilon parameter", cxxopts::value<float>()->default_value("1e-8"))
("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"));
("horizontal_parallel_size", "Horizontal 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
.add_options("ORT configuration")
("ort_log_severity", "ORT minimum logging severity (see onnxruntime::logging::Severity values)",
@ -197,6 +199,7 @@ Status ParseArguments(int argc, char* argv[], GPT2Parameters& params, OrtParamet
}
params.partition_optimizer = flags["partition_optimizer"].as<bool>();
params.enable_grad_norm_clip = flags["enable_grad_norm_clip"].as<bool>();
float alpha = flags["alpha"].as<float>();
float beta = flags["beta"].as<float>();
float lambda = flags["lambda"].as<float>();

View file

@ -109,6 +109,7 @@ Status TrainingRunner::Initialize() {
opt.use_nccl = params_.use_nccl;
opt.partition_optimizer = params_.partition_optimizer;
opt.adasum_reduction_type = params_.GetAdasumReductionType();
opt.enable_grad_norm_clip = params_.enable_grad_norm_clip;
config.optimizer_config = opt;
}

View file

@ -149,6 +149,8 @@ class TrainingRunner {
int data_parallel_size = 1;
int horizontal_parallel_size = 1;
// Enable gradient clipping.
bool enable_grad_norm_clip=true;
};
TrainingRunner(Parameters params);

View file

@ -393,7 +393,7 @@ def create_ort_training_session_with_optimizer(model, device, training_optimizer
gradient_accumulation_steps=1, bind_parameters=False,
use_mixed_precision=False, allreduce_post_accumulation=False,
loss_scale_input_name='', scaled_loss_output_name='',
partition_optimizer=False):
partition_optimizer=False, enable_grad_norm_clip=True):
output_name = model.graph.output[0].name
ort_parameters = ort.TrainingParameters()
ort_parameters.loss_output_name = output_name
@ -410,6 +410,7 @@ def create_ort_training_session_with_optimizer(model, device, training_optimizer
ort_parameters.scaled_loss_output_name = scaled_loss_output_name
ort_parameters.allreduce_post_accumulation = allreduce_post_accumulation
ort_parameters.partition_optimizer = partition_optimizer
ort_parameters.enable_grad_norm_clip = enable_grad_norm_clip
output_types = {}
for output in model.graph.output:
@ -515,7 +516,7 @@ class ORTTrainer():
def __init__(self, model, loss_fn, model_desc, training_optimizer_name, map_optimizer_attributes,
learning_rate_description, device, gradient_accumulation_steps=1, postprocess_model=None,
world_rank=0, world_size=1, use_mixed_precision=False, allreduce_post_accumulation=False,
global_step=0, get_lr_this_step=None, loss_scaler=None, partition_optimizer=False):
global_step=0, get_lr_this_step=None, loss_scaler=None, partition_optimizer=False, enable_grad_norm_clip=True):
super(ORTTrainer, self).__init__()
"""
Initializes ORTTrainer.
@ -596,6 +597,7 @@ class ORTTrainer():
self.input_desc_with_lr_and_loss_scale = [*self.input_desc_with_lr, IODescription(self.loss_scale_input_name, [], torch.float32)]
else:
self.loss_scale_input_name, self.scaled_loss_output_name = '', ''
self.enable_grad_norm_clip_ = enable_grad_norm_clip
self.verify_fully_optimized_model(self.onnx_model_)
self.session, self.train_io_binding, self.eval_io_binding, self.output_name, _, self.output_types = \
@ -606,7 +608,7 @@ class ORTTrainer():
self.gradient_accumulation_steps, bind_parameters=False,
use_mixed_precision=use_mixed_precision, allreduce_post_accumulation=allreduce_post_accumulation,
loss_scale_input_name=self.loss_scale_input_name, scaled_loss_output_name=self.scaled_loss_output_name,
partition_optimizer=partition_optimizer)
partition_optimizer=partition_optimizer, enable_grad_norm_clip=self.enable_grad_norm_clip_)
# ORT backend has modified model output dtype from float32 to float16.
for o_desc in self.model_desc_.outputs_:

View file

@ -52,6 +52,7 @@ struct TrainingParameters {
int horizontal_parallel_size = 1;
bool partition_optimizer = false;
int seed = -1;
bool enable_grad_norm_clip = true;
};
// TODO: this method does not handle parallel optimization.
@ -133,6 +134,9 @@ void ConfigureSessionForTraining(
// an allreduce_post_accumulation option and remove the use_nccl option.
opt.use_nccl = parameters.allreduce_post_accumulation;
opt.partition_optimizer = parameters.partition_optimizer;
// TODO: The norm clipping value is 1.0f which is the default used in most frameworks.
// Need to have another option to support more values in the future.
opt.enable_grad_norm_clip = parameters.enable_grad_norm_clip;
config.optimizer_config = opt;
}
@ -166,7 +170,9 @@ void addObjectMethodsForTraining(py::module& m) {
.def_readwrite("world_rank", &TrainingParameters::world_rank)
.def_readwrite("world_size", &TrainingParameters::world_size)
.def_readwrite("gradient_accumulation_steps", &TrainingParameters::gradient_accumulation_steps)
.def_readwrite("partition_optimizer", &TrainingParameters::partition_optimizer);
.def_readwrite("partition_optimizer", &TrainingParameters::partition_optimizer)
.def_readwrite("enable_grad_norm_clip", &TrainingParameters::enable_grad_norm_clip);
py::class_<onnxruntime::training::TrainingSession, InferenceSession> training_session(m, "TrainingSession");
training_session.def(py::init<SessionOptions, SessionObjectInitializer>())