diff --git a/orttraining/orttraining/core/graph/adasum_optimizer_graph_builder.cc b/orttraining/orttraining/core/graph/adasum_optimizer_graph_builder.cc index b3add47d14..05aad1d166 100644 --- a/orttraining/orttraining/core/graph/adasum_optimizer_graph_builder.cc +++ b/orttraining/orttraining/core/graph/adasum_optimizer_graph_builder.cc @@ -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(adasum_reduction_type))}, + static_cast(adasum_reduction_type))}, "AdasumAllReduce")}); gradient_argdefs = std::move(adasum_output_argdefs); return Status::OK(); @@ -168,7 +168,6 @@ Status AdasumOptimizerGraphBuilder::BuildInternal( std::vector& gradient_argdefs, std::unordered_map>& weight_to_opt_mapping, OptimizerOutputKeyMap& 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 diff --git a/orttraining/orttraining/core/graph/optimizer/adam_optimizer_builder.h b/orttraining/orttraining/core/graph/optimizer/adam_optimizer_builder.h index d2fb363675..529d8c94dc 100644 --- a/orttraining/orttraining/core/graph/optimizer/adam_optimizer_builder.h +++ b/orttraining/orttraining/core/graph/optimizer/adam_optimizer_builder.h @@ -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>& weight_to_opt_mapping, std::vector& output_weight_argdefs, std::vector& output_gradient_argdefs) const override; - }; } // namespace training diff --git a/orttraining/orttraining/core/graph/optimizer/lamb_optimizer_builder.cc b/orttraining/orttraining/core/graph/optimizer/lamb_optimizer_builder.cc index 27fc9ebb8a..cb29e411b0 100644 --- a/orttraining/orttraining/core/graph/optimizer/lamb_optimizer_builder.cc +++ b/orttraining/orttraining/core/graph/optimizer/lamb_optimizer_builder.cc @@ -89,6 +89,7 @@ Status LambOptimizerBuilder::Build( std::vector beta; std::vector lambda; std::vector epsilon; + std::vector max_norm_clip; float ratio_min = -std::numeric_limits::infinity(); float ratio_max = std::numeric_limits::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)); diff --git a/orttraining/orttraining/core/graph/optimizer/lamb_optimizer_builder.h b/orttraining/orttraining/core/graph/optimizer/lamb_optimizer_builder.h index 38f3fd164a..144fa6e4d9 100644 --- a/orttraining/orttraining/core/graph/optimizer/lamb_optimizer_builder.h +++ b/orttraining/orttraining/core/graph/optimizer/lamb_optimizer_builder.h @@ -15,6 +15,7 @@ class LambOptimizerBuilder final : public OptimizerBuilder { "beta", "lambda", "epsilon", + "max_norm_clip", "ratio_min", "ratio_max", "do_bias_correction"}) {} diff --git a/orttraining/orttraining/core/graph/optimizer_builder.h b/orttraining/orttraining/core/graph/optimizer_builder.h index f578986df9..81399384f2 100644 --- a/orttraining/orttraining/core/graph/optimizer_builder.h +++ b/orttraining/orttraining/core/graph/optimizer_builder.h @@ -49,10 +49,10 @@ Status IsMatchingTypeAndShape( const int32_t element_type, const std::vector& expected_shape); - /** +/** * The configuration for optimizer builder. */ -struct OptimizerBuilderConfig{ +struct OptimizerBuilderConfig { //The ArgDefs of the weights to optimize. std::vector weight_argdefs; @@ -70,11 +70,11 @@ struct OptimizerBuilderConfig{ // The per weight optimizer configuration. std::vector 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 enable_grad_clipping; - // The initial state for optimizer params + // The initial state for optimizer params // shared by all weights. NameMLValMap shared_optimizer_states{}; }; diff --git a/orttraining/orttraining/core/graph/optimizer_config.h b/orttraining/orttraining/core/graph/optimizer_config.h index be48cfba5c..d6b31809c5 100644 --- a/orttraining/orttraining/core/graph/optimizer_config.h +++ b/orttraining/orttraining/core/graph/optimizer_config.h @@ -50,7 +50,7 @@ struct OptimizerNodeConfig { std::unordered_map attributes{}; std::unordered_map 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; + } } }; diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc index ceff4c2cd6..84bc92b2ba 100644 --- a/orttraining/orttraining/core/graph/training_op_defs.cc +++ b/orttraining/orttraining/core/graph/training_op_defs.cc @@ -172,6 +172,11 @@ OpSchema& RegisterLambOpSchema(OpSchema&& op_schema) { "Small scalar to avoid dividing by zero.", AttributeProto::FLOATS, std::vector(1024, 1e-6f)) + .Attr( + "max_norm_clip", + "clip threshold of gradients.", + AttributeProto::FLOATS, + std::vector(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(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(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(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(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); } diff --git a/orttraining/orttraining/core/session/training_session.cc b/orttraining/orttraining/core/session/training_session.cc index 515b84d857..2d37474527 100644 --- a/orttraining/orttraining/core/session/training_session.cc +++ b/orttraining/orttraining/core/session/training_session.cc @@ -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 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 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{""} : optional{}; 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{}; diff --git a/orttraining/orttraining/python/orttraining_pybind_state.cc b/orttraining/orttraining/python/orttraining_pybind_state.cc index 16e194ccf1..515be8dce1 100644 --- a/orttraining/orttraining/python/orttraining_pybind_state.cc +++ b/orttraining/orttraining/python/orttraining_pybind_state.cc @@ -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> ConvertORTTensorMapToNumpy(std::unordered_map c_tensor_state, const DataTransferManager& data_transfer_manager) { std::unordered_map> 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(); diff --git a/orttraining/orttraining/python/training/optim/config.py b/orttraining/orttraining/python/training/optim/config.py index 62349cd561..8e97cf8cc3 100644 --- a/orttraining/orttraining/python/training/optim/config.py +++ b/orttraining/orttraining/python/training/optim/config.py @@ -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 diff --git a/orttraining/orttraining/test/gradient/optimizer_ops_test.cc b/orttraining/orttraining/test/gradient/optimizer_ops_test.cc index fc7441e6ed..6888e5690a 100644 --- a/orttraining/orttraining/test/gradient/optimizer_ops_test.cc +++ b/orttraining/orttraining/test/gradient/optimizer_ops_test.cc @@ -9,6 +9,8 @@ namespace onnxruntime { namespace test { +namespace { + TEST(OptimizerTest, SGDOptimizerTest) { OpTester test("SGDOptimizer", 1, onnxruntime::kMSDomain); test.AddInput("ETA", {}, {0.5f}); @@ -131,7 +133,7 @@ TEST(OptimizerTest, AdamBiasCorrection) { test.AddInput("ETA", {}, {1.f}); test.AddInput("Update_Count", {}, {1}); - test.AddInput("W", {3}, {-0.4634f, 0.3584f, -0.2121f}); + test.AddInput("W", {3}, {-0.4634f, 0.3584f, -0.2121f}); test.AddInput("G", {3}, {0.4171f, 0.9485f, 1.2289f}); test.AddInput("Moment_1", {3}, {0.f, 0.f, 0.f}); test.AddInput("Moment_2", {3}, {0.f, 0.f, 0.f}); @@ -153,7 +155,7 @@ TEST(OptimizerTest, AdamWeightDecayMode0NoBiasCorrection) { test.AddInput("ETA", {}, {1.f}); test.AddInput("Update_Count", {}, {1}); - test.AddInput("W", {3}, {-0.4634f, 0.3584f, -0.2121f}); + test.AddInput("W", {3}, {-0.4634f, 0.3584f, -0.2121f}); test.AddInput("G", {3}, {0.4171f, 0.9485f, 1.2289f}); test.AddInput("Moment_1", {3}, {0.f, 0.f, 0.f}); test.AddInput("Moment_2", {3}, {0.f, 0.f, 0.f}); @@ -164,7 +166,6 @@ TEST(OptimizerTest, AdamWeightDecayMode0NoBiasCorrection) { test.AddOutput("W_Out", {3}, {-3.6210f, -2.8075f, -3.3723f}); test.AddOutput("G_Out", {3}, {-3.1576f, -3.1658f, -3.1601f}); - test.AddAttribute("do_bias_correction", static_cast(0)); test.AddAttribute("lambda", 0.01f); test.AddAttribute("weight_decay_mode", static_cast(0)); @@ -178,7 +179,7 @@ TEST(OptimizerTest, AdamWeightDecayMode0WithBiasCorrection) { test.AddInput("ETA", {}, {1.f}); test.AddInput("Update_Count", {}, {1}); - test.AddInput("W", {3}, {-0.4634f, 0.3584f, -0.2121f}); + test.AddInput("W", {3}, {-0.4634f, 0.3584f, -0.2121f}); test.AddInput("G", {3}, {0.4171f, 0.9485f, 1.2289f}); test.AddInput("Moment_1", {3}, {0.f, 0.f, 0.f}); test.AddInput("Moment_2", {3}, {0.f, 0.f, 0.f}); @@ -202,7 +203,7 @@ TEST(OptimizerTest, AdamWeightDecayMode1NoBiasCorrection) { test.AddInput("ETA", {}, {1.f}); test.AddInput("Update_Count", {}, {1}); - test.AddInput("W", {3}, {-0.4634f, 0.3584f, -0.2121f}); + test.AddInput("W", {3}, {-0.4634f, 0.3584f, -0.2121f}); test.AddInput("G", {3}, {0.4171f, 0.9485f, 1.2289f}); test.AddInput("Moment_1", {3}, {0.f, 0.f, 0.f}); test.AddInput("Moment_2", {3}, {0.f, 0.f, 0.f}); @@ -225,7 +226,7 @@ TEST(OptimizerTest, AdamWeightDecayMode1WithBiasCorrection) { test.AddInput("ETA", {}, {1.f}); test.AddInput("Update_Count", {}, {1}); - test.AddInput("W", {3}, {-0.4634f, 0.3584f, -0.2121f}); + test.AddInput("W", {3}, {-0.4634f, 0.3584f, -0.2121f}); test.AddInput("G", {3}, {0.4171f, 0.9485f, 1.2289f}); test.AddInput("Moment_1", {3}, {0.f, 0.f, 0.f}); test.AddInput("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& 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("ETA", {}, data.eta_half); + test.AddInput("Update_Count", {}, {3}); + test.AddInput("W", {3}, data.w); + test.AddInput("G", {3}, data.g_half); + test.AddInput("Moment_1", {3}, data.m1_half); + test.AddInput("Moment_2", {3}, data.m2_half); + test.AddInput("FP16_W", {3}, data.w_half); + test.AddInput("loss_scale", {1}, {1.0f}); + // grad clipping should not take effect because default max_norm is 1.0f + test.AddInput("grad_norm", {1}, {0.01f}); + // Verify AdamOptimizer outputs + test.AddOutput("Update_Count_Out", {}, {4}); + test.AddOutput("Moment_1_Out", {3}, data.m1_new_half); + test.AddOutput("Moment_2_Out", {3}, data.m2_new_half); + test.AddOutput("W_Out", {3}, data.w_new); + test.AddMissingOptionalOutput(); + test.AddOutput("FP16_W_Out", {3}, data.w_new_half); + + test.AddAttribute("do_bias_correction", static_cast(0)); + test.AddAttribute("weight_decay_mode", static_cast(0)); + test.Run(); +} + +TEST(OptimizerTest, AdamOptimizerMixPrecision_FP16Weight_ClipNorm_Test) { + OpTester test("AdamOptimizer", 1, onnxruntime::kMSDomain); + AdamOptimizerInputOutput data; + + // Expected FP32 Outputs + std::vector m1_new = {0.13f, 0.23f, 0.33f}; + std::vector m2_new = {0.3997f, 0.4998f, 0.6001f}; + std::vector w_new = {0.8972168f, 1.8369141f, 2.7871094f}; + // FP16 Outputs + std::vector m1_new_half; + std::vector m2_new_half; + std::vector 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("ETA", {}, data.eta_half); + test.AddInput("Update_Count", {}, {3}); + test.AddInput("W", {3}, data.w); + test.AddInput("G", {3}, data.g_half); + test.AddInput("Moment_1", {3}, data.m1_half); + test.AddInput("Moment_2", {3}, data.m2_half); + test.AddInput("FP16_W", {3}, data.w_half); + test.AddInput("loss_scale", {1}, {1.0f}); + test.AddInput("grad_norm", {1}, {0.01f}); + // Verify AdamOptimizer outputs + test.AddOutput("Update_Count_Out", {}, {4}); + test.AddOutput("Moment_1_Out", {3}, m1_new_half); + test.AddOutput("Moment_2_Out", {3}, m2_new_half); + test.AddOutput("W_Out", {3}, w_new); + test.AddMissingOptionalOutput(); + test.AddOutput("FP16_W_Out", {3}, w_new_half); + + test.AddAttribute("do_bias_correction", static_cast(0)); + test.AddAttribute("weight_decay_mode", static_cast(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& m, /* 2nd-order momentum */ const std::vector& 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& w_new, /* updated gradients */ std::vector& g_new, /* updated momentum */ std::vector& m_new, /* updated 2nd-order momentum */ std::vector& 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::infinity(), const float ratio_max = std::numeric_limits::infinity()) { // Element counts of all vector-typed arguments. @@ -400,9 +482,12 @@ void compute_lamb( // Buffer to store update direction. std::vector 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(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& w_new, const std::vector& g_new, const std::vector& m_new, @@ -470,13 +556,19 @@ void run_lamb_test_with_baseline( const std::vector& 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::infinity(), const float ratio_max = std::numeric_limits::infinity()) { OpTester test("LambOptimizer", 1, onnxruntime::kMSDomain, true); test.AddInput("update_signal", {1}, {do_update}); - test.AddMissingOptionalInput(); - test.AddMissingOptionalInput(); + test.AddInput("loss_scale", {}, {loss_scale}); + if (p_g_norm == nullptr) { + test.AddMissingOptionalInput(); + } else { + test.AddInput("gradient_norm", {}, {T2(*p_g_norm)}); + } test.AddInput("ETA", {1}, eta); if (step > 0) { test.AddInput("Step", {}, {step}); @@ -497,6 +589,7 @@ void run_lamb_test_with_baseline( test.AddAttribute("beta", std::vector(1, beta)); test.AddAttribute("lambda", std::vector(1, lambda)); test.AddAttribute("epsilon", std::vector(1, epsilon)); + test.AddAttribute("max_norm_clip", std::vector(1, max_norm)); test.AddAttribute("ratio_min", ratio_min); test.AddAttribute("ratio_max", ratio_max); @@ -530,8 +623,6 @@ template void run_multi_tensor_lamb_test_with_baseline( const std::vector>& shapes, const T1 eta, - const T1 loss_scale, - const T1 g_norm, const std::vector>& ws, const std::vector>& gs, const std::vector>& ms, @@ -540,6 +631,7 @@ void run_multi_tensor_lamb_test_with_baseline( const std::vector& betas, const std::vector& lambdas, const std::vector& epsilons, + const std::vector& max_norms, const std::vector>& w_news, const std::vector>& g_news, const std::vector>& m_news, @@ -548,6 +640,8 @@ void run_multi_tensor_lamb_test_with_baseline( const std::vector>& 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::infinity(), const float ratio_max = std::numeric_limits::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(ws.size()); test.AddInput("update_signal", {}, {do_update}); - test.AddInput("loss_scale", {}, {loss_scale}); - test.AddInput("gradient_norm", {}, {g_norm}); + test.AddInput("loss_scale", {}, {loss_scale}); + if (p_g_norm == nullptr) { + test.AddMissingOptionalInput(); + } else { + test.AddInput("gradient_norm", {}, {T2(*p_g_norm)}); + } test.AddInput("ETA", {}, {eta}); if (step > 0) { test.AddInput("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> shapes, const float eta, - const float loss_scale, - const float g_norm, const std::vector> ws, const std::vector> gs, const std::vector> ms, @@ -655,7 +753,10 @@ void run_multi_tensor_lamb_test( const std::vector alphas, const std::vector betas, const std::vector epsilons, + const std::vector 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::infinity(), const float ratio_max = std::numeric_limits::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(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 w_new(w.size(), 0); std::vector g_new(g.size(), 0); std::vector 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 eta_half(eta.size()); std::vector 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::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::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 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 shape = {(int64_t)1}; + const float eta = 0.5f; + const std::vector w = {1.0f}; + const std::vector g = {3.0f}; + const std::vector m = {-10.0f}; + const std::vector 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 m_new = {0.0f}; + std::vector v_new = {0.0f}; + std::vector 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 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 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 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 shape = {1}; + const std::vector eta = {0.1f}; + const std::vector w = {-1.5f}; + const std::vector g = {-0.75f}; + const std::vector m = {0.87f}; + const std::vector 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 betas(group_count); std::vector lambdas(group_count); std::vector epsilons(group_count); + std::vector 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 diff --git a/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py b/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py index 463bf9fb62..6af5a240d2 100644 --- a/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py +++ b/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py @@ -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) \ No newline at end of file diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/adam.cc b/orttraining/orttraining/training_ops/cuda/optimizer/adam.cc index c169955523..3045a549b3 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/adam.cc +++ b/orttraining/orttraining/training_ops/cuda/optimizer/adam.cc @@ -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(1) /* Keep step count in CPU */ \ - .InputMemoryType(9) /* Keep do_update in CPU */ \ - .OutputMemoryType(0) /* Keep step count in CPU */ \ - .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T3", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T4", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T_GRAD", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T_MIXED_PRECISION_FP", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T_GRAD_NORM", DataTypeImpl::GetTensorType()), \ +#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(1) /* Keep step count in CPU */ \ + .InputMemoryType(9) /* Keep do_update in CPU */ \ + .OutputMemoryType(0) /* Keep step count in CPU */ \ + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T3", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T4", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T_GRAD", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T_MIXED_PRECISION_FP", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T_GRAD_NORM", DataTypeImpl::GetTensorType()), \ AdamOptimizer); REGISTER_ADAM_KERNEL_TYPED(float, int64_t, float, float, float, float, MLFloat16) @@ -106,7 +106,7 @@ Status AdamOptimizer: const T2* S_in = S.template Data(); T2* S_out = NS.template MutableData(); - + const CudaT_GRAD_NORM* G_norm = nullptr; if (gradient_norm_tensor != nullptr) { G_norm = reinterpret_cast(gradient_norm_tensor->template Data()); @@ -148,6 +148,7 @@ Status AdamOptimizer: ToCudaType::FromFloat(beta_), ToCudaType::FromFloat(lambda_), ToCudaType::FromFloat(epsilon_), + ToCudaType::FromFloat(max_norm_clip_), do_bias_correction_, weight_decay_mode_, reinterpret_cast(NM1.template MutableData()), diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/adam.cu b/orttraining/orttraining/training_ops/cuda/optimizer/adam.cu index bf26206e96..d892cd446d 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/adam.cu +++ b/orttraining/orttraining/training_ops/cuda/optimizer/adam.cu @@ -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(loss_scale, grad_norm); + const T4 actual_scale = _ComputeGradScale(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(loss_scale, grad_norm); + const T4 actual_scale = _ComputeGradScale(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, \ diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/adam.h b/orttraining/orttraining/training_ops/cuda/optimizer/adam.h index 9326966749..f979056e38 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/adam.h +++ b/orttraining/orttraining/training_ops/cuda/optimizer/adam.h @@ -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(0); ORT_ENFORCE(info.GetAttr("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(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_; }; diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/common.cuh b/orttraining/orttraining/training_ops/cuda/optimizer/common.cuh index 4305e9cec7..ea21edb6fa 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/common.cuh +++ b/orttraining/orttraining/training_ops/cuda/optimizer/common.cuh @@ -12,16 +12,20 @@ namespace cuda { // _ComputeGradScale -- helper to calculate gradient scales based on global norms // --------------------------------------------------------------------------- -template +template __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 \ No newline at end of file diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc index 0c37031fa9..e27903e89d 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc +++ b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc @@ -189,6 +189,7 @@ Status launch_lamb_compute_direction( const std::vector& betas, const std::vector& lambdas, const std::vector& epsilons, + const std::vector& max_norms, const int64_t do_bias_correction) { ORT_ENFORCE(group_count == static_cast(tensor_sizes.size())); @@ -209,8 +210,8 @@ Status launch_lamb_compute_direction( const int max_tensor_size = compute_max_tensor_size_per_launch(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::vector>> buckets; - std::map, std::vector> tensor_sizes_in_buckets; + std::map, std::vector>> buckets; + std::map, std::vector> 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::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( diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu index 0e89da0129..b8c8171509 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu +++ b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu @@ -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(loss_scale, g_norm); + const T1 scale = _ComputeGradScale(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(chunk_group.tensor_ptrs[3][group_index]) + chunk_start; T3* m1_new = reinterpret_cast(chunk_group.tensor_ptrs[4][group_index]) + chunk_start; T3* m2_new = reinterpret_cast(chunk_group.tensor_ptrs[5][group_index]) + chunk_start; - const T1 scale = _ComputeGradScale(loss_scale, g_norm); + const T1 scale = _ComputeGradScale(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::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::operator() alpha, beta, epsilon, + max_norm, alpha_correction, beta_correction); } @@ -384,6 +391,7 @@ void LambMultiTensorComputeDirectionFunctor::operator() const T3 alpha, \ const T3 beta, \ const T3 epsilon, \ + const T1 max_norm, \ const T3 alpha_correction, \ const T3 beta_correction); diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.h b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.h index 1b136ec07c..7882a94759 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.h +++ b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.h @@ -17,8 +17,12 @@ class LambOptimizer final : public CudaKernel { beta_ = info.GetAttrsOrDefault("beta", std::vector(1024, 0.999f)); lambda_ = info.GetAttrsOrDefault("lambda", std::vector(1024, 0.0f)); epsilon_ = info.GetAttrsOrDefault("epsilon", std::vector(1024, 1e-6f)); + max_norm_clip_ = info.GetAttrsOrDefault("max_norm_clip", std::vector(1024, 1.0f)); ORT_ENFORCE(info.GetAttr("ratio_min", &ratio_min_).IsOK(), "Missing/Invalid 'ratio_min' attribute value"); ORT_ENFORCE(info.GetAttr("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(0); ORT_ENFORCE(info.GetAttr("do_bias_correction", &tmp_flag).IsOK(), "Missing/Invalid do_bias_correction"); @@ -33,6 +37,7 @@ class LambOptimizer final : public CudaKernel { std::vector beta_; std::vector lambda_; std::vector epsilon_; + std::vector 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); }; diff --git a/orttraining/orttraining/training_ops/rocm/optimizer/adam.cc b/orttraining/orttraining/training_ops/rocm/optimizer/adam.cc index 3387da83a9..2d378bff5b 100644 --- a/orttraining/orttraining/training_ops/rocm/optimizer/adam.cc +++ b/orttraining/orttraining/training_ops/rocm/optimizer/adam.cc @@ -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(1) /* Keep step count in CPU */ \ - .InputMemoryType(9) /* Keep do_update in CPU */ \ - .OutputMemoryType(0) /* Keep step count in CPU */ \ - .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T3", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T4", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T_GRAD", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T_MIXED_PRECISION_FP", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("T_GRAD_NORM", DataTypeImpl::GetTensorType()), \ +#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(1) /* Keep step count in CPU */ \ + .InputMemoryType(9) /* Keep do_update in CPU */ \ + .OutputMemoryType(0) /* Keep step count in CPU */ \ + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T3", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T4", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T_GRAD", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T_MIXED_PRECISION_FP", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T_GRAD_NORM", DataTypeImpl::GetTensorType()), \ AdamOptimizer); REGISTER_ADAM_KERNEL_TYPED(float, int64_t, float, float, float, float, MLFloat16) @@ -94,7 +94,7 @@ Status AdamOptimizer: const T2* S_in = S.template Data(); T2* S_out = NS.template MutableData(); - + const HipT_GRAD_NORM* G_norm = nullptr; if (gradient_norm_tensor != nullptr) { G_norm = reinterpret_cast(gradient_norm_tensor->template Data()); @@ -136,6 +136,7 @@ Status AdamOptimizer: ToHipType::FromFloat(beta_), ToHipType::FromFloat(lambda_), ToHipType::FromFloat(epsilon_), + ToHipType::FromFloat(max_norm_clip_), do_bias_correction_, weight_decay_mode_, reinterpret_cast(NM1.template MutableData()), diff --git a/orttraining/orttraining/training_ops/rocm/optimizer/adam.cu b/orttraining/orttraining/training_ops/rocm/optimizer/adam.cu index 91e27f2ed7..cfda9ddbca 100644 --- a/orttraining/orttraining/training_ops/rocm/optimizer/adam.cu +++ b/orttraining/orttraining/training_ops/rocm/optimizer/adam.cu @@ -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(loss_scale, grad_norm); + const T4 actual_scale = _ComputeGradScale(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(loss_scale, grad_norm); + const T4 actual_scale = _ComputeGradScale(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, \ diff --git a/orttraining/orttraining/training_ops/rocm/optimizer/lamb.cc b/orttraining/orttraining/training_ops/rocm/optimizer/lamb.cc index 49bfccaa51..2c51b6ec9c 100644 --- a/orttraining/orttraining/training_ops/rocm/optimizer/lamb.cc +++ b/orttraining/orttraining/training_ops/rocm/optimizer/lamb.cc @@ -178,6 +178,7 @@ Status launch_lamb_compute_direction( const std::vector& betas, const std::vector& lambdas, const std::vector& epsilons, + const std::vector& max_norms, const int64_t do_bias_correction) { ORT_ENFORCE(group_count == static_cast(tensor_sizes.size())); @@ -198,8 +199,8 @@ Status launch_lamb_compute_direction( const int max_tensor_size = compute_max_tensor_size_per_launch(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::vector>> buckets; - std::map, std::vector> tensor_sizes_in_buckets; + std::map, std::vector>> buckets; + std::map, std::vector> 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::Compute ORT_ENFORCE(beta_.size() >= static_cast(group_count)); ORT_ENFORCE(lambda_.size() >= static_cast(group_count)); ORT_ENFORCE(epsilon_.size() >= static_cast(group_count)); + ORT_ENFORCE(max_norm_clip_.size() >= static_cast(group_count)); // If gradient norm is not finite, we copy inputs to outputs directly. if (ctx->Input(0)) { @@ -647,7 +650,7 @@ Status LambOptimizer::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( diff --git a/tools/ci_build/github/pai/pai-excluded-tests.txt b/tools/ci_build/github/pai/pai-excluded-tests.txt index 33cd15ac51..6ed3d24ce7 100644 --- a/tools/ci_build/github/pai/pai-excluded-tests.txt +++ b/tools/ci_build/github/pai/pai-excluded-tests.txt @@ -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