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

Add max_norm for gradient clipping. (#6289)

Browse source 
* add max_norm as user option for gradient clipping

* add adam and lamb test cases for clip norm

* add frontend tests
This commit is contained in:
pengwa 2021-01-21 01:01:11 +08:00 committed by GitHub
parent a1b5bfc4f8
commit 453431f7bb
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
23 changed files with 652 additions and 236 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
orttraining/orttraining/core/graph/adasum_optimizer_graph_builder.cc
View file

@ -129,10 +129,10 @@ static Status AddNcclAllReduceForGradientsWithGroups(
allreduce_outputs[i] = ArgDef(gradient_argdefs[i].name + "_AllReduce_Out", allreduced_gradient_type_proto);
}
graph_defs.AddNodeDefs({NodeDef(OpDef{"View", kMSDomain, 1},
view_inputs,
allreduce_outputs,
NodeAttributes(),
"AllReduceOutputView")});
view_inputs,
allreduce_outputs,
NodeAttributes(),
"AllReduceOutputView")});
gradient_argdefs = allreduce_outputs;
return Status::OK();
@ -153,7 +153,7 @@ static Status AddAdasumAllReduceForGradients(
gradient_argdefs,
adasum_output_argdefs,
{ONNX_NAMESPACE::MakeAttribute("reduce_algo",
static_cast<int64_t>(adasum_reduction_type))},
static_cast<int64_t>(adasum_reduction_type))},
"AdasumAllReduce")});
gradient_argdefs = std::move(adasum_output_argdefs);
return Status::OK();
@ -168,7 +168,6 @@ Status AdasumOptimizerGraphBuilder::BuildInternal(
std::vector<ArgDef>& gradient_argdefs,
std::unordered_map<std::string, std::unordered_map<std::string, std::string>>& weight_to_opt_mapping,
OptimizerOutputKeyMap<std::string>& optimizer_graph_outputs) {
// Set weight update to false for optimizer
for (auto& opt_config : opt_configs_) {
opt_config.update_weight = false;
@ -191,7 +190,7 @@ Status AdasumOptimizerGraphBuilder::BuildInternal(
const float scale = 1.0f / scale_divisor;
// Only fuse if using hierarchical reduce.
const bool fuse_scaling_outputs = opt_graph_config_.adasum_reduction_type == AdasumReductionType::GpuHierarchicalReduction ? true: false;
const bool fuse_scaling_outputs = opt_graph_config_.adasum_reduction_type == AdasumReductionType::GpuHierarchicalReduction ? true : false;
ORT_RETURN_IF_ERROR(AddGradientScalingNodes(nodearg_name_generator, scale, gradient_argdefs, fused_gradient_argdef, graph_defs,
opt_graph_config_.AllReduceDataType(), fuse_scaling_outputs));
@ -200,7 +199,7 @@ Status AdasumOptimizerGraphBuilder::BuildInternal(
if (opt_graph_config_.adasum_reduction_type == AdasumReductionType::GpuHierarchicalReduction) {
#ifdef ORT_USE_NCCL
ORT_RETURN_IF_ERROR(AddNcclAllReduceForGradientsWithGroups(gradient_argdefs, fused_gradient_argdef, graph_defs,
reduced_fused_gradient_argdef, WorkerGroupType::NodeLocalDataParallel));
reduced_fused_gradient_argdef, WorkerGroupType::NodeLocalDataParallel));
#else
ORT_THROW("ORT is not built with NCCL.");
#endif

2
orttraining/orttraining/core/graph/optimizer/adam_optimizer_builder.h
View file

@ -15,6 +15,7 @@ class AdamOptimizerBuilder final : public OptimizerBuilder {
"beta",
"lambda",
"epsilon",
"max_norm_clip",
"do_bias_correction",
"weight_decay_mode"}) {}
@ -25,7 +26,6 @@ class AdamOptimizerBuilder final : public OptimizerBuilder {
std::unordered_map<std::string, std::unordered_map<std::string, std::string>>& weight_to_opt_mapping,
std::vector<ArgDef>& output_weight_argdefs,
std::vector<ArgDef>& output_gradient_argdefs) const override;
};
} // namespace training

20
orttraining/orttraining/core/graph/optimizer/lamb_optimizer_builder.cc
View file

@ -89,6 +89,7 @@ Status LambOptimizerBuilder::Build(
std::vector<float> beta;
std::vector<float> lambda;
std::vector<float> epsilon;
std::vector<float> max_norm_clip;
float ratio_min = -std::numeric_limits<float>::infinity();
float ratio_max = std::numeric_limits<float>::infinity();
int64_t do_bias_correction = 0;
@ -157,6 +158,12 @@ Status LambOptimizerBuilder::Build(
else
epsilon.emplace_back(1e-6f);
auto max_norm_clip_iter = attrs.find("max_norm_clip");
if (max_norm_clip_iter != attrs.end())
max_norm_clip.emplace_back(max_norm_clip_iter->second);
else
max_norm_clip.emplace_back(1.0f);
auto ratio_min_iter = attrs.find("ratio_min");
if (ratio_min_iter != attrs.end()) {
// All weight tensors should have the same min ratio.
@ -202,9 +209,7 @@ Status LambOptimizerBuilder::Build(
output_argdefs.push_back(output_gradient_argdef); // g_new
}
const auto element_type = opt_configs[i].use_mixed_precision_moments ?
ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT16 :
ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT;
const auto element_type = opt_configs[i].use_mixed_precision_moments ? ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT16 : ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT;
weight_to_opt_mapping[weight_name] = {};
// m1 & m2 & m1_new & m2_new
@ -245,11 +250,11 @@ Status LambOptimizerBuilder::Build(
// w_mixed_precision & w_mixed_precision_new
if (opt_configs[i].update_weight && opt_configs[i].mixed_precision_weight_arg != nullptr) {
input_argdefs.emplace_back(ArgDef(
opt_configs[i].mixed_precision_weight_arg->Name(),
opt_configs[i].mixed_precision_weight_arg->TypeAsProto()));
opt_configs[i].mixed_precision_weight_arg->Name(),
opt_configs[i].mixed_precision_weight_arg->TypeAsProto()));
output_weight_argdef = ArgDef(
opt_configs[i].mixed_precision_weight_arg->Name() + "_Lamb_out",
opt_configs[i].mixed_precision_weight_arg->TypeAsProto());
opt_configs[i].mixed_precision_weight_arg->Name() + "_Lamb_out",
opt_configs[i].mixed_precision_weight_arg->TypeAsProto());
output_argdefs.push_back(output_weight_argdef);
} else {
input_argdefs.emplace_back(ArgDef());
@ -266,6 +271,7 @@ Status LambOptimizerBuilder::Build(
attribute_protos.emplace_back(ONNX_NAMESPACE::MakeAttribute("beta", beta));
attribute_protos.emplace_back(ONNX_NAMESPACE::MakeAttribute("lambda", lambda));
attribute_protos.emplace_back(ONNX_NAMESPACE::MakeAttribute("epsilon", epsilon));
attribute_protos.emplace_back(ONNX_NAMESPACE::MakeAttribute("max_norm_clip", max_norm_clip));
attribute_protos.emplace_back(ONNX_NAMESPACE::MakeAttribute("ratio_min", ratio_min));
attribute_protos.emplace_back(ONNX_NAMESPACE::MakeAttribute("ratio_max", ratio_max));
attribute_protos.emplace_back(ONNX_NAMESPACE::MakeAttribute("do_bias_correction", do_bias_correction));

1
orttraining/orttraining/core/graph/optimizer/lamb_optimizer_builder.h
View file

@ -15,6 +15,7 @@ class LambOptimizerBuilder final : public OptimizerBuilder {
"beta",
"lambda",
"epsilon",
"max_norm_clip",
"ratio_min",
"ratio_max",
"do_bias_correction"}) {}

10
orttraining/orttraining/core/graph/optimizer_builder.h
View file

@ -49,10 +49,10 @@ Status IsMatchingTypeAndShape(
const int32_t element_type,
const std::vector<int64_t>& expected_shape);
/**
/**
* The configuration for optimizer builder.
*/
struct OptimizerBuilderConfig{
struct OptimizerBuilderConfig {
//The ArgDefs of the weights to optimize.
std::vector<ArgDef> weight_argdefs;
@ -70,11 +70,11 @@ struct OptimizerBuilderConfig{
// The per weight optimizer configuration.
std::vector<OptimizerNodeConfig> opt_configs;
// (Optional) The flag to force gradient clipping. If planning
// to use the default behavior of each sub-class, should not be set.
// (Optional) The flag to force gradient clipping. If planning
// to use the default behavior of each sub-class, should not be set.
optional<bool> enable_grad_clipping;
// The initial state for optimizer params
// The initial state for optimizer params
// shared by all weights.
NameMLValMap shared_optimizer_states{};
};

16
orttraining/orttraining/core/graph/optimizer_config.h
View file

@ -50,7 +50,7 @@ struct OptimizerNodeConfig {
std::unordered_map<std::string, float> attributes{};
std::unordered_map<std::string, int64_t> int_attributes{};
std::string loss_scale_input_name{};
NameMLValMap initial_states{}; // initial states for optimizer initializers
NameMLValMap initial_states{}; // initial states for optimizer initializers
bool use_mixed_precision_moments{false};
bool update_weight{true}; // indicates whether Optimizer should do weight update, or output new gradient
bool enabled{true}; // indicates whether this weight is included in the Optimizer
@ -71,7 +71,8 @@ struct OptimizerGraphConfig {
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};
NameMLValMap shared_optimizer_states{}; // initial states for shared params, eg. 'Step' for lamb
NameMLValMap shared_optimizer_states{}; // initial states for shared params, eg. 'Step' for lamb
ONNX_NAMESPACE::TensorProto_DataType AllReduceDataType() const {
if (!allreduce_in_mixed_precision_type) {
@ -79,10 +80,13 @@ struct OptimizerGraphConfig {
}
switch (mixed_precision_type) {
case MixedPrecisionDataType::FP16: return ONNX_NAMESPACE::TensorProto_DataType_FLOAT16;
case MixedPrecisionDataType::BF16: return ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16;
default: return ONNX_NAMESPACE::TensorProto_DataType_UNDEFINED;
}
case MixedPrecisionDataType::FP16:
return ONNX_NAMESPACE::TensorProto_DataType_FLOAT16;
case MixedPrecisionDataType::BF16:
return ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16;
default:
return ONNX_NAMESPACE::TensorProto_DataType_UNDEFINED;
}
}
};

35
orttraining/orttraining/core/graph/training_op_defs.cc
View file

@ -172,6 +172,11 @@ OpSchema& RegisterLambOpSchema(OpSchema&& op_schema) {
"Small scalar to avoid dividing by zero.",
AttributeProto::FLOATS,
std::vector<float>(1024, 1e-6f))
.Attr(
"max_norm_clip",
"clip threshold of gradients.",
AttributeProto::FLOATS,
std::vector<float>(1024, 1.f))
.Attr(
"do_bias_correction",
"Compute unbiased 1st and 2nd momentums.",
@ -663,6 +668,11 @@ void RegisterTrainingOpSchemas() {
"Small scalar to avoid dividing by zero.",
AttributeProto::FLOAT,
1e-8f)
.Attr(
"max_norm_clip",
"clip threshold of gradients.",
AttributeProto::FLOAT,
1.0f)
.Attr(
"do_bias_correction",
"Compute unbiased 1st and 2nd momentums.",
@ -926,9 +936,10 @@ Example 4:
ONNX_CONTRIB_OPERATOR_SCHEMA(NcclAllReduce)
.SetDomain(kMSDomain)
.SinceVersion(1)
.Attr("group_type", "0 - global parallel group, 1 - data parallel group, "
"2 - node local data parallel group, 3 - cross node data parallel group, "
"4 - horozontal parallel, 5 - model parallel.",
.Attr("group_type",
"0 - global parallel group, 1 - data parallel group, "
"2 - node local data parallel group, 3 - cross node data parallel group, "
"4 - horozontal parallel, 5 - model parallel.",
AttributeProto::INT,
static_cast<int64_t>(0))
.Input(0, "input", "tensors to be reduced", "T", OpSchema::Variadic)
@ -945,9 +956,10 @@ Example 4:
ONNX_CONTRIB_OPERATOR_SCHEMA(NcclAllGather)
.SetDomain(kMSDomain)
.SinceVersion(1)
.Attr("group_type", "0 - global parallel group, 1 - data parallel group, "
"2 - node local data parallel group, 3 - cross node data parallel group, "
"4 - horozontal parallel, 5 - model parallel.",
.Attr("group_type",
"0 - global parallel group, 1 - data parallel group, "
"2 - node local data parallel group, 3 - cross node data parallel group, "
"4 - horozontal parallel, 5 - model parallel.",
AttributeProto::INT,
static_cast<int64_t>(0))
.Input(0, "input", "tensors to be sent", "T", OpSchema::Variadic)
@ -964,9 +976,10 @@ Example 4:
ONNX_CONTRIB_OPERATOR_SCHEMA(NcclReduceScatter)
.SetDomain(kMSDomain)
.SinceVersion(1)
.Attr("group_type", "0 - global parallel group, 1 - data parallel group, "
"2 - node local data parallel group, 3 - cross node data parallel group, "
"4 - horozontal parallel, 5 - model parallel.",
.Attr("group_type",
"0 - global parallel group, 1 - data parallel group, "
"2 - node local data parallel group, 3 - cross node data parallel group, "
"4 - horozontal parallel, 5 - model parallel.",
AttributeProto::INT,
static_cast<int64_t>(0))
.Input(0, "input", "tensors to be reduced and scattered", "T", OpSchema::Variadic)
@ -980,7 +993,7 @@ Example 4:
assert(getAttribute(ctx, "group_type", 0) < static_cast<int64_t>(WorkerGroupType::WorkerGroupTypeCount));
})
#endif
;
;
ONNX_CONTRIB_OPERATOR_SCHEMA(AdasumAllReduce)
.SetDomain(kMSDomain)
@ -1002,7 +1015,7 @@ Example 4:
propagateElemTypeFromInputToOutput(ctx, i, i);
auto typeProto = ctx.getInputType(i);
if (!hasShape(*typeProto)) {
continue;
continue;
}
propagateShapeFromInputToOutput(ctx, i, i);
}

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

@ -381,7 +381,7 @@ Status TrainingSession::ConfigureForTraining(
std::string loss_name{};
if (config.pipeline_config.has_value()) {
// if use pipeline, first check if model contains send op. If it does, set the
// If use pipeline, first check if model contains send op. If it does, set the
// send node's output as the start tensor to build gradient graph
GetPipelineSendOutput(model_->MainGraph(), loss_name);
}
@ -425,14 +425,14 @@ Status TrainingSession::ConfigureForTraining(
ORT_RETURN_IF_ERROR(ApplyModelParallelTransformationsToMainGraph(trainable_initializers, config_result));
weight_partition_info_ = config_result.weight_partition_info;
// Save the model after graph transformations
if (IsRootNode(config) && config.model_after_graph_transforms_path.has_value()) {
ORT_IGNORE_RETURN_VALUE(Save(
config.model_after_graph_transforms_path.value(), SaveOption::NO_RELOAD));
}
// derive actual set of weights to train
// Derive actual set of weights to train
std::unordered_set<std::string> weight_names_to_train =
!filtered_config_weight_names_to_train.empty()
? trainable_initializers
@ -467,7 +467,7 @@ Status TrainingSession::ConfigureForTraining(
ORT_RETURN_IF_ERROR(BuildGradientGraph(
weight_names_to_train, loss_name, config.gradient_graph_config, *session_logger_));
if (IsRootNode(config) && config.model_with_gradient_graph_path.has_value()) {
ORT_IGNORE_RETURN_VALUE(Save(
config.model_with_gradient_graph_path.value(), SaveOption::NO_RELOAD));
@ -495,7 +495,7 @@ Status TrainingSession::ConfigureForTraining(
}
}
// add optimizer or gradient accumulation
// Add optimizer or gradient accumulation
if (config.optimizer_config.has_value()) {
OptimizerGraphConfig opt_graph_config{};
std::unordered_map<std::string, OptimizerNodeConfig> opt_node_configs{};
@ -516,7 +516,7 @@ Status TrainingSession::ConfigureForTraining(
// Set eval feed names for nodes that differ between training and inferencing.
ORT_RETURN_IF_ERROR(SetEvalFeedNames());
// add Tensorboard
// Add Tensorboard
if (config.tensorboard_config.has_value()) {
const auto& tensorboard_config = config.tensorboard_config.value();
@ -526,7 +526,7 @@ Status TrainingSession::ConfigureForTraining(
tensorboard_scalar_names.emplace_back(loss_scale_input_name.value());
}
// add some tensors from optimizer graph outputs
// Add some tensors from optimizer graph outputs
if (config_result.opt_config_result.has_value()) {
const auto& opt_output_key_to_graph_output_name =
config_result.opt_config_result.value().output_key_to_graph_output_name;
@ -549,7 +549,7 @@ Status TrainingSession::ConfigureForTraining(
tensorboard_config.dump_convergence_metrics));
}
// add GIST encoding
// Add GIST encoding
if (config.gist_config.has_value()) {
ORT_RETURN_IF_ERROR(AddGistEncoding());
}
@ -595,7 +595,7 @@ static Status AddLossScaling(
return Status::OK();
}
// add node to scale loss_name by loss_scale_input_name
// Add node to scale loss_name by loss_scale_input_name
GraphAugmenter::GraphDefs defs{};
*loss_scale_input_name = graph.GenerateNodeArgName("loss_scale");
const auto* loss_scale_input_type =
@ -621,7 +621,7 @@ static Status ConfigureLossFunctionInternal(
Graph& graph,
std::string* loss_scale_input_name,
std::string& actual_loss_name) {
// build loss function or use external one
// Build loss function or use external one
ORT_RETURN_IF_NOT(
(loss_func_info.has_value() && loss_graph_builder) ^ external_loss_name.has_value(),
"Either loss function information should be provided or an external "
@ -808,7 +808,6 @@ Status TrainingSession::ApplyModelParallelTransformationsToMainGraph(std::unorde
graph_transformation_mgr.Register(std::move(entry), TransformerLevel::Level1);
}
// apply transformers
Graph& graph = model_->MainGraph();
ORT_RETURN_IF_ERROR(graph_transformation_mgr.ApplyTransformers(
graph, TransformerLevel::Level1, *session_logger_));
@ -1688,10 +1687,10 @@ Status PipelineTrainingSession::BuildLossAndLossScaling(
loss_scale_input_name = enable_true_loss_scale ? optional<std::string>{""} : optional<std::string>{};
ORT_RETURN_IF_ERROR(BuildLoss(
external_loss_name,
loss_name,
loss_function_config,
loss_scale_input_name));
external_loss_name,
loss_name,
loss_function_config,
loss_scale_input_name));
if (enable_true_loss_scale) {
TrainingConfigurationResult::MixedPrecisionConfigurationResult mp_result{};

6
orttraining/orttraining/python/orttraining_pybind_state.cc
View file

@ -201,8 +201,6 @@ TrainingConfigurationResult ConfigureSessionForTraining(
// an allreduce_post_accumulation option and remove the use_nccl option.
opt.use_nccl = parameters.allreduce_post_accumulation;
opt.deepspeed_zero = onnxruntime::training::ZeROConfig(parameters.deepspeed_zero_stage);
// 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;
// TODO reduction types
@ -276,9 +274,9 @@ void CopyMPIContextToTrainingParameters(TrainingParameters& parameters, const lo
std::unordered_map<std::string, std::unordered_map<std::string, py::object>> ConvertORTTensorMapToNumpy(std::unordered_map<std::string, NameMLValMap> c_tensor_state, const DataTransferManager& data_transfer_manager) {
std::unordered_map<std::string, std::unordered_map<std::string, py::object>> py_tensor_state;
for (const auto& layer1_item: c_tensor_state) {
for (const auto& layer1_item : c_tensor_state) {
py_tensor_state[layer1_item.first] = {};
for (const auto& layer2_item: layer1_item.second) {
for (const auto& layer2_item : layer1_item.second) {
assert(layer2_item.second.IsTensor());
py::object obj;
const Tensor& rtensor = layer2_item.second.Get<Tensor>();

10
orttraining/orttraining/python/training/optim/config.py
View file

@ -148,13 +148,14 @@ class AdamConfig(_OptimizerConfig):
BEFORE_WEIGHT_UPDATE = 0,
AFTER_WEIGHT_UPDATE = 1
def __init__(self, params=[], lr=0.001, alpha=0.9, beta=0.999, lambda_coef=0.0, epsilon=1e-8,
def __init__(self, params=[], lr=0.001, alpha=0.9, beta=0.999, lambda_coef=0.0, epsilon=1e-8, max_norm_clip=1.0,
do_bias_correction=True, weight_decay_mode=DecayMode.BEFORE_WEIGHT_UPDATE):
assert lr >= 0, "'lr' must be a positive number"
assert alpha >= 0, "'alpha' must be a positive number"
assert beta >= 0, "'beta' must be a positive number"
assert lambda_coef >= 0, "'lambda_coef' must be a positive number"
assert epsilon >= 0, "'epsilon' must be a positive number"
assert max_norm_clip != 0, "'max_norm_clip' must not be 0"
assert isinstance(do_bias_correction, bool), "'do_bias_correction' must be a boolean"
assert isinstance(weight_decay_mode, AdamConfig.DecayMode), "'weight_decay_mode' must be a AdamConfig.DecayMode"
for param in params:
@ -165,6 +166,7 @@ class AdamConfig(_OptimizerConfig):
'beta': beta,
'lambda': lambda_coef,
'epsilon': epsilon,
'max_norm_clip': max_norm_clip,
'do_bias_correction': do_bias_correction,
'weight_decay_mode': weight_decay_mode}
super().__init__(name='AdamOptimizer',
@ -174,6 +176,7 @@ class AdamConfig(_OptimizerConfig):
self.beta = beta
self.lambda_coef = lambda_coef
self.epsilon = epsilon
self.max_norm_clip = max_norm_clip
self.do_bias_correction = do_bias_correction
self.weight_decay_mode = weight_decay_mode
@ -211,7 +214,7 @@ class LambConfig(_OptimizerConfig):
"""
def __init__(self, params=[], lr=0.001, alpha=0.9, beta=0.999, lambda_coef=0.0,
ratio_min=float('-inf'), ratio_max=float('inf'), epsilon=1e-6, do_bias_correction=False):
ratio_min=float('-inf'), ratio_max=float('inf'), epsilon=1e-6, max_norm_clip=1.0, do_bias_correction=False):
assert lr >= 0, "'lr' must be a positive number"
assert alpha >= 0, "'alpha' must be a positive number"
assert beta >= 0, "'beta' must be a positive number"
@ -219,6 +222,7 @@ class LambConfig(_OptimizerConfig):
assert isinstance(ratio_min, float), "'ratio_min' must be a valid float"
assert isinstance(ratio_max, float), "'ratio_max' must be a valid float"
assert epsilon >= 0, "'epsilon' must be a positive number"
assert max_norm_clip != 0, "'max_norm_clip' must not be 0"
assert isinstance(do_bias_correction, bool), "'do_bias_correction' must be a boolean"
for param in params:
assert 'lr' not in param, "'lr' is not supported inside params"
@ -230,6 +234,7 @@ class LambConfig(_OptimizerConfig):
'ratio_min': ratio_min,
'ratio_max': ratio_max,
'epsilon': epsilon,
'max_norm_clip': max_norm_clip,
'do_bias_correction': do_bias_correction}
super().__init__(name='LambOptimizer',
params=params,
@ -240,4 +245,5 @@ class LambConfig(_OptimizerConfig):
self.ratio_min = ratio_min
self.ratio_max = ratio_max
self.epsilon = epsilon
self.max_norm_clip = max_norm_clip
self.do_bias_correction = do_bias_correction

466
orttraining/orttraining/test/gradient/optimizer_ops_test.cc
View file

@ -9,6 +9,8 @@
namespace onnxruntime {
namespace test {
namespace {
TEST(OptimizerTest, SGDOptimizerTest) {
OpTester test("SGDOptimizer", 1, onnxruntime::kMSDomain);
test.AddInput<float>("ETA", {}, {0.5f});
@ -131,7 +133,7 @@ TEST(OptimizerTest, AdamBiasCorrection) {
test.AddInput<float>("ETA", {}, {1.f});
test.AddInput<int64_t>("Update_Count", {}, {1});
test.AddInput<float>("W", {3}, {-0.4634f, 0.3584f, -0.2121f});
test.AddInput<float>("W", {3}, {-0.4634f, 0.3584f, -0.2121f});
test.AddInput<float>("G", {3}, {0.4171f, 0.9485f, 1.2289f});
test.AddInput<float>("Moment_1", {3}, {0.f, 0.f, 0.f});
test.AddInput<float>("Moment_2", {3}, {0.f, 0.f, 0.f});
@ -153,7 +155,7 @@ TEST(OptimizerTest, AdamWeightDecayMode0NoBiasCorrection) {
test.AddInput<float>("ETA", {}, {1.f});
test.AddInput<int64_t>("Update_Count", {}, {1});
test.AddInput<float>("W", {3}, {-0.4634f, 0.3584f, -0.2121f});
test.AddInput<float>("W", {3}, {-0.4634f, 0.3584f, -0.2121f});
test.AddInput<float>("G", {3}, {0.4171f, 0.9485f, 1.2289f});
test.AddInput<float>("Moment_1", {3}, {0.f, 0.f, 0.f});
test.AddInput<float>("Moment_2", {3}, {0.f, 0.f, 0.f});
@ -164,7 +166,6 @@ TEST(OptimizerTest, AdamWeightDecayMode0NoBiasCorrection) {
test.AddOutput<float>("W_Out", {3}, {-3.6210f, -2.8075f, -3.3723f});
test.AddOutput<float>("G_Out", {3}, {-3.1576f, -3.1658f, -3.1601f});
test.AddAttribute("do_bias_correction", static_cast<int64_t>(0));
test.AddAttribute("lambda", 0.01f);
test.AddAttribute("weight_decay_mode", static_cast<int64_t>(0));
@ -178,7 +179,7 @@ TEST(OptimizerTest, AdamWeightDecayMode0WithBiasCorrection) {
test.AddInput<float>("ETA", {}, {1.f});
test.AddInput<int64_t>("Update_Count", {}, {1});
test.AddInput<float>("W", {3}, {-0.4634f, 0.3584f, -0.2121f});
test.AddInput<float>("W", {3}, {-0.4634f, 0.3584f, -0.2121f});
test.AddInput<float>("G", {3}, {0.4171f, 0.9485f, 1.2289f});
test.AddInput<float>("Moment_1", {3}, {0.f, 0.f, 0.f});
test.AddInput<float>("Moment_2", {3}, {0.f, 0.f, 0.f});
@ -202,7 +203,7 @@ TEST(OptimizerTest, AdamWeightDecayMode1NoBiasCorrection) {
test.AddInput<float>("ETA", {}, {1.f});
test.AddInput<int64_t>("Update_Count", {}, {1});
test.AddInput<float>("W", {3}, {-0.4634f, 0.3584f, -0.2121f});
test.AddInput<float>("W", {3}, {-0.4634f, 0.3584f, -0.2121f});
test.AddInput<float>("G", {3}, {0.4171f, 0.9485f, 1.2289f});
test.AddInput<float>("Moment_1", {3}, {0.f, 0.f, 0.f});
test.AddInput<float>("Moment_2", {3}, {0.f, 0.f, 0.f});
@ -225,7 +226,7 @@ TEST(OptimizerTest, AdamWeightDecayMode1WithBiasCorrection) {
test.AddInput<float>("ETA", {}, {1.f});
test.AddInput<int64_t>("Update_Count", {}, {1});
test.AddInput<float>("W", {3}, {-0.4634f, 0.3584f, -0.2121f});
test.AddInput<float>("W", {3}, {-0.4634f, 0.3584f, -0.2121f});
test.AddInput<float>("G", {3}, {0.4171f, 0.9485f, 1.2289f});
test.AddInput<float>("Moment_1", {3}, {0.f, 0.f, 0.f});
test.AddInput<float>("Moment_2", {3}, {0.f, 0.f, 0.f});
@ -243,6 +244,16 @@ TEST(OptimizerTest, AdamWeightDecayMode1WithBiasCorrection) {
}
#if defined(USE_CUDA) || defined(USE_ROCM)
float GetGradientL2Norm(const std::vector<float>& gradient_vector) {
float gradient_norm = 0.0f;
for (const auto g_value : gradient_vector) {
gradient_norm += g_value * g_value;
}
gradient_norm = std::sqrt(gradient_norm);
return gradient_norm;
}
TEST(OptimizerTest, AdamOptimizerMixPrecisionTest) {
OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain);
AdamOptimizerInputOutput data;
@ -292,6 +303,76 @@ TEST(OptimizerTest, AdamOptimizerMixPrecision_FP16Weight_Test) {
test.Run();
}
TEST(OptimizerTest, AdamOptimizerMixPrecision_FP16Weight_NoClipNorm_Test) {
OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain);
AdamOptimizerInputOutput data;
test.AddInput<MLFloat16>("ETA", {}, data.eta_half);
test.AddInput<int64_t>("Update_Count", {}, {3});
test.AddInput<float>("W", {3}, data.w);
test.AddInput<MLFloat16>("G", {3}, data.g_half);
test.AddInput<MLFloat16>("Moment_1", {3}, data.m1_half);
test.AddInput<MLFloat16>("Moment_2", {3}, data.m2_half);
test.AddInput<MLFloat16>("FP16_W", {3}, data.w_half);
test.AddInput<float>("loss_scale", {1}, {1.0f});
// grad clipping should not take effect because default max_norm is 1.0f
test.AddInput<float>("grad_norm", {1}, {0.01f});
// Verify AdamOptimizer outputs
test.AddOutput<int64_t>("Update_Count_Out", {}, {4});
test.AddOutput<MLFloat16>("Moment_1_Out", {3}, data.m1_new_half);
test.AddOutput<MLFloat16>("Moment_2_Out", {3}, data.m2_new_half);
test.AddOutput<float>("W_Out", {3}, data.w_new);
test.AddMissingOptionalOutput<MLFloat16>();
test.AddOutput<MLFloat16>("FP16_W_Out", {3}, data.w_new_half);
test.AddAttribute("do_bias_correction", static_cast<int64_t>(0));
test.AddAttribute("weight_decay_mode", static_cast<int64_t>(0));
test.Run();
}
TEST(OptimizerTest, AdamOptimizerMixPrecision_FP16Weight_ClipNorm_Test) {
OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain);
AdamOptimizerInputOutput data;
// Expected FP32 Outputs
std::vector<float> m1_new = {0.13f, 0.23f, 0.33f};
std::vector<float> m2_new = {0.3997f, 0.4998f, 0.6001f};
std::vector<float> w_new = {0.8972168f, 1.8369141f, 2.7871094f};
// FP16 Outputs
std::vector<MLFloat16> m1_new_half;
std::vector<MLFloat16> m2_new_half;
std::vector<MLFloat16> w_new_half;
m1_new_half.resize(m1_new.size());
m2_new_half.resize(m2_new.size());
w_new_half.resize(w_new.size());
ConvertFloatToMLFloat16(m1_new.data(), m1_new_half.data(), int(m1_new.size()));
ConvertFloatToMLFloat16(m2_new.data(), m2_new_half.data(), int(m2_new.size()));
ConvertFloatToMLFloat16(w_new.data(), w_new_half.data(), int(w_new.size()));
test.AddInput<MLFloat16>("ETA", {}, data.eta_half);
test.AddInput<int64_t>("Update_Count", {}, {3});
test.AddInput<float>("W", {3}, data.w);
test.AddInput<MLFloat16>("G", {3}, data.g_half);
test.AddInput<MLFloat16>("Moment_1", {3}, data.m1_half);
test.AddInput<MLFloat16>("Moment_2", {3}, data.m2_half);
test.AddInput<MLFloat16>("FP16_W", {3}, data.w_half);
test.AddInput<float>("loss_scale", {1}, {1.0f});
test.AddInput<float>("grad_norm", {1}, {0.01f});
// Verify AdamOptimizer outputs
test.AddOutput<int64_t>("Update_Count_Out", {}, {4});
test.AddOutput<MLFloat16>("Moment_1_Out", {3}, m1_new_half);
test.AddOutput<MLFloat16>("Moment_2_Out", {3}, m2_new_half);
test.AddOutput<float>("W_Out", {3}, w_new);
test.AddMissingOptionalOutput<MLFloat16>();
test.AddOutput<MLFloat16>("FP16_W_Out", {3}, w_new_half);
test.AddAttribute("do_bias_correction", static_cast<int64_t>(0));
test.AddAttribute("weight_decay_mode", static_cast<int64_t>(0));
test.AddAttribute("max_norm_clip", 0.001f);
test.Run();
}
TEST(OptimizerTest, AdamOptimizerMixPrecision_FP16Weight_SkipUpdate_Test) {
OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain);
AdamOptimizerInputOutput data;
@ -377,17 +458,18 @@ void compute_lamb(
/* momentum */ const std::vector<float>& m,
/* 2nd-order momentum */ const std::vector<float>& v,
const float eta,
const float loss_scale,
const float g_norm,
const float lambda,
const float alpha,
const float beta,
const float epsilon,
const float max_norm_clip,
/* updated weights */ std::vector<float>& w_new,
/* updated gradients */ std::vector<float>& g_new,
/* updated momentum */ std::vector<float>& m_new,
/* updated 2nd-order momentum */ std::vector<float>& v_new,
const int64_t step = 0,
const float loss_scale = 1.0f,
const float* p_scaled_g_norm = nullptr,
const float ratio_min = -std::numeric_limits<float>::infinity(),
const float ratio_max = std::numeric_limits<float>::infinity()) {
// Element counts of all vector-typed arguments.
@ -400,9 +482,12 @@ void compute_lamb(
// Buffer to store update direction.
std::vector<float> r(size, 0.0f);
float g_scale = loss_scale;
if (g_norm > loss_scale) {
g_scale *= g_norm / loss_scale;
float scaled_g_scaling_factor = loss_scale;
if (p_scaled_g_norm != nullptr) {
const float scaled_g_max_norm = loss_scale * max_norm_clip;
if (*p_scaled_g_norm > scaled_g_max_norm) {
scaled_g_scaling_factor = *p_scaled_g_norm / max_norm_clip;
}
}
const float alpha_correction = step > 0 ? 1.f - std::pow(alpha, static_cast<float>(step)) : 1.f;
@ -410,7 +495,7 @@ void compute_lamb(
// Compute new 1st-, 2nd-order momentums, and the update direction.
for (int i = 0; i < size; ++i) {
const float g_scaled = g[i] / g_scale;
const float g_scaled = g[i] / scaled_g_scaling_factor;
m_new[i] = alpha * m[i] + (1.0f - alpha) * g_scaled;
v_new[i] = beta * v[i] + (1.0f - beta) * g_scaled * g_scaled;
const float m_new_tmp = m_new[i] / alpha_correction;
@ -462,6 +547,7 @@ void run_lamb_test_with_baseline(
const float beta,
const float lambda,
const float epsilon,
const float max_norm,
const std::vector<T2>& w_new,
const std::vector<T3>& g_new,
const std::vector<T4>& m_new,
@ -470,13 +556,19 @@ void run_lamb_test_with_baseline(
const std::vector<MLFloat16>& w_new_half = {},
const bool do_update = true,
const int64_t step = 0,
const float loss_scale = 1.0f,
const float* p_g_norm = nullptr,
const float ratio_min = -std::numeric_limits<float>::infinity(),
const float ratio_max = std::numeric_limits<float>::infinity()) {
OpTester test("LambOptimizer", 1, onnxruntime::kMSDomain, true);
test.AddInput<bool>("update_signal", {1}, {do_update});
test.AddMissingOptionalInput<T2>();
test.AddMissingOptionalInput<T2>();
test.AddInput<T2>("loss_scale", {}, {loss_scale});
if (p_g_norm == nullptr) {
test.AddMissingOptionalInput<T2>();
} else {
test.AddInput<T2>("gradient_norm", {}, {T2(*p_g_norm)});
}
test.AddInput<T1>("ETA", {1}, eta);
if (step > 0) {
test.AddInput<int64_t>("Step", {}, {step});
@ -497,6 +589,7 @@ void run_lamb_test_with_baseline(
test.AddAttribute("beta", std::vector<float>(1, beta));
test.AddAttribute("lambda", std::vector<float>(1, lambda));
test.AddAttribute("epsilon", std::vector<float>(1, epsilon));
test.AddAttribute("max_norm_clip", std::vector<float>(1, max_norm));
test.AddAttribute("ratio_min", ratio_min);
test.AddAttribute("ratio_max", ratio_max);
@ -530,8 +623,6 @@ template <typename T1, typename T2, typename T3, typename T4>
void run_multi_tensor_lamb_test_with_baseline(
const std::vector<std::vector<int64_t>>& shapes,
const T1 eta,
const T1 loss_scale,
const T1 g_norm,
const std::vector<std::vector<T2>>& ws,
const std::vector<std::vector<T3>>& gs,
const std::vector<std::vector<T4>>& ms,
@ -540,6 +631,7 @@ void run_multi_tensor_lamb_test_with_baseline(
const std::vector<float>& betas,
const std::vector<float>& lambdas,
const std::vector<float>& epsilons,
const std::vector<float>& max_norms,
const std::vector<std::vector<T2>>& w_news,
const std::vector<std::vector<T3>>& g_news,
const std::vector<std::vector<T4>>& m_news,
@ -548,6 +640,8 @@ void run_multi_tensor_lamb_test_with_baseline(
const std::vector<std::vector<MLFloat16>>& w_new_halfs = {},
const bool do_update = true,
const int64_t step = 0,
const float loss_scale = 1.0f,
const float* p_g_norm = nullptr,
const float ratio_min = -std::numeric_limits<float>::infinity(),
const float ratio_max = std::numeric_limits<float>::infinity()) {
OpTester test("LambOptimizer", 1, onnxruntime::kMSDomain, true);
@ -560,6 +654,7 @@ void run_multi_tensor_lamb_test_with_baseline(
ORT_ENFORCE(shapes.size() == betas.size());
ORT_ENFORCE(shapes.size() == lambdas.size());
ORT_ENFORCE(shapes.size() == epsilons.size());
ORT_ENFORCE(shapes.size() == max_norms.size());
if (!w_news.empty()) {
ORT_ENFORCE(shapes.size() == w_news.size());
}
@ -578,8 +673,12 @@ void run_multi_tensor_lamb_test_with_baseline(
const int group_count = static_cast<int>(ws.size());
test.AddInput<bool>("update_signal", {}, {do_update});
test.AddInput<T1>("loss_scale", {}, {loss_scale});
test.AddInput<T1>("gradient_norm", {}, {g_norm});
test.AddInput<T2>("loss_scale", {}, {loss_scale});
if (p_g_norm == nullptr) {
test.AddMissingOptionalInput<T2>();
} else {
test.AddInput<float>("gradient_norm", {}, {T2(*p_g_norm)});
}
test.AddInput<T1>("ETA", {}, {eta});
if (step > 0) {
test.AddInput<int64_t>("Step", {}, {step});
@ -633,6 +732,7 @@ void run_multi_tensor_lamb_test_with_baseline(
test.AddAttribute("beta", betas);
test.AddAttribute("lambda", lambdas);
test.AddAttribute("epsilon", epsilons);
test.AddAttribute("max_norm_clip", max_norms);
test.AddAttribute("ratio_min", ratio_min);
test.AddAttribute("ratio_max", ratio_max);
@ -645,8 +745,6 @@ void run_multi_tensor_lamb_test_with_baseline(
void run_multi_tensor_lamb_test(
const std::vector<std::vector<int64_t>> shapes,
const float eta,
const float loss_scale,
const float g_norm,
const std::vector<std::vector<float>> ws,
const std::vector<std::vector<float>> gs,
const std::vector<std::vector<float>> ms,
@ -655,7 +753,10 @@ void run_multi_tensor_lamb_test(
const std::vector<float> alphas,
const std::vector<float> betas,
const std::vector<float> epsilons,
const std::vector<float> max_norms,
const int64_t step = 0,
const float loss_scale = 1.0f,
const float* p_scaled_g_norm = nullptr,
const float ratio_min = -std::numeric_limits<float>::infinity(),
const float ratio_max = std::numeric_limits<float>::infinity()) {
// Check if parallel vectors have the same length.
@ -667,6 +768,7 @@ void run_multi_tensor_lamb_test(
ORT_ENFORCE(shapes.size() == betas.size());
ORT_ENFORCE(shapes.size() == lambdas.size());
ORT_ENFORCE(shapes.size() == epsilons.size());
ORT_ENFORCE(shapes.size() == max_norms.size());
const int group_count = static_cast<int>(ws.size());
@ -685,9 +787,8 @@ void run_multi_tensor_lamb_test(
// Invoke LAMB's reference implementation to compute baseline output.
compute_lamb(
shapes[i], ws[i], gs[i], ms[i], vs[i],
eta, loss_scale, g_norm,
lambdas[i], alphas[i], betas[i], epsilons[i],
w_news[i], g_news[i], m_news[i], v_news[i], step,
eta, lambdas[i], alphas[i], betas[i], epsilons[i], max_norms[i],
w_news[i], g_news[i], m_news[i], v_news[i], step, loss_scale, p_scaled_g_norm,
ratio_min, ratio_max);
}
@ -695,18 +796,18 @@ void run_multi_tensor_lamb_test(
// Output new weights.
run_multi_tensor_lamb_test_with_baseline(
shapes, eta, loss_scale, g_norm,
shapes, eta,
ws, gs, ms, vs,
alphas, betas, lambdas, epsilons,
w_news, {}, m_news, v_news, {}, {}, true, step,
alphas, betas, lambdas, epsilons, max_norms,
w_news, {}, m_news, v_news, {}, {}, true, step, loss_scale, p_scaled_g_norm,
ratio_min, ratio_max);
// Output new gradients.
run_multi_tensor_lamb_test_with_baseline(
shapes, eta, loss_scale, g_norm,
shapes, eta,
ws, gs, ms, vs,
alphas, betas, lambdas, epsilons,
{}, g_news, m_news, v_news, {}, {}, true, step,
alphas, betas, lambdas, epsilons, max_norms,
{}, g_news, m_news, v_news, {}, {}, true, step, loss_scale, p_scaled_g_norm,
ratio_min, ratio_max);
}
@ -721,7 +822,10 @@ void run_lamb_mix_precision_test(
const float alpha,
const float beta,
const float epsilon,
const int64_t step = 0) {
const float max_norm,
const int64_t step = 0,
const float loss_scale = 1.0f,
const float* p_g_norm = nullptr) {
std::vector<float> w_new(w.size(), 0);
std::vector<float> g_new(g.size(), 0);
std::vector<float> m_new(m.size(), 0);
@ -730,8 +834,8 @@ void run_lamb_mix_precision_test(
// Invoke LAMB's reference implementation to compute output.
compute_lamb(
shape, w, g, m, v,
eta[0], 1.f, 1.f, lambda, alpha, beta, epsilon,
w_new, g_new, m_new, v_new, step);
eta[0], lambda, alpha, beta, epsilon, max_norm,
w_new, g_new, m_new, v_new, step, loss_scale, p_g_norm);
std::vector<MLFloat16> eta_half(eta.size());
std::vector<MLFloat16> g_half(w.size());
@ -755,47 +859,58 @@ void run_lamb_mix_precision_test(
// Half momentums, without fp16 weight
run_lamb_test_with_baseline(
shape, eta_half, w, g_half, m_half, v_half, alpha, beta, lambda, epsilon, w_new, {}, m_new_half, v_new_half, {}, {}, true, step);
shape, eta_half, w, g_half, m_half, v_half, alpha, beta, lambda, epsilon, max_norm,
w_new, {}, m_new_half, v_new_half, {}, {}, true, step, loss_scale, p_g_norm);
// Float momentums, without fp16 weight
run_lamb_test_with_baseline(
shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, w_new, {}, m_new, v_new, {}, {}, true, step);
shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, max_norm,
w_new, {}, m_new, v_new, {}, {}, true, step, loss_scale, p_g_norm);
// Half momentums, with fp16 weight
run_lamb_test_with_baseline(
shape, eta_half, w, g_half, m_half, v_half, alpha, beta, lambda, epsilon, w_new, {}, m_new_half, v_new_half, {}, {}, true, step);
shape, eta_half, w, g_half, m_half, v_half, alpha, beta, lambda, epsilon, max_norm,
w_new, {}, m_new_half, v_new_half, {}, {}, true, step, loss_scale, p_g_norm);
// Float momentums, with fp16 weight
run_lamb_test_with_baseline(
shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, w_new, {}, m_new, v_new, w_half, w_new_half, true, step);
shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, max_norm,
w_new, {}, m_new, v_new, w_half, w_new_half, true, step, loss_scale, p_g_norm);
// Half momentums, with fp16 weight, skip weight update
run_lamb_test_with_baseline(
shape, eta_half, w, g_half, m_half, v_half, alpha, beta, lambda, epsilon, w, {}, m_half, v_half, w_half, w_half, false, step);
shape, eta_half, w, g_half, m_half, v_half, alpha, beta, lambda, epsilon, max_norm,
w, {}, m_half, v_half, w_half, w_half, false, step, loss_scale, p_g_norm);
// Float momentums, with fp16 weight, skip weight update
run_lamb_test_with_baseline(
shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, w, {}, m, v, w_half, w_half, false, step);
shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, max_norm,
w, {}, m, v, w_half, w_half, false, step, loss_scale, p_g_norm);
// Float eta, float momentums, with fp16 weight
run_lamb_test_with_baseline(
shape, eta, w, g_half, m, v, alpha, beta, lambda, epsilon, w_new, {}, m_new, v_new, w_half, w_new_half, true, step);
shape, eta, w, g_half, m, v, alpha, beta, lambda, epsilon, max_norm,
w_new, {}, m_new, v_new, w_half, w_new_half, true, step, loss_scale, p_g_norm);
// Float eta, float momentums, with fp16 weight, skip weight update
run_lamb_test_with_baseline(
shape, eta, w, g_half, m, v, alpha, beta, lambda, epsilon, w, {}, m, v, w_half, w_half, false, step);
shape, eta, w, g_half, m, v, alpha, beta, lambda, epsilon, max_norm,
w, {}, m, v, w_half, w_half, false, step, loss_scale, p_g_norm);
// Float momentums, without fp16 weight, output gradients only
run_lamb_test_with_baseline(
shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, {}, g_new_half, m_new, v_new, {}, {}, true, step);
shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, max_norm,
{}, g_new_half, m_new, v_new, {}, {}, true, step, loss_scale, p_g_norm);
// Float momentums, with fp16 weight, output gradients only
run_lamb_test_with_baseline(
shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, {}, g_new_half, m_new, v_new, w_half, {}, true, step);
shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, max_norm,
{}, g_new_half, m_new, v_new, w_half, {}, true, step, loss_scale, p_g_norm);
// Float momentums, with fp16 weight, output gradients only, skip weight update
run_lamb_test_with_baseline(
shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, {}, g_half, m, v, w_half, {}, false, step);
shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, max_norm,
{}, g_half, m, v, w_half, {}, false, step, loss_scale, p_g_norm);
}
// A optimizer test with an 2-element vector.
@ -811,12 +926,14 @@ TEST(OptimizerTest, LambOptimizerTestVector) {
const float alpha = 0.2f;
const float beta = 0.8f;
const float epsilon = 1e-6f;
const float max_norm = 1.0f;
const int64_t step = 0;
const float loss_scale = 1.f;
const float scaled_g_norm = 1.f;
run_multi_tensor_lamb_test(
{shape},
eta,
1.f,
1.f,
{w},
{g},
{m},
@ -824,7 +941,11 @@ TEST(OptimizerTest, LambOptimizerTestVector) {
{lambda},
{alpha},
{beta},
{epsilon});
{epsilon},
{max_norm},
step,
loss_scale,
&scaled_g_norm);
}
TEST(OptimizerTest, LambOptimizerTestVectorWithZeroWeight) {
@ -839,12 +960,13 @@ TEST(OptimizerTest, LambOptimizerTestVectorWithZeroWeight) {
const float alpha = 0.2f;
const float beta = 0.8f;
const float epsilon = 1e-6f;
const float max_norm = 1.0f;
const int64_t step = 0;
const float loss_scale = 1.f;
const float scaled_g_norm = 1.f;
run_multi_tensor_lamb_test(
{shape},
eta,
1.f,
1.f,
{w},
{g},
{m},
@ -852,7 +974,11 @@ TEST(OptimizerTest, LambOptimizerTestVectorWithZeroWeight) {
{lambda},
{alpha},
{beta},
{epsilon});
{epsilon},
{max_norm},
step,
loss_scale,
&scaled_g_norm);
}
TEST(OptimizerTest, LambOptimizerRatioMin) {
@ -867,14 +993,16 @@ TEST(OptimizerTest, LambOptimizerRatioMin) {
const float alpha = 0.2f;
const float beta = 0.8f;
const float epsilon = 1e-6f;
const float max_norm = 1.0f;
const int64_t step = 0;
const float loss_scale = 1.f;
const float scaled_g_norm = 1.f;
const float ratio_min = -std::numeric_limits<float>::infinity();
const float ratio_max = 0.1f;
run_multi_tensor_lamb_test(
{shape},
eta,
1.f,
1.f,
{w},
{g},
{m},
@ -883,7 +1011,10 @@ TEST(OptimizerTest, LambOptimizerRatioMin) {
{alpha},
{beta},
{epsilon},
0,
{max_norm},
step,
loss_scale,
&scaled_g_norm,
ratio_min,
ratio_max);
}
@ -900,14 +1031,16 @@ TEST(OptimizerTest, LambOptimizerRatioMax) {
const float alpha = 0.2f;
const float beta = 0.8f;
const float epsilon = 1e-6f;
const float max_norm = 1.0f;
const int64_t step = 0;
const float loss_scale = 1.f;
const float scaled_g_norm = 1.f;
const float ratio_min = 1.0f;
const float ratio_max = std::numeric_limits<float>::infinity();
run_multi_tensor_lamb_test(
{shape},
eta,
1.f,
1.f,
{w},
{g},
{m},
@ -916,7 +1049,10 @@ TEST(OptimizerTest, LambOptimizerRatioMax) {
{alpha},
{beta},
{epsilon},
0,
{max_norm},
step,
loss_scale,
&scaled_g_norm,
ratio_min,
ratio_max);
}
@ -933,12 +1069,14 @@ TEST(OptimizerTest, LambOptimizerTestBiasCorrectionFirst) {
const float alpha = 0.2f;
const float beta = 0.8f;
const float epsilon = 1e-6f;
const float max_norm = 1.0f;
const int64_t step = 1;
const float loss_scale = 1.f;
const float scaled_g_norm = 1.f;
run_multi_tensor_lamb_test(
{shape},
eta,
1.f,
1.f,
{w},
{g},
{m},
@ -947,7 +1085,10 @@ TEST(OptimizerTest, LambOptimizerTestBiasCorrectionFirst) {
{alpha},
{beta},
{epsilon},
1);
{max_norm},
step,
loss_scale,
&scaled_g_norm);
}
TEST(OptimizerTest, LambOptimizerTestBiasCorrectionThird) {
@ -962,12 +1103,14 @@ TEST(OptimizerTest, LambOptimizerTestBiasCorrectionThird) {
const float alpha = 0.2f;
const float beta = 0.8f;
const float epsilon = 1e-6f;
const float max_norm = 1.0f;
const int64_t step = 3;
const float loss_scale = 1.f;
const float scaled_g_norm = 1.f;
run_multi_tensor_lamb_test(
{shape},
eta,
1.f,
1.f,
{w},
{g},
{m},
@ -976,7 +1119,10 @@ TEST(OptimizerTest, LambOptimizerTestBiasCorrectionThird) {
{alpha},
{beta},
{epsilon},
3);
{max_norm},
step,
loss_scale,
&scaled_g_norm);
}
// A optimizer test with an 2-by-1-by-1-by-1 tensor.
@ -992,12 +1138,14 @@ TEST(OptimizerTest, LambOptimizerTest4DTensor) {
const float alpha = 0.2f;
const float beta = 0.8f;
const float epsilon = 1e-6f;
const float max_norm = 1.0f;
const int64_t step = 0;
const float loss_scale = 1.f;
const float scaled_g_norm = 1.f;
run_multi_tensor_lamb_test(
{shape},
eta,
1.f,
1.f,
{w},
{g},
{m},
@ -1005,7 +1153,11 @@ TEST(OptimizerTest, LambOptimizerTest4DTensor) {
{lambda},
{alpha},
{beta},
{epsilon});
{epsilon},
{max_norm},
step,
loss_scale,
&scaled_g_norm);
}
// A optimizer test with an 2-by-3 tensor.
@ -1021,12 +1173,14 @@ TEST(OptimizerTest, LambOptimizerTest2by3Tensor) {
const float alpha = 0.2f;
const float beta = 0.8f;
const float epsilon = 1e-6f;
const float max_norm = 1.0f;
const int64_t step = 0;
const float loss_scale = 1.f;
const float scaled_g_norm = 1.f;
run_multi_tensor_lamb_test(
{shape},
eta,
1.f,
1.f,
{w},
{g},
{m},
@ -1034,7 +1188,11 @@ TEST(OptimizerTest, LambOptimizerTest2by3Tensor) {
{lambda},
{alpha},
{beta},
{epsilon});
{epsilon},
{max_norm},
step,
loss_scale,
&scaled_g_norm);
}
// A optimizer test with an 1-element tensor.
@ -1050,6 +1208,10 @@ TEST(OptimizerTest, LambOptimizerTestScalar) {
const float alpha = 0.2f;
const float beta = 0.8f;
const float epsilon = 1e-6f;
const float max_norm = 1.0f;
const int64_t step = 0;
const float loss_scale = 1.f;
const float scaled_g_norm = 1.f;
// Intermediate and output buffers of the optimizer.
std::vector<float> m_new = {0.0f};
@ -1059,8 +1221,6 @@ TEST(OptimizerTest, LambOptimizerTestScalar) {
run_multi_tensor_lamb_test(
{shape},
eta,
1.f,
1.f,
{w},
{g},
{m},
@ -1068,7 +1228,51 @@ TEST(OptimizerTest, LambOptimizerTestScalar) {
{lambda},
{alpha},
{beta},
{epsilon});
{epsilon},
{max_norm},
step,
loss_scale,
&scaled_g_norm);
}
// A optimizer test with an 1-element tensor.
TEST(OptimizerTest, LambOptimizerTestScalar_NonDefaultMaxNormClipping) {
// Input tensors and attributes.
const std::vector<int64_t> shape = {(int64_t)1};
const float eta = 0.5f;
const std::vector<float> w = {1.0f};
const std::vector<float> g = {3.0f};
const std::vector<float> m = {-10.0f};
const std::vector<float> v = {1.0f};
const float lambda = 0.5f;
const float alpha = 0.2f;
const float beta = 0.8f;
const float epsilon = 1e-6f;
const float max_norm = 0.1f;
// Intermediate and output buffers of the optimizer.
std::vector<float> m_new = {0.0f};
std::vector<float> v_new = {0.0f};
std::vector<float> w_new = {0.0f};
const int64_t step = 0;
const float loss_scale = 1.0f;
const float scaled_g_norm = g[0];
run_multi_tensor_lamb_test(
{shape},
eta,
{w},
{g},
{m},
{v},
{lambda},
{alpha},
{beta},
{epsilon},
{max_norm},
step,
loss_scale,
&scaled_g_norm);
}
TEST(OptimizerTest, LambOptimizerTestScalarScaling) {
@ -1083,6 +1287,10 @@ TEST(OptimizerTest, LambOptimizerTestScalarScaling) {
const float alpha = 0.2f;
const float beta = 0.8f;
const float epsilon = 1e-6f;
const float max_norm = 1.0f;
const int64_t step = 0;
const float loss_scale = 8.f;
const float scaled_g_norm = 4.f;
// Intermediate and output buffers of the optimizer.
std::vector<float> m_new = {0.0f};
@ -1092,8 +1300,6 @@ TEST(OptimizerTest, LambOptimizerTestScalarScaling) {
run_multi_tensor_lamb_test(
{shape},
eta,
8.f,
4.f,
{w},
{g},
{m},
@ -1101,7 +1307,11 @@ TEST(OptimizerTest, LambOptimizerTestScalarScaling) {
{lambda},
{alpha},
{beta},
{epsilon});
{epsilon},
{max_norm},
step,
loss_scale,
&scaled_g_norm);
}
TEST(OptimizerTest, LambOptimizerTestExternalBaseline) {
@ -1125,6 +1335,7 @@ TEST(OptimizerTest, LambOptimizerTestExternalBaseline) {
const float alpha = 0.1f;
const float beta = 0.01f;
const float epsilon = 0.1f;
const float max_norm = 1.0f;
std::vector<float> w_new = {
0.02979828f, 0.13677707f, -0.22708717f, -0.20361158f, -0.15338624f, 0.1081504f,
@ -1141,11 +1352,11 @@ TEST(OptimizerTest, LambOptimizerTestExternalBaseline) {
// Output new weights
run_lamb_test_with_baseline(
shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, w_new, {}, m_new, v_new);
shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, max_norm, w_new, {}, m_new, v_new);
// Output new gradients
run_lamb_test_with_baseline(
shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, {}, g_new, m_new, v_new);
shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, max_norm, {}, g_new, m_new, v_new);
}
TEST(OptimizerTest, LambOptimizerTestExternalBaselineDouble) {
@ -1169,6 +1380,7 @@ TEST(OptimizerTest, LambOptimizerTestExternalBaselineDouble) {
const float alpha = 0.1f;
const float beta = 0.01f;
const float epsilon = 0.1f;
const float max_norm = 1.0f;
std::vector<double> w_new = {
0.02979828, 0.13677707, -0.22708717, -0.20361158, -0.15338624, 0.1081504,
@ -1185,11 +1397,11 @@ TEST(OptimizerTest, LambOptimizerTestExternalBaselineDouble) {
// Output new weights
run_lamb_test_with_baseline(
shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, w_new, {}, m_new, v_new);
shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, max_norm, w_new, {}, m_new, v_new);
// Output new gradients
run_lamb_test_with_baseline(
shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, {}, g_new, m_new, v_new);
shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, max_norm, {}, g_new, m_new, v_new);
}
TEST(OptimizerTest, LambOptimizerTest5DTensorMixPrecision32_16) {
@ -1204,10 +1416,16 @@ TEST(OptimizerTest, LambOptimizerTest5DTensorMixPrecision32_16) {
const float alpha = 1.5f;
const float beta = 1.5f;
const float epsilon = 1.0f;
const float max_norm = 1.0f;
const float loss_scale = 1.0f;
run_lamb_mix_precision_test(
shape, eta, w, g, m, v, lambda, alpha, beta, epsilon, max_norm);
float gradient_norm = GetGradientL2Norm(g);
// gradient clipping
run_lamb_mix_precision_test(
shape, eta, w, g, m, v,
lambda, alpha, beta, epsilon);
lambda, alpha, beta, epsilon, max_norm, 0, loss_scale, &gradient_norm);
}
TEST(OptimizerTest, LambOptimizerTestSimpleBaselineMixPrecision32_16) {
@ -1222,10 +1440,17 @@ TEST(OptimizerTest, LambOptimizerTestSimpleBaselineMixPrecision32_16) {
const float alpha = 1.0f;
const float beta = 1.0f;
const float epsilon = 1.0f;
const float max_norm = 1.0f;
const float loss_scale = 1.0f;
run_lamb_mix_precision_test(
shape, eta, w, g, m, v,
lambda, alpha, beta, epsilon);
lambda, alpha, beta, epsilon, max_norm);
// gradient clipping
float gradient_norm = GetGradientL2Norm(g);
run_lamb_mix_precision_test(
shape, eta, w, g, m, v,
lambda, alpha, beta, epsilon, max_norm, 0, loss_scale, &gradient_norm);
}
TEST(OptimizerTest, LambOptimizerTestBaselineMixPrecision32_16) {
@ -1240,10 +1465,18 @@ TEST(OptimizerTest, LambOptimizerTestBaselineMixPrecision32_16) {
const float alpha = 0.9f;
const float beta = 0.95f;
const float epsilon = 0.33f;
const float max_norm = 1.0f;
const float loss_scale = 1.0f;
run_lamb_mix_precision_test(
shape, eta, w, g, m, v,
lambda, alpha, beta, epsilon);
lambda, alpha, beta, epsilon, max_norm);
// gradient clipping
float gradient_norm = GetGradientL2Norm(g);
run_lamb_mix_precision_test(
shape, eta, w, g, m, v,
lambda, alpha, beta, epsilon, max_norm, 0, loss_scale, &gradient_norm);
}
TEST(OptimizerTest, LambOptimizerTestScalarMixPrecision32_16) {
@ -1258,14 +1491,44 @@ TEST(OptimizerTest, LambOptimizerTestScalarMixPrecision32_16) {
const float alpha = 0.9f;
const float beta = 0.95f;
const float epsilon = 0.33f;
const float max_norm = 1.0f;
const float loss_scale = 1.0f;
run_lamb_mix_precision_test(
shape, eta, w, g, m, v,
lambda, alpha, beta, epsilon);
lambda, alpha, beta, epsilon, max_norm);
// gradient clipping
float gradient_norm = GetGradientL2Norm(g);
run_lamb_mix_precision_test(
shape, eta, w, g, m, v,
lambda, alpha, beta, epsilon, max_norm, 2, loss_scale, &gradient_norm);
}
TEST(OptimizerTest, LambOptimizerTestScalarMixPrecision32_16_NoDefaultMaxNormClipping) {
const std::vector<int64_t> shape = {1};
const std::vector<float> eta = {0.1f};
const std::vector<float> w = {-1.5f};
const std::vector<float> g = {-0.75f};
const std::vector<float> m = {0.87f};
const std::vector<float> v = {0.12f};
const float lambda = 0.25f;
const float alpha = 0.9f;
const float beta = 0.95f;
const float epsilon = 0.33f;
const float max_norm = 0.1f;
const float loss_scale = 1.0f;
run_lamb_mix_precision_test(
shape, eta, w, g, m, v,
lambda, alpha, beta, epsilon, 2);
lambda, alpha, beta, epsilon, max_norm, 2);
// gradient clipping
float gradient_norm = GetGradientL2Norm(g);
run_lamb_mix_precision_test(
shape, eta, w, g, m, v,
lambda, alpha, beta, epsilon, max_norm, 2, loss_scale, &gradient_norm);
}
TEST(OptimizerTest, LambOptimizerTestLarge) {
@ -1292,12 +1555,14 @@ TEST(OptimizerTest, LambOptimizerTestLarge) {
const float alpha = 0.2f;
const float beta = 0.8f;
const float epsilon = 1e-6f;
const float max_norm = 1.0f;
const int64_t step = 0;
const float loss_scale = 1.f;
const float scaled_g_norm = 1.f;
run_multi_tensor_lamb_test(
{shape},
eta,
1.f,
1.f,
{w},
{g},
{m},
@ -1305,7 +1570,11 @@ TEST(OptimizerTest, LambOptimizerTestLarge) {
{lambda},
{alpha},
{beta},
{epsilon});
{epsilon},
{max_norm},
step,
loss_scale,
&scaled_g_norm);
}
}
@ -1327,6 +1596,7 @@ TEST(OptimizerTest, LambOptimizerMultiTensorRatio) {
std::vector<float> betas(group_count);
std::vector<float> lambdas(group_count);
std::vector<float> epsilons(group_count);
std::vector<float> max_norms(group_count);
const float eta = dist(random_engine);
@ -1351,20 +1621,26 @@ TEST(OptimizerTest, LambOptimizerMultiTensorRatio) {
betas[i] = dist(random_engine);
lambdas[i] = dist(random_engine);
epsilons[i] = dist(random_engine);
max_norms[i] = dist(random_engine);
}
run_multi_tensor_lamb_test(
shapes, eta, 1.f, 1.f,
ws, gs, ms, vs,
lambdas, alphas, betas, epsilons,
0, 0.3f, 0.7f);
const int64_t step = 0;
float loss_scale = 1.f;
const float scaled_g_norm = 1.f;
run_multi_tensor_lamb_test(
shapes, eta, 1.f, 1.f,
shapes, eta,
ws, gs, ms, vs,
lambdas, alphas, betas, epsilons);
lambdas, alphas, betas, epsilons, max_norms,
step, loss_scale, &scaled_g_norm, 0.3f, 0.7f);
run_multi_tensor_lamb_test(
shapes, eta,
ws, gs, ms, vs,
lambdas, alphas, betas, epsilons, max_norms,
step, loss_scale, &scaled_g_norm);
}
#endif
}
} // namespace test
} // namespace onnxruntime

68
orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py
View file

@ -426,6 +426,7 @@ def testOptimizerConfigAdam():
assert_allclose(0.0, cfg.lambda_coef, rtol=rtol,
err_msg="lambda_coef mismatch")
assert_allclose(1e-8, cfg.epsilon, rtol=rtol, err_msg="epsilon mismatch")
assert_allclose(1.0, cfg.max_norm_clip, rtol=rtol, err_msg="max_norm_clip mismatch")
assert cfg.do_bias_correction == True, "lambda_coef mismatch"
assert cfg.weight_decay_mode == optim.AdamConfig.DecayMode.BEFORE_WEIGHT_UPDATE, "weight_decay_mode mismatch"
@ -443,6 +444,7 @@ def testOptimizerConfigLamb():
assert cfg.ratio_min == float('-inf'), "ratio_min mismatch"
assert cfg.ratio_max == float('inf'), "ratio_max mismatch"
assert_allclose(1e-6, cfg.epsilon, rtol=rtol, err_msg="epsilon mismatch")
assert_allclose(1.0, cfg.max_norm_clip, rtol=rtol, err_msg="max_norm_clip mismatch")
assert cfg.do_bias_correction == False, "do_bias_correction mismatch"
@ -1468,3 +1470,69 @@ def testTrainingGraphExport(debug_files):
os.remove(path)
else:
assert not os.path.isfile(path)
@pytest.mark.parametrize("seed,device,max_norm_clip,gradient_accumulation_steps,total_steps,expected_loss", [
(0, 'cuda', 1.0, 1, 12, [10.536802, 9.95102, 9.495312, 9.067217, 8.735067, 8.447508,\
8.179443, 7.903837, 7.655049, 7.409669, 7.135822, 6.931838]),
(0, 'cuda', 0.1, 1, 12, [10.536802, 9.951735, 9.496659, 9.069328, 8.7381115, 8.4513855,\
8.184143, 7.9093056, 7.661127, 7.4162436, 7.142842, 6.9388437]),
(42, 'cuda', 1.0, 1, 12, [10.645588, 10.0333, 9.52253, 9.108369, 8.766306, 8.497426,\
8.199408, 7.958235, 7.659668, 7.459833, 7.170661, 6.9139776]),
(42, 'cuda', 0.1, 1, 12, [10.645588, 10.03406, 9.524019, 9.110594, 8.769308, 8.501322,\
8.204281, 7.963957, 7.6660814, 7.46682, 7.1780496, 6.92159]),
])
def testORTTrainerAdamMaxNormClip(seed, device, max_norm_clip, gradient_accumulation_steps,total_steps, expected_loss):
rtol = 1e-5
torch.manual_seed(seed)
set_seed(seed)
# Setup ORTTrainer
options = orttrainer.ORTTrainerOptions({'device' : {'id' : device},
'batch' : {'gradient_accumulation_steps' : gradient_accumulation_steps},
'debug' : {'deterministic_compute' : True}})
model, model_desc, my_loss, batcher_fn, train_data, _, _ = _test_commons._load_pytorch_transformer_model(device)
optim_config = optim.AdamConfig(lr=0.001, max_norm_clip=max_norm_clip)
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, loss_fn=my_loss, options=options)
# Training loop
actual_loss = []
for i in range(total_steps):
data, targets = batcher_fn(train_data, i)
loss, _ = trainer.train_step(data, targets)
actual_loss.append(loss.cpu())
# Compare legacy vs experimental APIs
_test_helpers.assert_model_outputs(expected_loss, actual_loss, rtol=rtol)
@pytest.mark.parametrize("seed,device,max_norm_clip, gradient_accumulation_steps,total_steps,expected_loss", [
(0, 'cuda', 1.0, 1, 12, [10.536802, 10.409792, 10.354762, 10.253063, 10.213676, 10.113361,\
10.066136, 9.977713, 9.924597, 9.858974, 9.796471, 9.794921]),
(0, 'cuda', 0.1, 1, 12, [10.536802, 10.3714695, 10.276415, 10.13743, 10.063246, 9.93144,\
9.854875, 9.739198, 9.661381, 9.570321, 9.482681, 9.457669]),
(42, 'cuda', 1.0, 1, 12, [10.645588, 10.51151, 10.438802, 10.356055, 10.291667, 10.232069,\
10.168237, 10.074414, 9.990586, 9.9324, 9.891901, 9.788895]),
(42, 'cuda', 0.1, 1, 12, [10.645588, 10.473022, 10.359108, 10.238948, 10.141735, 10.049339,\
9.953887, 9.832249, 9.722989, 9.640278, 9.572205, 9.448381]),
])
def testORTTrainerLambMaxNormClip(seed, device, max_norm_clip, gradient_accumulation_steps, total_steps, expected_loss):
rtol = 1e-3
torch.manual_seed(seed)
set_seed(seed)
# Setup ORTTrainer
options = orttrainer.ORTTrainerOptions({'device' : {'id' : device},
'batch' : {'gradient_accumulation_steps' : gradient_accumulation_steps},
'debug' : {'deterministic_compute' : True}})
model, model_desc, my_loss, batcher_fn, train_data, _, _ = _test_commons._load_pytorch_transformer_model(device)
optim_config = optim.LambConfig(lr=0.001, max_norm_clip=max_norm_clip)
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, loss_fn=my_loss, options=options)
# Training loop
actual_loss = []
for i in range(total_steps):
data, targets = batcher_fn(train_data, i)
loss, _ = trainer.train_step(data, targets)
actual_loss.append(loss.cpu())
# Compare legacy vs experimental APIs
_test_helpers.assert_model_outputs(expected_loss, actual_loss, rtol=rtol)

51
orttraining/orttraining/training_ops/cuda/optimizer/adam.cc
View file

@ -11,30 +11,30 @@ namespace onnxruntime {
namespace cuda {
// TODO: Once Schema is checked in to onnx lets fix this to match that
#define REGISTER_ADAM_KERNEL_TYPED(T1, T2, T3, T4, T_GRAD, T_GRAD_NORM, T_MIXED_PRECISION_FP) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
AdamOptimizer, \
kMSDomain, \
1, \
T1##_##T2##_##T3##_##T4##_##T_GRAD##_##T_GRAD_NORM##_##T_MIXED_PRECISION_FP, \
kCudaExecutionProvider, \
KernelDefBuilder() \
.Alias(1, 0) /* Update step count in-place */ \
.Alias(2, 3) /* Update weights in-place */ \
.Alias(3, 4) /* Update gradients in-place */ \
.Alias(4, 1) /* Update moment-1 in-place */ \
.Alias(5, 2) /* Update moment-2 in-place */ \
.Alias(6, 5) /* Update mixed_precision weights in-place */ \
.InputMemoryType<OrtMemTypeCPUInput>(1) /* Keep step count in CPU */ \
.InputMemoryType<OrtMemTypeCPUInput>(9) /* Keep do_update in CPU */ \
.OutputMemoryType<OrtMemTypeCPUOutput>(0) /* Keep step count in CPU */ \
.TypeConstraint("T1", DataTypeImpl::GetTensorType<T1>()) \
.TypeConstraint("T2", DataTypeImpl::GetTensorType<T2>()) \
.TypeConstraint("T3", DataTypeImpl::GetTensorType<T3>()) \
.TypeConstraint("T4", DataTypeImpl::GetTensorType<T4>()) \
.TypeConstraint("T_GRAD", DataTypeImpl::GetTensorType<T_GRAD>()) \
.TypeConstraint("T_MIXED_PRECISION_FP", DataTypeImpl::GetTensorType<T_MIXED_PRECISION_FP>()) \
.TypeConstraint("T_GRAD_NORM", DataTypeImpl::GetTensorType<T_GRAD_NORM>()), \
#define REGISTER_ADAM_KERNEL_TYPED(T1, T2, T3, T4, T_GRAD, T_GRAD_NORM, T_MIXED_PRECISION_FP) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
AdamOptimizer, \
kMSDomain, \
1, \
T1##_##T2##_##T3##_##T4##_##T_GRAD##_##T_GRAD_NORM##_##T_MIXED_PRECISION_FP, \
kCudaExecutionProvider, \
KernelDefBuilder() \
.Alias(1, 0) /* Update step count in-place */ \
.Alias(2, 3) /* Update weights in-place */ \
.Alias(3, 4) /* Update gradients in-place */ \
.Alias(4, 1) /* Update moment-1 in-place */ \
.Alias(5, 2) /* Update moment-2 in-place */ \
.Alias(6, 5) /* Update mixed_precision weights in-place */ \
.InputMemoryType<OrtMemTypeCPUInput>(1) /* Keep step count in CPU */ \
.InputMemoryType<OrtMemTypeCPUInput>(9) /* Keep do_update in CPU */ \
.OutputMemoryType<OrtMemTypeCPUOutput>(0) /* Keep step count in CPU */ \
.TypeConstraint("T1", DataTypeImpl::GetTensorType<T1>()) \
.TypeConstraint("T2", DataTypeImpl::GetTensorType<T2>()) \
.TypeConstraint("T3", DataTypeImpl::GetTensorType<T3>()) \
.TypeConstraint("T4", DataTypeImpl::GetTensorType<T4>()) \
.TypeConstraint("T_GRAD", DataTypeImpl::GetTensorType<T_GRAD>()) \
.TypeConstraint("T_MIXED_PRECISION_FP", DataTypeImpl::GetTensorType<T_MIXED_PRECISION_FP>()) \
.TypeConstraint("T_GRAD_NORM", DataTypeImpl::GetTensorType<T_GRAD_NORM>()), \
AdamOptimizer<T1, T2, T3, T4, T_GRAD, T_GRAD_NORM, T_MIXED_PRECISION_FP>);
REGISTER_ADAM_KERNEL_TYPED(float, int64_t, float, float, float, float, MLFloat16)
@ -106,7 +106,7 @@ Status AdamOptimizer<T1, T2, T3, T4, T_GRAD, T_GRAD_NORM, T_MIXED_PRECISION_FP>:
const T2* S_in = S.template Data<T2>();
T2* S_out = NS.template MutableData<T2>();
const CudaT_GRAD_NORM* G_norm = nullptr;
if (gradient_norm_tensor != nullptr) {
G_norm = reinterpret_cast<const CudaT_GRAD_NORM*>(gradient_norm_tensor->template Data<T_GRAD_NORM>());
@ -148,6 +148,7 @@ Status AdamOptimizer<T1, T2, T3, T4, T_GRAD, T_GRAD_NORM, T_MIXED_PRECISION_FP>:
ToCudaType<T4>::FromFloat(beta_),
ToCudaType<T4>::FromFloat(lambda_),
ToCudaType<T4>::FromFloat(epsilon_),
ToCudaType<T4>::FromFloat(max_norm_clip_),
do_bias_correction_,
weight_decay_mode_,
reinterpret_cast<CudaT4*>(NM1.template MutableData<T4>()),

11
orttraining/orttraining/training_ops/cuda/optimizer/adam.cu
View file

@ -22,6 +22,7 @@ __global__ void _AdamOptimizer_mode0(
const T4 beta,
const T4 lambda,
const T4 epsilon,
const T4 max_norm,
const T4 alpha_correction,
const T4 beta_correction,
T4* moment_1_out,
@ -31,7 +32,7 @@ __global__ void _AdamOptimizer_mode0(
T_MIXED_PRECISION_FP* mixed_precision_weights_out,
CUDA_LONG N) {
CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);
const T4 actual_scale = _ComputeGradScale<T3, T_GRAD_NORM, T4>(loss_scale, grad_norm);
const T4 actual_scale = _ComputeGradScale<T3, T_GRAD_NORM, T4>(loss_scale, grad_norm, max_norm);
// Gradient scaling/clipping.
const T4 g = T4(grads[id]) / actual_scale;
@ -83,6 +84,7 @@ __global__ void _AdamOptimizer_mode1(
const T4 beta,
const T4 lambda,
const T4 epsilon,
const T4 max_norm,
const T4 alpha_correction,
const T4 beta_correction,
T4* moment_1_out,
@ -92,7 +94,7 @@ __global__ void _AdamOptimizer_mode1(
T_MIXED_PRECISION_FP* mixed_precision_weights_out,
CUDA_LONG N) {
CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);
const T4 actual_scale = _ComputeGradScale<T3, T_GRAD_NORM, T4>(loss_scale, grad_norm);
const T4 actual_scale = _ComputeGradScale<T3, T_GRAD_NORM, T4>(loss_scale, grad_norm, max_norm);
// Gradient scaling/clipping.
const T4 g = T4(grads[id]) / actual_scale;
@ -149,6 +151,7 @@ void AdamOptimizerImpl(
const T4 beta,
const T4 lambda,
const T4 epsilon,
const T4 max_norm,
const bool do_bias_correction,
const int64_t weight_decay_mode,
T4* moment_1_out,
@ -185,8 +188,10 @@ void AdamOptimizerImpl(
beta,
lambda,
epsilon,
max_norm,
alpha_correction,
beta_correction,
moment_1_out,
moment_2_out,
weights_out,
@ -207,6 +212,7 @@ void AdamOptimizerImpl(
beta,
lambda,
epsilon,
max_norm,
alpha_correction,
beta_correction,
moment_1_out,
@ -236,6 +242,7 @@ void AdamOptimizerImpl(
const T4 beta, \
const T4 lambda, \
const T4 epsilon, \
const T4 max_norm, \
const bool do_bias_correction, \
const int64_t weight_decay_mode, \
T4* moment_1_out, \

4
orttraining/orttraining/training_ops/cuda/optimizer/adam.h
View file

@ -22,6 +22,7 @@ void AdamOptimizerImpl(
const T4 beta,
const T4 lambda,
const T4 epsilon,
const T4 max_norm,
const bool do_bias_correction,
const int64_t weight_decay_mode,
T4* moment_1_out,
@ -39,10 +40,12 @@ class AdamOptimizer final : public CudaKernel {
info.GetAttrOrDefault("beta", &beta_, 0.999f);
info.GetAttrOrDefault("lambda", &lambda_, 0.0f);
info.GetAttrOrDefault("epsilon", &epsilon_, 1e-8f);
info.GetAttrOrDefault("max_norm_clip", &max_norm_clip_, 1.0f);
int64_t tmp_flag = static_cast<int64_t>(0);
ORT_ENFORCE(info.GetAttr<int64_t>("do_bias_correction", &tmp_flag).IsOK(), "Missing/Invalid do_bias_correction");
ORT_ENFORCE(tmp_flag == 0 || tmp_flag == 1, "do_bias_correction must be either 0 or 1.");
ORT_ENFORCE(max_norm_clip_ != 0, "max_norm_clip must NOT be 0.");
do_bias_correction_ = tmp_flag != 0 ? true : false;
info.GetAttrOrDefault("weight_decay_mode", &weight_decay_mode_, static_cast<int64_t>(0));
}
@ -54,6 +57,7 @@ class AdamOptimizer final : public CudaKernel {
float beta_;
float lambda_;
float epsilon_;
float max_norm_clip_;
bool do_bias_correction_;
int64_t weight_decay_mode_;
};

22
orttraining/orttraining/training_ops/cuda/optimizer/common.cuh
View file

@ -12,16 +12,20 @@ namespace cuda {
// _ComputeGradScale -- helper to calculate gradient scales based on global norms
// ---------------------------------------------------------------------------
template<typename TLossScale, typename TGradNorm, typename TFinalScale>
template <typename TLossScale, typename TGradNorm, typename TFinalScale>
__device__ __forceinline__ TFinalScale _ComputeGradScale(
const TLossScale* loss_scale,
const TGradNorm* g_norm) {
TFinalScale scale = loss_scale != nullptr ? TFinalScale(*loss_scale) : TFinalScale(1.f);
if (g_norm != nullptr && TFinalScale(*g_norm) > scale) {
const TFinalScale actual_g_norm = TFinalScale(*g_norm) / scale;
scale *= actual_g_norm;
}
return scale;
const TLossScale* loss_scale, // Scale of the gradient (called "scaled_g_norm" below)
const TGradNorm* scaled_g_norm, // Scaled gradient norm is an optimizer input
const TFinalScale max_g_norm) {
const TFinalScale scale = loss_scale != nullptr ? TFinalScale(*loss_scale) : TFinalScale(1.f);
const TFinalScale scaled_max_g_norm = TFinalScale(scale * max_g_norm);
// This number is used to divide the scaled gradient before applying optimizers.
TFinalScale scaled_g_scaling_factor = scale;
if (scaled_g_norm != nullptr && TFinalScale(*scaled_g_norm) > scaled_max_g_norm) {
scaled_g_scaling_factor = TFinalScale(*scaled_g_norm) / max_g_norm;
}
return scaled_g_scaling_factor;
}
} // namespace cuda
} // namespace onnxruntime

16
orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc
View file

@ -189,6 +189,7 @@ Status launch_lamb_compute_direction(
const std::vector<float>& betas,
const std::vector<float>& lambdas,
const std::vector<float>& epsilons,
const std::vector<float>& max_norms,
const int64_t do_bias_correction) {
ORT_ENFORCE(group_count == static_cast<int>(tensor_sizes.size()));
@ -209,8 +210,8 @@ Status launch_lamb_compute_direction(
const int max_tensor_size = compute_max_tensor_size_per_launch<tensor_count_per_group>(4);
// Bucketize tensor groups by the associated optimizer configuration.
// If two tensor groups use different "alpha", they should be put into two distinct buckets.
std::map<std::tuple<float, float, float, float>, std::vector<std::vector<void*>>> buckets;
std::map<std::tuple<float, float, float, float>, std::vector<int>> tensor_sizes_in_buckets;
std::map<std::tuple<float, float, float, float, float>, std::vector<std::vector<void*>>> buckets;
std::map<std::tuple<float, float, float, float, float>, std::vector<int>> tensor_sizes_in_buckets;
for (int i = 0; i < group_count; ++i) {
if (tensor_sizes[i] > max_tensor_size) {
// For the first iteration (indexed by 0), the update count should be 2.
@ -230,6 +231,7 @@ Status launch_lamb_compute_direction(
CudaT4(betas[i]),
CudaT2(lambdas[i]),
CudaT4(epsilons[i]),
CudaT2(max_norms[i]),
CudaT4(alpha_correction),
CudaT4(beta_correction),
p_ds[i],
@ -245,7 +247,7 @@ Status launch_lamb_compute_direction(
ptrs[4] = p_m1_news[i]; // new 1st momentum
ptrs[5] = p_m2_news[i]; // new 2nd momentum
auto key = std::make_tuple(alphas[i], betas[i], lambdas[i], epsilons[i]);
auto key = std::make_tuple(alphas[i], betas[i], lambdas[i], epsilons[i], max_norms[i]);
buckets[key].push_back(ptrs);
tensor_sizes_in_buckets[key].push_back(tensor_sizes[i]);
}
@ -253,8 +255,8 @@ Status launch_lamb_compute_direction(
for (auto& pair : buckets) {
const auto key = pair.first;
float alpha = 0.f, beta = 0.f, lambda = 0.f, epsilon = 0.f;
std::tie(alpha, beta, lambda, epsilon) = key;
float alpha = 0.f, beta = 0.f, lambda = 0.f, epsilon = 0.f, max_norm = 0.f;
std::tie(alpha, beta, lambda, epsilon, max_norm) = key;
// For the first iteration (indexed by 0), the update count should be 1.
const float alpha_correction =
@ -270,7 +272,7 @@ Status launch_lamb_compute_direction(
tensor_sizes_in_buckets[key],
buckets[key],
lamb_stage1,
p_loss_scale, p_g_norm, lambda, alpha, beta, epsilon, alpha_correction, beta_correction);
p_loss_scale, p_g_norm, lambda, alpha, beta, epsilon, CudaT2(max_norm), alpha_correction, beta_correction);
}
return Status::OK();
@ -658,7 +660,7 @@ Status LambOptimizer<T1, T2, T3, T4, T_GRAD_NORM, T_MIXED_PRECISION_FP>::Compute
p_ws, p_gs, p_m1s, p_m2s,
p_ds,
p_m1_news, p_m2_news,
alpha_, beta_, lambda_, epsilon_,
alpha_, beta_, lambda_, epsilon_, max_norm_clip_,
do_bias_correction_));
ORT_RETURN_IF_ERROR(launch_lamb_reduction(

18
orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu
View file

@ -31,16 +31,16 @@ __device__ __forceinline__ void _LambComputeDirectionRule(
T3& m2_new) {
// Actual gradient. The scale is a product of loss' scale and
// global gradient norm (if the norm > 1).
const T3 g_scaled = T3(T1(g) / g_scale);
const T3 g_unscaled = T3(T1(g) / g_scale);
// A constant in Lamb's equation.
const T3 one = T3(1.0f);
// Update exponentially-averaged historical gradient
const T3 m1_new_tmp = alpha * m1 + (one - alpha) * g_scaled;
const T3 m1_new_tmp = alpha * m1 + (one - alpha) * g_unscaled;
// Update exponentially-averaged historical squared gradient
const T3 m2_new_tmp = beta * m2 + (one - beta) * g_scaled * g_scaled;
const T3 m2_new_tmp = beta * m2 + (one - beta) * g_unscaled * g_unscaled;
// Compute unbiased 1st-order momentom.
// The value alpha_correction is usually (1-alpha^t),
@ -80,6 +80,7 @@ __global__ void _LambComputeDirectionImpl(
T3 beta,
T1 lambda,
T3 epsilon,
T1 max_norm,
T3 alpha_correction,
T3 beta_correction,
T2* update_direction,
@ -88,7 +89,7 @@ __global__ void _LambComputeDirectionImpl(
CUDA_LONG N) {
CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);
const T1 scale = _ComputeGradScale<T1, T_GRAD_NORM, T1>(loss_scale, g_norm);
const T1 scale = _ComputeGradScale<T1, T_GRAD_NORM, T1>(loss_scale, g_norm, max_norm);
_LambComputeDirectionRule(
scale,
@ -119,6 +120,7 @@ void LambComputeDirection(
T3 beta,
T1 lambda,
T3 epsilon,
T1 max_norm,
T3 alpha_correction,
T3 beta_correction,
T2* update_direction,
@ -139,6 +141,7 @@ void LambComputeDirection(
beta,
lambda,
epsilon,
max_norm,
alpha_correction,
beta_correction,
update_direction,
@ -159,6 +162,7 @@ void LambComputeDirection(
T3 beta, \
T1 lambda, \
T3 epsilon, \
T1 max_norm, \
T3 alpha_correction, \
T3 beta_correction, \
T2* weights_out, \
@ -315,6 +319,7 @@ __global__ void LambMultiTensorComputeDirectionImpl(
const T3 alpha,
const T3 beta,
const T3 epsilon,
const T1 max_norm,
const T3 alpha_correction,
const T3 beta_correction) {
const int group_index = chunk_group.block_index_to_tensor_group_index[blockIdx.x];
@ -327,7 +332,7 @@ __global__ void LambMultiTensorComputeDirectionImpl(
const T3* m2 = reinterpret_cast<const T3*>(chunk_group.tensor_ptrs[3][group_index]) + chunk_start;
T3* m1_new = reinterpret_cast<T3*>(chunk_group.tensor_ptrs[4][group_index]) + chunk_start;
T3* m2_new = reinterpret_cast<T3*>(chunk_group.tensor_ptrs[5][group_index]) + chunk_start;
const T1 scale = _ComputeGradScale<T1, T_GRAD_NORM, T1>(loss_scale, g_norm);
const T1 scale = _ComputeGradScale<T1, T_GRAD_NORM, T1>(loss_scale, g_norm, max_norm);
#pragma unroll
for (int i = threadIdx.x; i < chunk_size && i + chunk_start < tensor_size; i += blockDim.x) {
@ -358,6 +363,7 @@ void LambMultiTensorComputeDirectionFunctor<T1, T2, T3, T_GRAD_NORM>::operator()
const T3 alpha,
const T3 beta,
const T3 epsilon,
const T1 max_norm,
const T3 alpha_correction,
const T3 beta_correction) {
const int thread_count = ChunkGroup<6>::thread_count_per_block;
@ -371,6 +377,7 @@ void LambMultiTensorComputeDirectionFunctor<T1, T2, T3, T_GRAD_NORM>::operator()
alpha,
beta,
epsilon,
max_norm,
alpha_correction,
beta_correction);
}
@ -384,6 +391,7 @@ void LambMultiTensorComputeDirectionFunctor<T1, T2, T3, T_GRAD_NORM>::operator()
const T3 alpha, \
const T3 beta, \
const T3 epsilon, \
const T1 max_norm, \
const T3 alpha_correction, \
const T3 beta_correction);

7
orttraining/orttraining/training_ops/cuda/optimizer/lamb.h
View file

@ -17,8 +17,12 @@ class LambOptimizer final : public CudaKernel {
beta_ = info.GetAttrsOrDefault("beta", std::vector<float>(1024, 0.999f));
lambda_ = info.GetAttrsOrDefault("lambda", std::vector<float>(1024, 0.0f));
epsilon_ = info.GetAttrsOrDefault("epsilon", std::vector<float>(1024, 1e-6f));
max_norm_clip_ = info.GetAttrsOrDefault("max_norm_clip", std::vector<float>(1024, 1.0f));
ORT_ENFORCE(info.GetAttr<float>("ratio_min", &ratio_min_).IsOK(), "Missing/Invalid 'ratio_min' attribute value");
ORT_ENFORCE(info.GetAttr<float>("ratio_max", &ratio_max_).IsOK(), "Missing/Invalid 'ratio_max' attribute value");
for (const auto& max_norm : max_norm_clip_) {
ORT_ENFORCE(max_norm != 0, "max_norm_clip must NOT be 0.");
}
int64_t tmp_flag = static_cast<int64_t>(0);
ORT_ENFORCE(info.GetAttr<int64_t>("do_bias_correction", &tmp_flag).IsOK(), "Missing/Invalid do_bias_correction");
@ -33,6 +37,7 @@ class LambOptimizer final : public CudaKernel {
std::vector<float> beta_;
std::vector<float> lambda_;
std::vector<float> epsilon_;
std::vector<float> max_norm_clip_;
float ratio_min_;
float ratio_max_;
bool do_bias_correction_;
@ -54,6 +59,7 @@ void LambComputeDirection(
T3 beta,
T1 lambda,
T3 epsilon,
T1 max_norm,
T3 alpha_correction,
T3 beta_correction,
T2* update_direction,
@ -107,6 +113,7 @@ struct LambMultiTensorComputeDirectionFunctor {
const T3 alpha,
const T3 beta,
const T3 epsilon,
const T1 max_norm,
const T3 alpha_correction,
const T3 beta_correction);
};

51
orttraining/orttraining/training_ops/rocm/optimizer/adam.cc
View file

@ -11,30 +11,30 @@ namespace onnxruntime {
namespace rocm {
// TODO: Once Schema is checked in to onnx lets fix this to match that
#define REGISTER_ADAM_KERNEL_TYPED(T1, T2, T3, T4, T_GRAD, T_GRAD_NORM, T_MIXED_PRECISION_FP) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
AdamOptimizer, \
kMSDomain, \
1, \
T1##_##T2##_##T3##_##T4##_##T_GRAD##_##T_GRAD_NORM##_##T_MIXED_PRECISION_FP, \
kRocmExecutionProvider, \
KernelDefBuilder() \
.Alias(1, 0) /* Update step count in-place */ \
.Alias(2, 3) /* Update weights in-place */ \
.Alias(3, 4) /* Update gradients in-place */ \
.Alias(4, 1) /* Update moment-1 in-place */ \
.Alias(5, 2) /* Update moment-2 in-place */ \
.Alias(6, 5) /* Update mixed_precision weights in-place */ \
.InputMemoryType<OrtMemTypeCPUInput>(1) /* Keep step count in CPU */ \
.InputMemoryType<OrtMemTypeCPUInput>(9) /* Keep do_update in CPU */ \
.OutputMemoryType<OrtMemTypeCPUOutput>(0) /* Keep step count in CPU */ \
.TypeConstraint("T1", DataTypeImpl::GetTensorType<T1>()) \
.TypeConstraint("T2", DataTypeImpl::GetTensorType<T2>()) \
.TypeConstraint("T3", DataTypeImpl::GetTensorType<T3>()) \
.TypeConstraint("T4", DataTypeImpl::GetTensorType<T4>()) \
.TypeConstraint("T_GRAD", DataTypeImpl::GetTensorType<T_GRAD>()) \
.TypeConstraint("T_MIXED_PRECISION_FP", DataTypeImpl::GetTensorType<T_MIXED_PRECISION_FP>()) \
.TypeConstraint("T_GRAD_NORM", DataTypeImpl::GetTensorType<T_GRAD_NORM>()), \
#define REGISTER_ADAM_KERNEL_TYPED(T1, T2, T3, T4, T_GRAD, T_GRAD_NORM, T_MIXED_PRECISION_FP) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
AdamOptimizer, \
kMSDomain, \
1, \
T1##_##T2##_##T3##_##T4##_##T_GRAD##_##T_GRAD_NORM##_##T_MIXED_PRECISION_FP, \
kRocmExecutionProvider, \
KernelDefBuilder() \
.Alias(1, 0) /* Update step count in-place */ \
.Alias(2, 3) /* Update weights in-place */ \
.Alias(3, 4) /* Update gradients in-place */ \
.Alias(4, 1) /* Update moment-1 in-place */ \
.Alias(5, 2) /* Update moment-2 in-place */ \
.Alias(6, 5) /* Update mixed_precision weights in-place */ \
.InputMemoryType<OrtMemTypeCPUInput>(1) /* Keep step count in CPU */ \
.InputMemoryType<OrtMemTypeCPUInput>(9) /* Keep do_update in CPU */ \
.OutputMemoryType<OrtMemTypeCPUOutput>(0) /* Keep step count in CPU */ \
.TypeConstraint("T1", DataTypeImpl::GetTensorType<T1>()) \
.TypeConstraint("T2", DataTypeImpl::GetTensorType<T2>()) \
.TypeConstraint("T3", DataTypeImpl::GetTensorType<T3>()) \
.TypeConstraint("T4", DataTypeImpl::GetTensorType<T4>()) \
.TypeConstraint("T_GRAD", DataTypeImpl::GetTensorType<T_GRAD>()) \
.TypeConstraint("T_MIXED_PRECISION_FP", DataTypeImpl::GetTensorType<T_MIXED_PRECISION_FP>()) \
.TypeConstraint("T_GRAD_NORM", DataTypeImpl::GetTensorType<T_GRAD_NORM>()), \
AdamOptimizer<T1, T2, T3, T4, T_GRAD, T_GRAD_NORM, T_MIXED_PRECISION_FP>);
REGISTER_ADAM_KERNEL_TYPED(float, int64_t, float, float, float, float, MLFloat16)
@ -94,7 +94,7 @@ Status AdamOptimizer<T1, T2, T3, T4, T_GRAD, T_GRAD_NORM, T_MIXED_PRECISION_FP>:
const T2* S_in = S.template Data<T2>();
T2* S_out = NS.template MutableData<T2>();
const HipT_GRAD_NORM* G_norm = nullptr;
if (gradient_norm_tensor != nullptr) {
G_norm = reinterpret_cast<const HipT_GRAD_NORM*>(gradient_norm_tensor->template Data<T_GRAD_NORM>());
@ -136,6 +136,7 @@ Status AdamOptimizer<T1, T2, T3, T4, T_GRAD, T_GRAD_NORM, T_MIXED_PRECISION_FP>:
ToHipType<T4>::FromFloat(beta_),
ToHipType<T4>::FromFloat(lambda_),
ToHipType<T4>::FromFloat(epsilon_),
ToHipType<T4>::FromFloat(max_norm_clip_),
do_bias_correction_,
weight_decay_mode_,
reinterpret_cast<HipT4*>(NM1.template MutableData<T4>()),

10
orttraining/orttraining/training_ops/rocm/optimizer/adam.cu
View file

@ -22,6 +22,7 @@ __global__ void _AdamOptimizer_mode0(
const T4 beta,
const T4 lambda,
const T4 epsilon,
const T4 max_norm,
const T4 alpha_correction,
const T4 beta_correction,
T4* moment_1_out,
@ -31,7 +32,7 @@ __global__ void _AdamOptimizer_mode0(
T_MIXED_PRECISION_FP* mixed_precision_weights_out,
HIP_LONG N) {
CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);
const T4 actual_scale = _ComputeGradScale<T3, T_GRAD_NORM, T4>(loss_scale, grad_norm);
const T4 actual_scale = _ComputeGradScale<T3, T_GRAD_NORM, T4>(loss_scale, grad_norm, max_norm);
// Gradient scaling/clipping.
const T4 g = T4(grads[id]) / actual_scale;
@ -83,6 +84,7 @@ __global__ void _AdamOptimizer_mode1(
const T4 beta,
const T4 lambda,
const T4 epsilon,
const T4 max_norm,
const T4 alpha_correction,
const T4 beta_correction,
T4* moment_1_out,
@ -92,7 +94,7 @@ __global__ void _AdamOptimizer_mode1(
T_MIXED_PRECISION_FP* mixed_precision_weights_out,
HIP_LONG N) {
CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);
const T4 actual_scale = _ComputeGradScale<T3, T_GRAD_NORM, T4>(loss_scale, grad_norm);
const T4 actual_scale = _ComputeGradScale<T3, T_GRAD_NORM, T4>(loss_scale, grad_norm, max_norm);
// Gradient scaling/clipping.
const T4 g = T4(grads[id]) / actual_scale;
@ -149,6 +151,7 @@ void AdamOptimizerImpl(
const T4 beta,
const T4 lambda,
const T4 epsilon,
const T4 max_norm,
const bool do_bias_correction,
const int64_t weight_decay_mode,
T4* moment_1_out,
@ -185,6 +188,7 @@ void AdamOptimizerImpl(
beta,
lambda,
epsilon,
max_norm,
alpha_correction,
beta_correction,
moment_1_out,
@ -207,6 +211,7 @@ void AdamOptimizerImpl(
beta,
lambda,
epsilon,
max_norm,
alpha_correction,
beta_correction,
moment_1_out,
@ -236,6 +241,7 @@ void AdamOptimizerImpl(
const T4 beta, \
const T4 lambda, \
const T4 epsilon, \
const T4 max_norm, \
const bool do_bias_correction, \
const int64_t weight_decay_mode, \
T4* moment_1_out, \

17
orttraining/orttraining/training_ops/rocm/optimizer/lamb.cc
View file

@ -178,6 +178,7 @@ Status launch_lamb_compute_direction(
const std::vector<float>& betas,
const std::vector<float>& lambdas,
const std::vector<float>& epsilons,
const std::vector<float>& max_norms,
const int64_t do_bias_correction) {
ORT_ENFORCE(group_count == static_cast<int>(tensor_sizes.size()));
@ -198,8 +199,8 @@ Status launch_lamb_compute_direction(
const int max_tensor_size = compute_max_tensor_size_per_launch<tensor_count_per_group>(4);
// Bucketize tensor groups by the associated optimizer configuration.
// If two tensor groups use different "alpha", they should be put into two distinct buckets.
std::map<std::tuple<float, float, float, float>, std::vector<std::vector<void*>>> buckets;
std::map<std::tuple<float, float, float, float>, std::vector<int>> tensor_sizes_in_buckets;
std::map<std::tuple<float, float, float, float, float>, std::vector<std::vector<void*>>> buckets;
std::map<std::tuple<float, float, float, float, float>, std::vector<int>> tensor_sizes_in_buckets;
for (int i = 0; i < group_count; ++i) {
if (tensor_sizes[i] > max_tensor_size) {
// For the first iteration (indexed by 0), the update count should be 2.
@ -217,6 +218,7 @@ Status launch_lamb_compute_direction(
HipT4(betas[i]),
HipT2(lambdas[i]),
HipT4(epsilons[i]),
HipT2(max_norms[i]),
HipT4(alpha_correction),
HipT4(beta_correction),
p_ds[i],
@ -232,7 +234,7 @@ Status launch_lamb_compute_direction(
ptrs[4] = p_m1_news[i]; // new 1st momentum
ptrs[5] = p_m2_news[i]; // new 2nd momentum
auto key = std::make_tuple(alphas[i], betas[i], lambdas[i], epsilons[i]);
auto key = std::make_tuple(alphas[i], betas[i], lambdas[i], epsilons[i], max_norms[i]);
buckets[key].push_back(ptrs);
tensor_sizes_in_buckets[key].push_back(tensor_sizes[i]);
}
@ -240,8 +242,8 @@ Status launch_lamb_compute_direction(
for (auto& pair : buckets) {
const auto key = pair.first;
float alpha = 0.f, beta = 0.f, lambda = 0.f, epsilon = 0.f;
std::tie(alpha, beta, lambda, epsilon) = key;
float alpha = 0.f, beta = 0.f, lambda = 0.f, epsilon = 0.f, max_norm = 0.f;
std::tie(alpha, beta, lambda, epsilon, max_norm) = key;
// For the first iteration (indexed by 0), the update count should be 1.
const float alpha_correction =
@ -257,7 +259,7 @@ Status launch_lamb_compute_direction(
tensor_sizes_in_buckets[key],
buckets[key],
lamb_stage1,
p_loss_scale, p_g_norm, lambda, alpha, beta, epsilon, alpha_correction, beta_correction);
p_loss_scale, p_g_norm, lambda, alpha, beta, epsilon, HipT2(max_norm), alpha_correction, beta_correction);
}
return Status::OK();
@ -480,6 +482,7 @@ Status LambOptimizer<T1, T2, T3, T4, T_GRAD_NORM, T_MIXED_PRECISION_FP>::Compute
ORT_ENFORCE(beta_.size() >= static_cast<size_t>(group_count));
ORT_ENFORCE(lambda_.size() >= static_cast<size_t>(group_count));
ORT_ENFORCE(epsilon_.size() >= static_cast<size_t>(group_count));
ORT_ENFORCE(max_norm_clip_.size() >= static_cast<size_t>(group_count));
// If gradient norm is not finite, we copy inputs to outputs directly.
if (ctx->Input<Tensor>(0)) {
@ -647,7 +650,7 @@ Status LambOptimizer<T1, T2, T3, T4, T_GRAD_NORM, T_MIXED_PRECISION_FP>::Compute
p_ws, p_gs, p_m1s, p_m2s,
p_ds,
p_m1_news, p_m2_news,
alpha_, beta_, lambda_, epsilon_,
alpha_, beta_, lambda_, epsilon_, max_norm_clip_,
do_bias_correction_);
launch_lamb_reduction(

3
tools/ci_build/github/pai/pai-excluded-tests.txt
View file

@ -1,12 +1,15 @@
OptimizerTest.AdamOptimizerMixPrecisionTest
OptimizerTest.AdamOptimizerMixPrecision_FP16Weight_Test
OptimizerTest.AdamOptimizerMixPrecision_FP16Weight_SkipUpdate_Test
OptimizerTest.AdamOptimizerMixPrecision_FP16Weight_NoClipNorm_Test
OptimizerTest.AdamOptimizerMixPrecision_FP16Weight_ClipNorm_Test
OptimizerTest.AdamOptimizerMixPrecisionTestFloatEta
OptimizerTest.AdamOptimizerMixPrecisionTest_Gradient
OptimizerTest.LambOptimizerTest5DTensorMixPrecision32_16
OptimizerTest.LambOptimizerTestSimpleBaselineMixPrecision32_16
OptimizerTest.LambOptimizerTestBaselineMixPrecision32_16
OptimizerTest.LambOptimizerTestScalarMixPrecision32_16
OptimizerTest.LambOptimizerTestScalarMixPrecision32_16_NoDefaultMaxNormClipping
OptimizerTest.LambOptimizerTestLarge
CudaKernelTest.SoftmaxCrossEntropy_TinySizeTensor
CudaKernelTest.SoftmaxCrossEntropy_SmallSizeTensor