mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Make gradient clipping configurable. (#3243)
* Make gradient clipping configurable. add control flag to c++ and python frontend
This commit is contained in:
parent
5c52332223
commit
7f610caca0
10 changed files with 28 additions and 7 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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>();
|
||||
|
|
|
|||
|
|
@ -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>();
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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_:
|
||||
|
|
|
|||
|
|
@ -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>())
|
||||
|
|
|
|||
Loading…
Reference in a new issue