From 44f7b1bf2ced2331d81e3801922881acd899d3b0 Mon Sep 17 00:00:00 2001 From: pengwa Date: Fri, 27 May 2022 19:52:04 +0800 Subject: [PATCH] MTA AdamWOptimizer (#11506) * skeleton change * adam compute kernels * add rtol/atol for tests * some clean up * optional outputs * more clean up * add tests * adamw mode=1 test pass * clean up tests * add HF AdamW test cases * refactor adam test file * make test pass * all test pass, fix comments * rename to adamw * make test pass again * fix cpplint * minor fixes * fix python lint * Fix build and tests * fix builds * fix windows build * fix win build * minor fix * Refine based on comments * resolve comments * formatting * resolve comments * add ut --- .../providers/cuda/multi_tensor/common.cuh | 3 +- .../shared_library/provider_interfaces.h | 18 +- .../shared_library/provider_wrappedtypes.h | 2 +- .../core/session/provider_bridge_ort.cc | 2 +- .../sequence/concat_from_sequence_op_test.cc | 6 + .../test/providers/provider_test_utils.h | 55 +- .../adamw_test/adamw_test_data_generator.py | 192 +++ .../adamw_test_multiple_weights_mode_0.json | 1117 +++++++++++++++++ .../adamw_test_multiple_weights_mode_1.json | 1117 +++++++++++++++++ .../adamw_test_single_weight_mode_0.json | 362 ++++++ .../adamw_test_single_weight_mode_1.json | 362 ++++++ .../core/graph/training_op_defs.cc | 159 ++- .../training_ops/cuda/optimizer/adamw_test.cc | 321 +++++ .../training_ops/cuda/optimizer/common.cc | 243 ++++ .../test/training_ops/cuda/optimizer/common.h | 116 ++ .../training_ops/cpu/optimizer/common.h | 3 +- .../cuda/cuda_training_kernels.cc | 2 + .../training_ops/cuda/optimizer/adam_impl.cu | 1 + .../cuda/optimizer/adamw/adamw.cc | 158 +++ .../training_ops/cuda/optimizer/adamw/adamw.h | 45 + .../cuda/optimizer/adamw/adamw_impl.cu | 222 ++++ .../cuda/optimizer/adamw/adamw_impl.h | 29 + .../training_ops/cuda/optimizer/lamb_impl.cu | 32 +- 23 files changed, 4519 insertions(+), 48 deletions(-) create mode 100644 onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_data_generator.py create mode 100644 onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_multiple_weights_mode_0.json create mode 100644 onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_multiple_weights_mode_1.json create mode 100644 onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_single_weight_mode_0.json create mode 100644 onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_single_weight_mode_1.json create mode 100644 orttraining/orttraining/test/training_ops/cuda/optimizer/adamw_test.cc create mode 100644 orttraining/orttraining/test/training_ops/cuda/optimizer/common.cc create mode 100644 orttraining/orttraining/test/training_ops/cuda/optimizer/common.h create mode 100644 orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw.cc create mode 100644 orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw.h create mode 100644 orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw_impl.cu create mode 100644 orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw_impl.h diff --git a/onnxruntime/core/providers/cuda/multi_tensor/common.cuh b/onnxruntime/core/providers/cuda/multi_tensor/common.cuh index 508cf56297..79cc6b7e1b 100644 --- a/onnxruntime/core/providers/cuda/multi_tensor/common.cuh +++ b/onnxruntime/core/providers/cuda/multi_tensor/common.cuh @@ -81,6 +81,7 @@ void launch_multi_tensor_functor( std::vector>& grouped_tensor_pointers, TMultiTensorFunctor multipleTensorKernel, TFunctorParams&&... kernelParams) { + // Check if 32-bit integer is enough. ORT_ENFORCE(tensor_sizes.size() > 0); ORT_ENFORCE(tensor_sizes.size() < static_cast(INT_MAX)); ORT_ENFORCE(grouped_tensor_pointers.size() > 0); @@ -93,8 +94,6 @@ void launch_multi_tensor_functor( int tensor_group_index = 0; int block_index = 0; - // Check if 32-bit integer is enough. - ORT_ENFORCE(tensor_sizes.size() < static_cast(INT_MAX)); ORT_ENFORCE(grouped_tensor_pointers.size() == tensor_sizes.size()); ORT_ENFORCE(group_size == ACTUAL_TENSOR_GROUP_SIZE[TensorGroupSize]); for (int i = 0; i < group_count; ++i) { diff --git a/onnxruntime/core/providers/shared_library/provider_interfaces.h b/onnxruntime/core/providers/shared_library/provider_interfaces.h index 60662637e2..ccb0a078b5 100644 --- a/onnxruntime/core/providers/shared_library/provider_interfaces.h +++ b/onnxruntime/core/providers/shared_library/provider_interfaces.h @@ -109,14 +109,16 @@ struct Node__EdgeIterator { virtual int GetDstArgIndex() const = 0; }; -// There are two ways to route a function, one is a virtual method and the other is a function pointer (or pointer to member function) -// The function pointers are nicer in that they directly call the target function, but they cannot be used in cases where we're calling -// a specific implementation of a virtual class member. Trying to get a pointer to member of a virtual function will return a thunk that -// calls the virtual function (which will lead to infinite recursion in the bridge). There is no known way to get the non virtual member -// function pointer implementation in this case. -//The suppressed warning is: "The type with a virtual function needs either public virtual or protected nonvirtual destructor." -//However, we do not allocate this type on heap. -//Please do not new or delete this type(and subtypes). +// There are two ways to route a function, one is a virtual method and the other is a function pointer (or pointer to +// member function). +// The function pointers are nicer in that they directly call the target function, but they cannot be used in cases +// where we're calling a specific implementation of a virtual class member. Trying to get a pointer to member of a +// virtual function will return a thunk that calls the virtual function (which will lead to infinite recursion in the +// bridge). There is no known way to get the non virtual member function pointer implementation in this case. +// The suppressed warning is: +// "The type with a virtual function needs either public virtual or protected nonvirtual destructor." +// However, we do not allocate this type on heap. +// Please do not new or delete this type(and subtypes). #if defined(_MSC_VER) && !defined(__clang__) #pragma warning(push) #pragma warning(disable : 26436) diff --git a/onnxruntime/core/providers/shared_library/provider_wrappedtypes.h b/onnxruntime/core/providers/shared_library/provider_wrappedtypes.h index cf1266b565..e4792cbea9 100644 --- a/onnxruntime/core/providers/shared_library/provider_wrappedtypes.h +++ b/onnxruntime/core/providers/shared_library/provider_wrappedtypes.h @@ -1026,7 +1026,7 @@ struct SparseTensor final { }; #endif -//TensorSeq +// TensorSeq struct TensorSeq final { MLDataType DataType() const noexcept { return g_host->TensorSeq__DataType(this); } void SetType(MLDataType elem_type) { g_host->TensorSeq__SetType(this, elem_type); } diff --git a/onnxruntime/core/session/provider_bridge_ort.cc b/onnxruntime/core/session/provider_bridge_ort.cc index 704a5a7f55..09f2943271 100644 --- a/onnxruntime/core/session/provider_bridge_ort.cc +++ b/onnxruntime/core/session/provider_bridge_ort.cc @@ -337,7 +337,7 @@ struct ProviderHostImpl : ProviderHost { ONNX_NAMESPACE::TensorShapeProto* TypeProto_Tensor__mutable_shape(ONNX_NAMESPACE::TypeProto_Tensor* p) override { return p->mutable_shape(); } int32_t TypeProto_Tensor__elem_type(const ONNX_NAMESPACE::TypeProto_Tensor* p) override { return p->elem_type(); } - //TypeProto_SparseTensor (wrapped) + // TypeProto_SparseTensor (wrapped) #if !defined(DISABLE_SPARSE_TENSORS) bool TypeProto_SparseTensor__has_shape(const ONNX_NAMESPACE::TypeProto_SparseTensor* p) override { return p->has_shape(); } const ONNX_NAMESPACE::TensorShapeProto& TypeProto_SparseTensor__shape(const ONNX_NAMESPACE::TypeProto_SparseTensor* p) override { diff --git a/onnxruntime/test/providers/cpu/sequence/concat_from_sequence_op_test.cc b/onnxruntime/test/providers/cpu/sequence/concat_from_sequence_op_test.cc index cca69effd5..dd5393f8cb 100644 --- a/onnxruntime/test/providers/cpu/sequence/concat_from_sequence_op_test.cc +++ b/onnxruntime/test/providers/cpu/sequence/concat_from_sequence_op_test.cc @@ -95,6 +95,9 @@ TEST(SequenceOpsTest, ConcatFromSequence_Concat_Axis1) { TEST(SequenceOpsTest, ConcatFromSequence_Concat_Axis2) { OpTester test("ConcatFromSequence", 11); + // Don't provide input shapes so it doesn't fail during shape inferencing, + // to test the expected failure in compute. + test.AddShapeToTensorData(false, -1); test.AddAttribute("axis", 2); test.AddAttribute("new_axis", 0); // concat mode SeqTensors input; @@ -120,6 +123,9 @@ TEST(SequenceOpsTest, ConcatFromSequence_Concat_Axis1_WithEmptyInput) { TEST(SequenceOpsTest, ConcatFromSequence_Concat_ScalarInputs) { OpTester test("ConcatFromSequence", 11); + // Don't provide input shapes so it doesn't fail during shape inferencing, + // to test the expected failure in compute. + test.AddShapeToTensorData(false, -1); test.AddAttribute("axis", 0); test.AddAttribute("new_axis", 0); // concat mode SeqTensors input; diff --git a/onnxruntime/test/providers/provider_test_utils.h b/onnxruntime/test/providers/provider_test_utils.h index ea1316dfaf..6d50bc2db6 100644 --- a/onnxruntime/test/providers/provider_test_utils.h +++ b/onnxruntime/test/providers/provider_test_utils.h @@ -3,7 +3,9 @@ #pragma once +#include #include + #include "gmock/gmock.h" #include "gtest/gtest.h" #include @@ -203,7 +205,7 @@ class OpTester { // Set whether the NodeArg created by AddInput/AddOutput should include shape information // for Tensor types. If not added, shape inferencing should resolve. If added, shape inferencing - // should validate. Default is to not add. + // should validate. Default is to add. // Additionally a symbolic dimension will be added if symbolic_dim matches a dimension in the input. OpTester& AddShapeToTensorData(bool add_shape = true, int symbolic_dim = -1) { add_shape_to_tensor_data_ = add_shape; @@ -400,8 +402,9 @@ class OpTester { } template - void AddSeqOutput(const char* name, const SeqTensors& seq_tensors) { - AddSeqData(output_data_, name, &seq_tensors); + void AddSeqOutput(const char* name, const SeqTensors& seq_tensors, + float rel_error = 0.0f, float abs_error = 0.0f) { + AddSeqData(output_data_, name, &seq_tensors, false, rel_error, abs_error); } #if !defined(DISABLE_OPTIONAL_TYPE) @@ -450,8 +453,9 @@ class OpTester { template void AddOptionalTypeSeqOutput(const char* name, - const SeqTensors* seq_tensors) { - AddSeqData(output_data_, name, seq_tensors, true); + const SeqTensors* seq_tensors, + float rel_error = 0.0f, float abs_error = 0.0f) { + AddSeqData(output_data_, name, seq_tensors, true, rel_error, abs_error); } #endif @@ -945,7 +949,8 @@ class OpTester { template void AddSeqData(std::vector& data, const char* name, const SeqTensors* seq_tensors, - bool is_optional_sequence_tensor_type = false) { + bool is_optional_sequence_tensor_type = false, + float rel_error = 0.0f, float abs_error = 0.0f) { #if defined(DISABLE_OPTIONAL_TYPE) if (is_optional_sequence_tensor_type) { ORT_THROW("Optional type is not supported in this build"); @@ -953,12 +958,14 @@ class OpTester { #endif std::unique_ptr ptr; + SequenceTensorTypeProto sequence_tensor_proto; if (seq_tensors) { auto num_tensors = seq_tensors->tensors.size(); std::vector tensors; tensors.resize(num_tensors); auto elem_type = DataTypeImpl::GetType(); + for (size_t i = 0; i < num_tensors; ++i) { TensorShape shape{seq_tensors->tensors[i].shape}; auto values_count = static_cast(seq_tensors->tensors[i].data.size()); @@ -976,6 +983,28 @@ class OpTester { for (int64_t x = 0; x < values_count; ++x) { data_ptr[x] = seq_tensors->tensors[i].data[x]; } + + if (add_shape_to_tensor_data_) { + auto* output_tensor_type = sequence_tensor_proto.proto.mutable_sequence_type() + ->mutable_elem_type() + ->mutable_tensor_type(); + if (i == 0) { + ONNX_NAMESPACE::TensorShapeProto* seq_input_shape = output_tensor_type->mutable_shape(); + output_tensor_type->set_elem_type(utils::ToTensorProtoElementType()); + for (size_t j = 0; j < shape.NumDimensions(); ++j) { + auto dim = seq_input_shape->add_dim(); + dim->set_dim_value(shape[j]); + } + } else { + ONNX_NAMESPACE::TensorShapeProto shape_proto; + for (size_t j = 0; j < shape.NumDimensions(); ++j) { + auto dim = shape_proto.add_dim(); + dim->set_dim_value(shape[j]); + } + + ONNX_NAMESPACE::UnionShapeInfo(shape_proto, *output_tensor_type); + } + } } ptr = std::make_unique(elem_type); @@ -988,7 +1017,6 @@ class OpTester { // nullptr means None OrtValue which we will skip inserting into the feeds value.Init(ptr ? ptr.release() : nullptr, mltype, mltype->GetDeleteFunc()); - SequenceTensorTypeProto sequence_tensor_proto; #if !defined(DISABLE_OPTIONAL_TYPE) OptionalTypeProto optional_type_proto(sequence_tensor_proto.proto); auto node_arg = NodeArg(name, !is_optional_sequence_tensor_type @@ -998,7 +1026,18 @@ class OpTester { auto node_arg = NodeArg(name, &sequence_tensor_proto.proto); #endif - data.push_back(Data(std::move(node_arg), std::move(value), optional(), optional())); + optional rel; + optional abs; + + if (rel_error != 0.0f) { + rel = rel_error; + } + + if (abs_error != 0.0f) { + abs = abs_error; + } + + data.push_back(Data(std::move(node_arg), std::move(value), std::move(rel), std::move(abs))); } std::vector GetDimsForProto(gsl::span dims); diff --git a/onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_data_generator.py b/onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_data_generator.py new file mode 100644 index 0000000000..96f8c34562 --- /dev/null +++ b/onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_data_generator.py @@ -0,0 +1,192 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +"""This file is used to generate test data for Adam optimizer tests in + orttraining/orttraining/test/training_ops/cuda/optimizer/adamw_test.cc.""" + +import torch + + +class SingleParameterModule(torch.nn.Module): + """A dummy module containing only one trainable parameter.""" + + def __init__(self, input_size, hidden_size): + super().__init__() + self.fc1 = torch.nn.Linear(input_size, hidden_size, bias=False) + + def forward(self, input1): + """Module forward call.""" + out = self.fc1(input1) + return out + + +class MultipleParametersModule(torch.nn.Module): + """A dummy module containing multiple trainable parameters.""" + + def __init__(self, input_size, hidden_size, num_classes): + super().__init__() + self.fc1 = torch.nn.Linear(input_size, hidden_size) + self.relu = torch.nn.ReLU() + self.fc2 = torch.nn.Linear(hidden_size, num_classes) + + def forward(self, input1): + """Module forward call.""" + out = self.fc1(input1) + out = self.relu(out) + out = self.fc2(out) + return out + + +def generate_adamw_test_data(seed, _model_setup_func, data_func, train_step_count, adam_mode, json_file_name): + """Generate test data using specified model/data and other configs.""" + + def _run_step(model, input, target): + prediction = model(input) + criterion = torch.nn.MSELoss() + loss = criterion(prediction, target) + loss.backward() + + def _torch_tensor_to_str(torch_tensor): + """Torch tensor to string.""" + return torch_tensor.detach().cpu().numpy().tolist() + + def _build_param_index_to_name_mapping(model, map_result): + """Build index to name mapping, which is used to retrieve data from optimizer group.""" + index = 0 + for param in model.named_parameters(): + map_result[index] = param[0] + index += 1 + + torch.manual_seed(seed) + + # Prepare model, zero gradient. + pt_model = _model_setup_func() + pt_model.zero_grad() + + # Prepare optimizer. + adamw_optimizer = None + if adam_mode == 0: + adamw_optimizer = torch.optim.AdamW(pt_model.parameters(), lr=1e-3) + elif adam_mode == 1: + from transformers import AdamW + + adamw_optimizer = AdamW(pt_model.parameters(), lr=1e-3) + else: + raise ValueError(f"invalid adam_model: {adam_mode}") + + param_index_to_name_mapping = {} + _build_param_index_to_name_mapping(pt_model, param_index_to_name_mapping) + + # Dump the optimizer configs, our adam tests should align with these. + for group in adamw_optimizer.param_groups: + print( + f"beta1, beta2 : {group['betas']}, lr={group['lr']}," + " weight_decay={group['weight_decay']}, eps={group['eps']})" + ) + + p_dict = {} + g_dict = {} + m1_dict = {} + m2_dict = {} + for step in range(train_step_count): + input, target = data_func() + _run_step(pt_model, input, target) + torch.cuda.synchronize() + + for name, param in pt_model.named_parameters(): + if name not in p_dict: + p_dict[name] = [] + g_dict[name] = [] + m1_dict[name] = [] + m2_dict[name] = [] + + # Collect flattened parameter data. + p_dict[name].append(_torch_tensor_to_str(param.view(-1))) + if step == 0: + print(f"Weight name: {name}, shape: {param.shape}, data type: {param.dtype}") + + if step != train_step_count - 1: + # Collect flattened gradient data. + # Skip collecting the last step's gradients. + g_dict[name].append(_torch_tensor_to_str(param.grad.view(-1))) + + torch.cuda.synchronize() + + for group in adamw_optimizer.param_groups: + p_index = 0 + for param in group["params"]: + state = adamw_optimizer.state[param] + name = param_index_to_name_mapping[p_index] + # Collect flattened optimizer state data. + if len(state) == 0: + m1_dict[name].append(_torch_tensor_to_str(torch.zeros_like(param).view(-1))) + m2_dict[name].append(_torch_tensor_to_str(torch.zeros_like(param).view(-1))) + else: + m1_dict[name].append(_torch_tensor_to_str(state["exp_avg"].view(-1))) + m2_dict[name].append(_torch_tensor_to_str(state["exp_avg_sq"].view(-1))) + p_index += 1 + + adamw_optimizer.step() + adamw_optimizer.zero_grad() + + torch.cuda.synchronize() + data = { + "Parameters": p_dict, + "Gradients": g_dict, + "Momentum1s": m1_dict, + "Momentum2s": m2_dict, + } + import json + + with open(json_file_name, "w", encoding="utf-8") as f: + json.dump(data, f, ensure_ascii=False, indent=4) + + +def generate_adamw_single_weight_tests(adam_mode, run_step_count): + """Generate test data using specified mode of adamw.""" + seed = 8888 + device = "cuda" + batch_size, dimension_in, dimension_hidden = 2, 2, 3 + + def _model_setup_func(): + pt_model = SingleParameterModule(dimension_in, dimension_hidden).to(device) + return pt_model + + def _data_func(): + input = torch.randn(batch_size, dimension_in, device=device, dtype=torch.float32) + target = torch.randn(batch_size, dimension_hidden, device=device, dtype=torch.float32) + return input, target + + json_file_name = f"adamw_test_single_weight_mode_{adam_mode}.json" + generate_adamw_test_data(seed, _model_setup_func, _data_func, run_step_count, adam_mode, json_file_name) + + +def generate_adamw_multiple_weights_tests(adam_mode, run_step_count): + """Generate test data using specified mode of adamw.""" + seed = 6666 + device = "cuda" + batch_size, dimension_in, dimension_hidden, dim_out = 2, 2, 3, 2 + + def _model_setup_func(): + pt_model = MultipleParametersModule(dimension_in, dimension_hidden, dim_out).to(device) + return pt_model + + def data_func(): + input = torch.randn(batch_size, dimension_in, device=device, dtype=torch.float32) + target = torch.randn(batch_size, dim_out, device=device, dtype=torch.float32) + return input, target + + json_file_name = f"adamw_test_multiple_weights_mode_{adam_mode}.json" + generate_adamw_test_data(seed, _model_setup_func, data_func, run_step_count, adam_mode, json_file_name) + + +def main(): + """Main entry.""" + test_data_step_count = 11 + for adam_mode in range(0, 2): + generate_adamw_single_weight_tests(adam_mode, test_data_step_count) + generate_adamw_multiple_weights_tests(adam_mode, test_data_step_count) + + +if __name__ == "__main__": + main() diff --git a/onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_multiple_weights_mode_0.json b/onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_multiple_weights_mode_0.json new file mode 100644 index 0000000000..781b9aa86a --- /dev/null +++ b/onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_multiple_weights_mode_0.json @@ -0,0 +1,1117 @@ +{ + "Parameters": { + "fc1.weight": [ + [ + 0.039699044078588486, + -0.14054380357265472, + 0.6361624002456665, + 0.36362218856811523, + -0.21526536345481873, + 0.3459897041320801 + ], + [ + 0.040698643773794174, + -0.14154240489006042, + 0.6361560225486755, + 0.36361855268478394, + -0.2142632007598877, + 0.3449862599372864 + ], + [ + 0.040811654180288315, + -0.14246702194213867, + 0.6368938088417053, + 0.3628707826137543, + -0.21388788521289825, + 0.3441329002380371 + ], + [ + 0.04036901518702507, + -0.14334096014499664, + 0.6377384662628174, + 0.36347779631614685, + -0.2136160433292389, + 0.3434220850467682 + ], + [ + 0.04042157530784607, + -0.14390942454338074, + 0.6384292244911194, + 0.36397436261177063, + -0.213392972946167, + 0.34283924102783203 + ], + [ + 0.04013028368353844, + -0.1445896029472351, + 0.6390120983123779, + 0.36439353227615356, + -0.21320411562919617, + 0.34234607219696045 + ], + [ + 0.04048721864819527, + -0.14453566074371338, + 0.6386719942092896, + 0.3640187978744507, + -0.21312062442302704, + 0.3417902886867523 + ], + [ + 0.04066716134548187, + -0.14445945620536804, + 0.6383741497993469, + 0.3636910617351532, + -0.2129487693309784, + 0.3412230610847473 + ], + [ + 0.04085984453558922, + -0.14448948204517365, + 0.6387458443641663, + 0.3631588816642761, + -0.21277391910552979, + 0.3405272662639618 + ], + [ + 0.04108747839927673, + -0.1445508450269699, + 0.6390751004219055, + 0.36268603801727295, + -0.21261847019195557, + 0.33990922570228577 + ], + [ + 0.041274573653936386, + -0.14464154839515686, + 0.6393682956695557, + 0.36226364970207214, + -0.21247950196266174, + 0.33935728669166565 + ] + ], + "fc1.bias": [ + [ + 0.517955482006073, + -0.435107946395874, + -0.12118204683065414 + ], + [ + 0.5169503092765808, + -0.4351035952568054, + -0.12218083441257477 + ], + [ + 0.5160060524940491, + -0.43435510993003845, + -0.12318029999732971 + ], + [ + 0.5150943994522095, + -0.4335651397705078, + -0.12397350370883942 + ], + [ + 0.5141466856002808, + -0.4329172670841217, + -0.12462301552295685 + ], + [ + 0.5137341618537903, + -0.43236902356147766, + -0.12517181038856506 + ], + [ + 0.5133119821548462, + -0.4326229393482208, + -0.1256912797689438 + ], + [ + 0.5132887363433838, + -0.4328441321849823, + -0.1262883096933365 + ], + [ + 0.513144850730896, + -0.4324527084827423, + -0.1269497573375702 + ], + [ + 0.5130518078804016, + -0.4321047365665436, + -0.127536803483963 + ], + [ + 0.5128913521766663, + -0.4317937195301056, + -0.12806057929992676 + ] + ], + "fc2.weight": [ + [ + -0.5107945203781128, + 0.04572092741727829, + 0.21493154764175415, + -0.03950735926628113, + -0.3590133488178253, + 0.30300942063331604 + ], + [ + -0.5097894072532654, + 0.04572046920657158, + 0.21592940390110016, + -0.040506962686777115, + -0.3590097427368164, + 0.30200639367103577 + ], + [ + -0.5088188052177429, + 0.04646414518356323, + 0.21664850413799286, + -0.04150328040122986, + -0.35975027084350586, + 0.30123963952064514 + ], + [ + -0.5078955888748169, + 0.047161854803562164, + 0.2173282355070114, + -0.042312245815992355, + -0.3605462908744812, + 0.30057719349861145 + ], + [ + -0.5069411396980286, + 0.047733284533023834, + 0.21788464486598969, + -0.04310775548219681, + -0.36119768023490906, + 0.3000340163707733 + ], + [ + -0.506521463394165, + 0.04821619391441345, + 0.21835459768772125, + -0.043949466198682785, + -0.36174771189689636, + 0.29957443475723267 + ], + [ + -0.5059956312179565, + 0.047795072197914124, + 0.21833021938800812, + -0.044635750353336334, + -0.3613414764404297, + 0.29929450154304504 + ], + [ + -0.5057910680770874, + 0.04742725193500519, + 0.21826116740703583, + -0.04537033289670944, + -0.36098626255989075, + 0.2990230321884155 + ], + [ + -0.5054700970649719, + 0.047109540551900864, + 0.21849225461483002, + -0.04618074744939804, + -0.36095497012138367, + 0.2985096275806427 + ], + [ + -0.5052005648612976, + 0.046827442944049835, + 0.21869714558124542, + -0.04699970409274101, + -0.3609268069267273, + 0.298053503036499 + ], + [ + -0.5048867464065552, + 0.0465756319463253, + 0.21887977421283722, + -0.047773923724889755, + -0.36090126633644104, + 0.29764610528945923 + ] + ], + "fc2.bias": [ + [ + 0.023823332041502, + 0.5375885367393494 + ], + [ + 0.02482309378683567, + 0.5365831851959229 + ], + [ + 0.025780897587537766, + 0.5355805158615112 + ], + [ + 0.0266813226044178, + 0.5347577333450317 + ], + [ + 0.027619561180472374, + 0.5339611172676086 + ], + [ + 0.02818547561764717, + 0.533092737197876 + ], + [ + 0.028722008690238, + 0.5324080586433411 + ], + [ + 0.028895270079374313, + 0.5316585302352905 + ], + [ + 0.0292242132127285, + 0.5308324694633484 + ], + [ + 0.02950773946940899, + 0.5300171375274658 + ], + [ + 0.0298329945653677, + 0.529287576675415 + ] + ] + }, + "Gradients": { + "fc1.weight": [ + [ + -0.09586180746555328, + 0.17629344761371613, + 0.0, + 0.0, + -0.32529348134994507, + 0.04030714929103851 + ], + [ + 0.06833842396736145, + 0.08454680442810059, + -0.03257494419813156, + 0.0038211802020668983, + 0.1191219612956047, + 0.22664979100227356 + ], + [ + 0.14586228132247925, + 0.06190881133079529, + -0.0484643280506134, + -0.08001512289047241, + 0.010100812651216984, + 0.019007522612810135 + ], + [ + -0.137947678565979, + -0.046707138419151306, + 0.0, + 0.0, + 0.0, + 0.0 + ], + [ + 0.17390619218349457, + 0.5486620664596558, + 0.0, + 0.0, + 0.0, + 0.0 + ], + [ + -0.5882855653762817, + -0.7570041418075562, + 0.1739857792854309, + 0.22667543590068817, + 0.052081093192100525, + 0.06785327196121216 + ], + [ + 0.163396418094635, + -0.05490700528025627, + 0.0, + 0.0, + -0.07044854760169983, + 0.04285266995429993 + ], + [ + -0.04709528759121895, + 0.18947599828243256, + -0.5040349960327148, + 0.18439726531505585, + -0.01673167198896408, + 0.17345011234283447 + ], + [ + -0.0806330144405365, + 0.07000920176506042, + 0.0, + 0.0, + 0.0, + 0.0 + ], + [ + 0.022376326844096184, + 0.0734390839934349, + 0.0, + 0.0, + 0.0, + 0.0 + ] + ], + "fc1.bias": [ + [ + 0.309002548456192, + 0.0, + 0.3643433749675751 + ], + [ + 0.16205920279026031, + -0.02531510405242443, + 0.3757982850074768 + ], + [ + 0.13360804319381714, + -0.07762981206178665, + 0.017482705414295197 + ], + [ + 0.30951613187789917, + 0.0, + 0.0 + ], + [ + -0.2870372533798218, + 0.0, + 0.0 + ], + [ + 0.07509520649909973, + 0.1636124551296234, + 0.048975929617881775 + ], + [ + -0.3837743401527405, + 0.0, + 0.18683575093746185 + ], + [ + 0.175801083445549, + -0.45883896946907043, + 0.19719897210597992 + ], + [ + -0.05067116767168045, + 0.0, + 0.0 + ], + [ + 0.11543156206607819, + 0.0, + 0.0 + ] + ], + "fc2.weight": [ + [ + -0.31826627254486084, + 0.0, + -0.1022794172167778, + 0.7872374057769775, + 0.0, + 0.2039588987827301 + ], + [ + -0.2035769671201706, + -0.012987848371267319, + -0.007285610307008028, + 0.7206295132637024, + 0.021516945213079453, + 0.02756710723042488 + ], + [ + -0.13538049161434174, + -0.11737341433763504, + -0.022564847022294998, + 0.06767335534095764, + 0.05762113258242607, + 0.023621685802936554 + ], + [ + -0.34123802185058594, + 0.0, + 0.0, + 0.30149734020233154, + 0.0, + 0.0 + ], + [ + 0.3098701536655426, + 0.0, + 0.0, + 1.4480323791503906, + 0.0, + 0.0 + ], + [ + -0.2343340367078781, + 0.5370868444442749, + 0.08742398768663406, + -0.13039062917232513, + -0.26853200793266296, + -0.0437101349234581 + ], + [ + 0.30305129289627075, + 0.0, + 0.012864593416452408, + 0.6519851684570312, + 0.0, + 0.01125416811555624 + ], + [ + -0.22530147433280945, + -0.007280223537236452, + -0.0945650264620781, + 1.1682872772216797, + 0.15443314611911774, + 0.19558396935462952 + ], + [ + 0.024412253871560097, + 0.0, + 0.0, + 0.5862233638763428, + 0.0, + 0.0 + ], + [ + -0.11912607401609421, + 0.0, + 0.0, + 0.22673514485359192, + 0.0, + 0.0 + ] + ], + "fc2.bias": [ + [ + -0.7384594678878784, + 1.726222276687622 + ], + [ + -0.44500619173049927, + 1.5997315645217896 + ], + [ + -0.2772650718688965, + 0.1799774169921875 + ], + [ + -0.6588927507400513, + 0.5939797759056091 + ], + [ + 0.39522403478622437, + 2.0108203887939453 + ], + [ + -0.11679565906524658, + -0.3625909686088562 + ], + [ + 0.6392423510551453, + 1.3513944149017334 + ], + [ + -0.5523114204406738, + 2.2823963165283203 + ], + [ + 0.02510213851928711, + 0.8224809169769287 + ], + [ + -0.2295447289943695, + 0.01137387752532959 + ] + ] + }, + "Momentum1s": { + "fc1.weight": [ + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + [ + -0.009586180560290813, + 0.017629345878958702, + 0.0, + 0.0, + -0.032529350370168686, + 0.004030715208500624 + ], + [ + -0.0017937193624675274, + 0.024321090430021286, + -0.003257494419813156, + 0.00038211801438592374, + -0.017364216968417168, + 0.026292623952031136 + ], + [ + 0.012971880845725536, + 0.028079861775040627, + -0.0077781775034964085, + -0.007657606154680252, + -0.01461771223694086, + 0.02556411363184452 + ], + [ + -0.0021200752817094326, + 0.020601162686944008, + -0.007000359706580639, + -0.0068918452598154545, + -0.013155940920114517, + 0.023007702082395554 + ], + [ + 0.015482551418244839, + 0.07340725511312485, + -0.0063003236427903175, + -0.006202660501003265, + -0.01184034626930952, + 0.02070693112909794 + ], + [ + -0.0448942631483078, + -0.00963388942182064, + 0.011728287674486637, + 0.017085149884223938, + -0.005448201671242714, + 0.025421565398573875 + ], + [ + -0.02406519465148449, + -0.014161201193928719, + 0.010555458255112171, + 0.015376634895801544, + -0.011948236264288425, + 0.027164675295352936 + ], + [ + -0.02636820264160633, + 0.006202519405633211, + -0.04090358689427376, + 0.032278697937726974, + -0.012426580302417278, + 0.04179321974515915 + ], + [ + -0.03179468587040901, + 0.012583187781274319, + -0.03681322559714317, + 0.029050827026367188, + -0.011183922179043293, + 0.037613898515701294 + ], + [ + -0.026377582922577858, + 0.018668776378035545, + -0.03313190117478371, + 0.02614574320614338, + -0.010065529495477676, + 0.033852506428956985 + ] + ], + "fc1.bias": [ + [ + 0.0, + 0.0, + 0.0 + ], + [ + 0.0309002548456192, + 0.0, + 0.03643433749675751 + ], + [ + 0.04401614889502525, + -0.0025315105449408293, + 0.07037073373794556 + ], + [ + 0.05297533795237541, + -0.010041340254247189, + 0.06508193165063858 + ], + [ + 0.07862941920757294, + -0.009037205949425697, + 0.05857373774051666 + ], + [ + 0.042062751948833466, + -0.008133484981954098, + 0.05271636322140694 + ], + [ + 0.04536599665880203, + 0.00904110912233591, + 0.05234232172369957 + ], + [ + 0.0024519632570445538, + 0.008136997930705547, + 0.06579166650772095 + ], + [ + 0.019786875694990158, + -0.03856059908866882, + 0.07893239706754684 + ], + [ + 0.012741071172058582, + -0.03470453992486, + 0.07103915512561798 + ], + [ + 0.023010119795799255, + -0.03123408555984497, + 0.06393523514270782 + ] + ], + "fc2.weight": [ + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + [ + -0.031826626509428024, + 0.0, + -0.010227941907942295, + 0.07872374355792999, + 0.0, + 0.02039588987827301 + ], + [ + -0.04900166019797325, + -0.0012987848604097962, + -0.009933708235621452, + 0.1429143249988556, + 0.0021516946144402027, + 0.021113011986017227 + ], + [ + -0.05763953924179077, + -0.01290624774992466, + -0.011196821928024292, + 0.13539022207260132, + 0.007698638364672661, + 0.02136388048529625 + ], + [ + -0.08599938452243805, + -0.011615622788667679, + -0.010077139362692833, + 0.15200093388557434, + 0.006928774528205395, + 0.019227491691708565 + ], + [ + -0.046412430703639984, + -0.010454060509800911, + -0.00906942505389452, + 0.2816040813922882, + 0.006235897075384855, + 0.017304742708802223 + ], + [ + -0.06520459055900574, + 0.044300030916929245, + 0.0005799162317998707, + 0.24040459096431732, + -0.021240893751382828, + 0.011203254573047161 + ], + [ + -0.02837899886071682, + 0.03987002745270729, + 0.0018083839677274227, + 0.2815626561641693, + -0.019116804003715515, + 0.01120834518224001 + ], + [ + -0.0480712465941906, + 0.035155002027750015, + -0.00782895740121603, + 0.3702351152896881, + -0.0017618080601096153, + 0.029645908623933792 + ], + [ + -0.04082289710640907, + 0.031639501452445984, + -0.00704606156796217, + 0.39183393120765686, + -0.001585627207532525, + 0.026681317016482353 + ], + [ + -0.04865321144461632, + 0.028475550934672356, + -0.006341455038636923, + 0.3753240406513214, + -0.0014270644169300795, + 0.024013184010982513 + ] + ], + "fc2.bias": [ + [ + 0.0, + 0.0 + ], + [ + -0.07384594529867172, + 0.17262223362922668 + ], + [ + -0.11096196621656418, + 0.31533315777778625 + ], + [ + -0.1275922656059265, + 0.3017975687980652 + ], + [ + -0.18072231113910675, + 0.33101576566696167 + ], + [ + -0.12312767654657364, + 0.49899622797966003 + ], + [ + -0.12249447405338287, + 0.41283750534057617 + ], + [ + -0.04632079228758812, + 0.5066931843757629 + ], + [ + -0.09691985696554184, + 0.6842634677886963 + ], + [ + -0.08471765369176865, + 0.6980851888656616 + ], + [ + -0.09920036047697067, + 0.6294140219688416 + ] + ] + }, + "Momentum2s": { + "fc1.weight": [ + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + [ + 9.189487172989175e-06, + 3.107938027824275e-05, + 0.0, + 0.0, + 0.00010581585229374468, + 1.6246663108177017e-06 + ], + [ + 1.3850438335794024e-05, + 3.8196463719941676e-05, + 1.0611270226945635e-06, + 1.4601418740767258e-08, + 0.00011990008351858705, + 5.2993171266280115e-05 + ], + [ + 3.5112392652081326e-05, + 4.199097020318732e-05, + 3.4088570828316733e-06, + 6.417007170966826e-06, + 0.000119882206490729, + 5.3301468142308295e-05 + ], + [ + 5.410684389062226e-05, + 4.4130538299214095e-05, + 3.4054482966894284e-06, + 6.4105902310984675e-06, + 0.00011976232781307772, + 5.3248168114805594e-05 + ], + [ + 8.429610170423985e-05, + 0.0003451164811849594, + 3.402042921152315e-06, + 6.404179657693021e-06, + 0.00011964256555074826, + 5.319491901900619e-05 + ], + [ + 0.00043029175139963627, + 0.0009178266627714038, + 3.3669693948468193e-05, + 5.7779529015533626e-05, + 0.00012223536032252014, + 5.774579040007666e-05 + ], + [ + 0.0004565598501358181, + 0.0009199236519634724, + 3.363602445460856e-05, + 5.7721750636119395e-05, + 0.0001270761276828125, + 5.952439460088499e-05 + ], + [ + 0.00045832127216272056, + 0.0009549048845656216, + 0.0002876536746043712, + 9.166638483293355e-05, + 0.00012722901010420173, + 8.95498160389252e-05 + ], + [ + 0.0004643646243494004, + 0.0009588512475602329, + 0.00028736601234413683, + 9.157472231891006e-05, + 0.00012710178270936012, + 8.946027082856745e-05 + ], + [ + 0.00046440097503364086, + 0.0009632856817916036, + 0.00028707864112220705, + 9.148314711637795e-05, + 0.00012697468628175557, + 8.937081292970106e-05 + ] + ], + "fc1.bias": [ + [ + 0.0, + 0.0, + 0.0 + ], + [ + 9.548258094582707e-05, + 0.0, + 0.00013274610682856292 + ], + [ + 0.00012165028601884842, + 6.408545232261531e-07, + 0.0002738377370405942 + ], + [ + 0.00013937974290456623, + 6.666601620963775e-06, + 0.00027386954752728343 + ], + [ + 0.00023504059936385602, + 6.659935024799779e-06, + 0.00027359568048268557 + ], + [ + 0.0003171959542669356, + 6.653275249846047e-06, + 0.0002733220753725618 + ], + [ + 0.00032251805532723665, + 3.3415657526347786e-05, + 0.00027544741169549525 + ], + [ + 0.0004694783128798008, + 3.338224269100465e-05, + 0.00031007957295514643 + ], + [ + 0.0004999148659408092, + 0.00024388206657022238, + 0.0003486569330561906 + ], + [ + 0.0005019825184717774, + 0.00024363819102291018, + 0.00034830826916731894 + ], + [ + 0.0005148049676790833, + 0.00024339456285815686, + 0.00034795995452441275 + ] + ], + "fc2.weight": [ + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + [ + 0.00010129342990694568, + 0.0, + 1.0461079000378959e-05, + 0.0006197427865117788, + 0.0, + 4.159923628321849e-05 + ], + [ + 0.00014263571938499808, + 1.6868420971150044e-07, + 1.0503697922104038e-05, + 0.0011384299723431468, + 4.6297896005853545e-07, + 4.2317580664530396e-05 + ], + [ + 0.00016082095680758357, + 1.3945034879725426e-05, + 1.100236659112852e-05, + 0.001141871209256351, + 3.782711246458348e-06, + 4.283324597054161e-05 + ], + [ + 0.0002771035360638052, + 1.3931089597463142e-05, + 1.0991364433721174e-05, + 0.0012316300999373198, + 3.7789286579936743e-06, + 4.2790412408066913e-05 + ], + [ + 0.00037284597055986524, + 1.3917158867116086e-05, + 1.0980373190250248e-05, + 0.003327196231111884, + 3.7751497075078078e-06, + 4.27476225013379e-05 + ], + [ + 0.00042738555930554867, + 0.00030236554448492825, + 1.8612347048474476e-05, + 0.003340870840474963, + 7.588081643916667e-05, + 4.4615451770368963e-05 + ], + [ + 0.0005187982460483909, + 0.00030206318479031324, + 1.8759232261800207e-05, + 0.003762614680454135, + 7.580493547720835e-05, + 4.469749183044769e-05 + ], + [ + 0.0005690401885658503, + 0.0003018141142092645, + 2.7683017833624035e-05, + 0.0051237475126981735, + 9.957873407984152e-05, + 8.290588448289782e-05 + ], + [ + 0.0005690670805051923, + 0.0003015123074874282, + 2.765533463389147e-05, + 0.005462281871587038, + 9.947915532393381e-05, + 8.282298222184181e-05 + ], + [ + 0.0005826890701428056, + 0.00030121079180389643, + 2.7627680537989363e-05, + 0.005508228670805693, + 9.93796784314327e-05, + 8.274015999631956e-05 + ] + ], + "fc2.bias": [ + [ + 0.0, + 0.0 + ], + [ + 0.0005453223711811006, + 0.0029798434115946293 + ], + [ + 0.0007428076351061463, + 0.005536004900932312 + ], + [ + 0.0008189407526515424, + 0.005562860984355211 + ], + [ + 0.0012522615725174546, + 0.00591011019423604 + ], + [ + 0.0014072112971916795, + 0.009947598911821842 + ], + [ + 0.0014194452669471502, + 0.010069123469293118 + ], + [ + 0.0018266566330567002, + 0.011885321699082851 + ], + [ + 0.002129877917468548, + 0.01708276942372322 + ], + [ + 0.0021283780224621296, + 0.017742162570357323 + ], + [ + 0.0021789404563605785, + 0.017724549397826195 + ] + ] + } +} \ No newline at end of file diff --git a/onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_multiple_weights_mode_1.json b/onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_multiple_weights_mode_1.json new file mode 100644 index 0000000000..a2b4792125 --- /dev/null +++ b/onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_multiple_weights_mode_1.json @@ -0,0 +1,1117 @@ +{ + "Parameters": { + "fc1.weight": [ + [ + 0.039699044078588486, + -0.14054380357265472, + 0.6361624002456665, + 0.36362218856811523, + -0.21526536345481873, + 0.3459897041320801 + ], + [ + 0.040698714554309845, + -0.14154362678527832, + 0.6361624002456665, + 0.36362218856811523, + -0.21426546573638916, + 0.34499049186706543 + ], + [ + 0.04081208258867264, + -0.14246951043605804, + 0.6369057893753052, + 0.36288416385650635, + -0.21389234066009521, + 0.3441406786441803 + ], + [ + 0.0403699055314064, + -0.14334474503993988, + 0.63775634765625, + 0.3634945750236511, + -0.21362267434597015, + 0.3434334099292755 + ], + [ + 0.04042285308241844, + -0.14391455054283142, + 0.6384530663490295, + 0.3639945983886719, + -0.21340177953243256, + 0.3428540825843811 + ], + [ + 0.04013199359178543, + -0.144596129655838, + 0.6390419602394104, + 0.3644172251224518, + -0.21321506798267365, + 0.3423644006252289 + ], + [ + 0.040489312261343, + -0.1445436030626297, + 0.6387082934379578, + 0.3640461564064026, + -0.21313372254371643, + 0.3418121337890625 + ], + [ + 0.04066965728998184, + -0.14446882903575897, + 0.6384168863296509, + 0.3637221157550812, + -0.2129640132188797, + 0.34124839305877686 + ], + [ + 0.040862735360860825, + -0.14450028538703918, + 0.6387949585914612, + 0.3631936311721802, + -0.212791308760643, + 0.3405560851097107 + ], + [ + 0.041090745478868484, + -0.14456306397914886, + 0.6391305923461914, + 0.3627244830131531, + -0.2126379907131195, + 0.33994150161743164 + ], + [ + 0.04127822443842888, + -0.14465518295764923, + 0.6394301056861877, + 0.36230576038360596, + -0.2125011533498764, + 0.3393929898738861 + ] + ], + "fc1.bias": [ + [ + 0.517955482006073, + -0.435107946395874, + -0.12118204683065414 + ], + [ + 0.5169556140899658, + -0.435107946395874, + -0.1221819594502449 + ], + [ + 0.5160166025161743, + -0.4343647360801697, + -0.12318258732557297 + ], + [ + 0.515110194683075, + -0.4335794150829315, + -0.12397698312997818 + ], + [ + 0.5141677260398865, + -0.4329361319541931, + -0.12462770193815231 + ], + [ + 0.5137603282928467, + -0.432392418384552, + -0.12517771124839783 + ], + [ + 0.5133432149887085, + -0.4326505959033966, + -0.12569840252399445 + ], + [ + 0.5133250951766968, + -0.4328760802745819, + -0.12629665434360504 + ], + [ + 0.5131863355636597, + -0.4324890077114105, + -0.12695933878421783 + ], + [ + 0.513098418712616, + -0.4321453869342804, + -0.12754763662815094 + ], + [ + 0.5129430294036865, + -0.431838721036911, + -0.12807267904281616 + ] + ], + "fc2.weight": [ + [ + -0.5107945203781128, + 0.04572092741727829, + 0.21493154764175415, + -0.03950735926628113, + -0.3590133488178253, + 0.30300942063331604 + ], + [ + -0.5097945928573608, + 0.04572092741727829, + 0.215931236743927, + -0.04050732031464577, + -0.3590133488178253, + 0.30200958251953125 + ], + [ + -0.5088291764259338, + 0.0464632585644722, + 0.2166522741317749, + -0.041504018008708954, + -0.35975638031959534, + 0.30124595761299133 + ], + [ + -0.5079110860824585, + 0.047161247581243515, + 0.21733398735523224, + -0.04231337830424309, + -0.36055558919906616, + 0.30058664083480835 + ], + [ + -0.5069617629051208, + 0.04773299768567085, + 0.21789240837097168, + -0.04310929775238037, + -0.36121025681495667, + 0.30004656314849854 + ], + [ + -0.5065471529960632, + 0.048216257244348526, + 0.2183644026517868, + -0.043951425701379776, + -0.361763596534729, + 0.2995900809764862 + ], + [ + -0.5060263872146606, + 0.04779564589262009, + 0.21834219992160797, + -0.04463813826441765, + -0.3613610565662384, + 0.29931318759918213 + ], + [ + -0.5058268308639526, + 0.047428324818611145, + 0.21827533841133118, + -0.045373160392045975, + -0.36100950837135315, + 0.29904475808143616 + ], + [ + -0.5055108666419983, + 0.047111108899116516, + 0.21850857138633728, + -0.04618402197957039, + -0.360981822013855, + 0.2985343635082245 + ], + [ + -0.5052463412284851, + 0.04682950675487518, + 0.2187156230211258, + -0.04700344055891037, + -0.3609572649002075, + 0.29808127880096436 + ], + [ + -0.5049375295639038, + 0.04657818377017975, + 0.2189004123210907, + -0.047778137028217316, + -0.3609353303909302, + 0.2976769208908081 + ] + ], + "fc2.bias": [ + [ + 0.023823332041502, + 0.5375885367393494 + ], + [ + 0.02482328936457634, + 0.5365885496139526 + ], + [ + 0.025781307369470596, + 0.5355912446975708 + ], + [ + 0.026681963354349136, + 0.5347738265991211 + ], + [ + 0.02762044407427311, + 0.5339825749397278 + ], + [ + 0.02818664163351059, + 0.5331195592880249 + ], + [ + 0.028723470866680145, + 0.5324403047561646 + ], + [ + 0.028897050768136978, + 0.531696081161499 + ], + [ + 0.029226312413811684, + 0.5308753252029419 + ], + [ + 0.029510175809264183, + 0.5300652980804443 + ], + [ + 0.029835781082510948, + 0.5293410420417786 + ] + ] + }, + "Gradients": { + "fc1.weight": [ + [ + -0.09586180746555328, + 0.17629344761371613, + 0.0, + 0.0, + -0.32529348134994507, + 0.04030714929103851 + ], + [ + 0.06834118068218231, + 0.08454809337854385, + -0.032576028257608414, + 0.0038213070947676897, + 0.11912338435649872, + 0.22665248811244965 + ], + [ + 0.14587241411209106, + 0.061909887939691544, + -0.048467256128787994, + -0.08001699298620224, + 0.01010121963918209, + 0.01900828815996647 + ], + [ + -0.1379551738500595, + -0.04670995473861694, + 0.0, + 0.0, + 0.0, + 0.0 + ], + [ + 0.1739121377468109, + 0.5486636757850647, + 0.0, + 0.0, + 0.0, + 0.0 + ], + [ + -0.5883157253265381, + -0.7570428252220154, + 0.17399102449417114, + 0.22668226063251495, + 0.052081551402807236, + 0.06785386800765991 + ], + [ + 0.16340069472789764, + -0.05490869656205177, + 0.0, + 0.0, + -0.07045339047908783, + 0.04285561293363571 + ], + [ + -0.04708762839436531, + 0.18948915600776672, + -0.5040743350982666, + 0.1844116598367691, + -0.016733329743146896, + 0.17346729338169098 + ], + [ + -0.08061128109693527, + 0.06998594105243683, + 0.0, + 0.0, + 0.0, + 0.0 + ], + [ + 0.022378630936145782, + 0.07342363893985748, + 0.0, + 0.0, + 0.0, + 0.0 + ] + ], + "fc1.bias": [ + [ + 0.309002548456192, + 0.0, + 0.3643433749675751 + ], + [ + 0.16206324100494385, + -0.025315944105386734, + 0.37580275535583496 + ], + [ + 0.13361552357673645, + -0.0776328444480896, + 0.017483409494161606 + ], + [ + 0.30953294038772583, + 0.0, + 0.0 + ], + [ + -0.2870349884033203, + 0.0, + 0.0 + ], + [ + 0.07510882616043091, + 0.16361738741397858, + 0.048976361751556396 + ], + [ + -0.3837846517562866, + 0.0, + 0.18684858083724976 + ], + [ + 0.17582830786705017, + -0.45887479186058044, + 0.19721850752830505 + ], + [ + -0.05064861476421356, + 0.0, + 0.0 + ], + [ + 0.11546661704778671, + 0.0, + 0.0 + ] + ], + "fc2.weight": [ + [ + -0.31826627254486084, + 0.0, + -0.1022794172167778, + 0.7872374057769775, + 0.0, + 0.2039588987827301 + ], + [ + -0.20358160138130188, + -0.01298887562006712, + -0.007285935338586569, + 0.7206387519836426, + 0.02151758410036564, + 0.027568239718675613 + ], + [ + -0.1353888362646103, + -0.11738108098506927, + -0.022565873339772224, + 0.06767717748880386, + 0.05762537568807602, + 0.02362271584570408 + ], + [ + -0.34125667810440063, + 0.0, + 0.0, + 0.3015143871307373, + 0.0, + 0.0 + ], + [ + 0.3098617494106293, + 0.0, + 0.0, + 1.4481031894683838, + 0.0, + 0.0 + ], + [ + -0.2343558967113495, + 0.5371115207672119, + 0.08742881566286087, + -0.13038989901542664, + -0.2685457170009613, + -0.04371276870369911 + ], + [ + 0.3030567765235901, + 0.0, + 0.012865260243415833, + 0.6520397067070007, + 0.0, + 0.011255647987127304 + ], + [ + -0.22533443570137024, + -0.007282924372702837, + -0.0945744663476944, + 1.168387532234192, + 0.15443821251392365, + 0.19560372829437256 + ], + [ + 0.024374589323997498, + 0.0, + 0.0, + 0.5862966775894165, + 0.0, + 0.0 + ], + [ + -0.11917039006948471, + 0.0, + 0.0, + 0.22678035497665405, + 0.0, + 0.0 + ] + ], + "fc2.bias": [ + [ + -0.7384594678878784, + 1.726222276687622 + ], + [ + -0.44501104950904846, + 1.5997353792190552 + ], + [ + -0.2772747874259949, + 0.17998221516609192 + ], + [ + -0.6589083671569824, + 0.5939946174621582 + ], + [ + 0.3951956629753113, + 2.0108397006988525 + ], + [ + -0.1168164610862732, + -0.3625774383544922 + ], + [ + 0.639214813709259, + 1.3514249324798584 + ], + [ + -0.5523422360420227, + 2.28243350982666 + ], + [ + 0.025046706199645996, + 0.8225188255310059 + ], + [ + -0.22959741950035095, + 0.011417150497436523 + ] + ] + }, + "Momentum1s": { + "fc1.weight": [ + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + [ + -0.009586180560290813, + 0.017629345878958702, + 0.0, + 0.0, + -0.032529350370168686, + 0.004030715208500624 + ], + [ + -0.0017934436909854412, + 0.024321218952536583, + -0.003257602918893099, + 0.0003821307036560029, + -0.017364075407385826, + 0.026292892172932625 + ], + [ + 0.01297314278781414, + 0.028080085292458534, + -0.00777856819331646, + -0.007657781708985567, + -0.014617545530200005, + 0.02556443214416504 + ], + [ + -0.002119689481332898, + 0.02060108073055744, + -0.007000711280852556, + -0.006892003584653139, + -0.013155790977180004, + 0.023007988929748535 + ], + [ + 0.01548349391669035, + 0.07340734452009201, + -0.0063006398268043995, + -0.0062028029933571815, + -0.011840211227536201, + 0.02070719003677368 + ], + [ + -0.04489642754197121, + -0.009637676179409027, + 0.011728527024388313, + 0.01708570308983326, + -0.005448034964501858, + 0.025421857833862305 + ], + [ + -0.024066712707281113, + -0.014164777472615242, + 0.010555674321949482, + 0.015377132222056389, + -0.011948570609092712, + 0.027165232226252556 + ], + [ + -0.026368804275989532, + 0.006200616713613272, + -0.040907327085733414, + 0.032280586659908295, + -0.01242704689502716, + 0.041795436292886734 + ], + [ + -0.03179305046796799, + 0.012579148635268211, + -0.03681659325957298, + 0.029052527621388435, + -0.011184342205524445, + 0.03761589154601097 + ], + [ + -0.02637588046491146, + 0.018663598224520683, + -0.033134933561086655, + 0.026147274300456047, + -0.01006590761244297, + 0.033854302018880844 + ] + ], + "fc1.bias": [ + [ + 0.0, + 0.0, + 0.0 + ], + [ + 0.0309002548456192, + 0.0, + 0.03643433749675751 + ], + [ + 0.044016554951667786, + -0.0025315943639725447, + 0.07037118077278137 + ], + [ + 0.052976448088884354, + -0.010041719302535057, + 0.06508240103721619 + ], + [ + 0.07863209396600723, + -0.009037546813488007, + 0.05857415869832039 + ], + [ + 0.04206538572907448, + -0.00813379231840372, + 0.05271674320101738 + ], + [ + 0.04536972939968109, + 0.009041326120495796, + 0.05234270170331001 + ], + [ + 0.002454288536682725, + 0.008137193508446217, + 0.06579329073429108 + ], + [ + 0.01979169063270092, + -0.03856400400400162, + 0.07893580943346024 + ], + [ + 0.012747658416628838, + -0.034707602113485336, + 0.07104222476482391 + ], + [ + 0.02301955409348011, + -0.031236840412020683, + 0.06393799930810928 + ] + ], + "fc2.weight": [ + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + [ + -0.031826626509428024, + 0.0, + -0.010227941907942295, + 0.07872374355792999, + 0.0, + 0.02039588987827301 + ], + [ + -0.04900212585926056, + -0.0012988875387236476, + -0.009933740831911564, + 0.1429152488708496, + 0.002151758410036564, + 0.02111312560737133 + ], + [ + -0.05764079466462135, + -0.01290710736066103, + -0.011196953244507313, + 0.1353914439678192, + 0.007699120324105024, + 0.021364083513617516 + ], + [ + -0.08600237965583801, + -0.011616396717727184, + -0.010077257640659809, + 0.15200373530387878, + 0.006929208058863878, + 0.01922767423093319 + ], + [ + -0.046415962278842926, + -0.010454757139086723, + -0.009069531224668026, + 0.28161367774009705, + 0.006236287299543619, + 0.017304906621575356 + ], + [ + -0.06520995497703552, + 0.044301871210336685, + 0.0005803040694445372, + 0.24041330814361572, + -0.021241914480924606, + 0.011203138157725334 + ], + [ + -0.028383279219269753, + 0.03987168148159981, + 0.0018087996868416667, + 0.2815759479999542, + -0.019117722287774086, + 0.011208388954401016 + ], + [ + -0.048078395426273346, + 0.035156216472387314, + -0.007829527370631695, + 0.37025710940361023, + -0.001762128435075283, + 0.02964792214334011 + ], + [ + -0.04083309695124626, + 0.03164059296250343, + -0.0070465742610394955, + 0.39186105132102966, + -0.0015859155682846904, + 0.026683129370212555 + ], + [ + -0.048666827380657196, + 0.028476532548666, + -0.006341916508972645, + 0.37535297870635986, + -0.0014273240230977535, + 0.02401481568813324 + ] + ], + "fc2.bias": [ + [ + 0.0, + 0.0 + ], + [ + -0.07384594529867172, + 0.17262223362922668 + ], + [ + -0.11096245050430298, + 0.3153335452079773 + ], + [ + -0.12759368121623993, + 0.30179840326309204 + ], + [ + -0.18072514235973358, + 0.33101800084114075 + ], + [ + -0.12313306331634521, + 0.4990001916885376 + ], + [ + -0.12250140309333801, + 0.41284242272377014 + ], + [ + -0.04632978141307831, + 0.5067006349563599 + ], + [ + -0.09693102538585663, + 0.684273898601532 + ], + [ + -0.08473324775695801, + 0.6980984210968018 + ], + [ + -0.0992196574807167, + 0.6294302940368652 + ] + ] + }, + "Momentum2s": { + "fc1.weight": [ + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + [ + 9.189487172989175e-06, + 3.107938027824275e-05, + 0.0, + 0.0, + 0.00010581585229374468, + 1.6246663108177017e-06 + ], + [ + 1.3850814866600558e-05, + 3.81966819986701e-05, + 1.0611976222207886e-06, + 1.460238863160157e-08, + 0.0001199004182126373, + 5.29943936271593e-05 + ], + [ + 3.511572504066862e-05, + 4.1991323087131605e-05, + 3.4092115583916893e-06, + 6.4173073042184114e-06, + 0.00011988255573669448, + 5.330271596903913e-05 + ], + [ + 5.411224265117198e-05, + 4.413115311763249e-05, + 3.4058023175020935e-06, + 6.410889909602702e-06, + 0.0001197626770590432, + 5.324941230355762e-05 + ], + [ + 8.4303566836752e-05, + 0.00034511886769905686, + 3.4023964872176293e-06, + 6.404478881449904e-06, + 0.00011964291479671374, + 5.319616320775822e-05 + ], + [ + 0.0004303346504457295, + 0.0009178876644000411, + 3.367187309777364e-05, + 5.778292324976064e-05, + 0.00012223576777614653, + 5.774711462436244e-05 + ], + [ + 0.0004566041170619428, + 0.0009199847700074315, + 3.36381999659352e-05, + 5.7725141232367605e-05, + 0.00012707721907645464, + 5.952597348368727e-05 + ], + [ + 0.0004583647532854229, + 0.0009549709502607584, + 0.0002876954968087375, + 9.16750796022825e-05, + 0.00012723014515358955, + 8.95573539310135e-05 + ], + [ + 0.0004644045839086175, + 0.0009589140536263585, + 0.0002874078054446727, + 9.158340253634378e-05, + 0.00012710291775874794, + 8.946779416874051e-05 + ], + [ + 0.0004644409636966884, + 0.0009633461595512927, + 0.00028712040511891246, + 9.149182005785406e-05, + 0.00012697582133114338, + 8.937832899391651e-05 + ] + ], + "fc1.bias": [ + [ + 0.0, + 0.0, + 0.0 + ], + [ + 9.548258094582707e-05, + 0.0, + 0.00013274610682856292 + ], + [ + 0.00012165159569121897, + 6.40897042103461e-07, + 0.00027384108398109674 + ], + [ + 0.0001393830607412383, + 6.6671150307229254e-06, + 0.0002738729235716164 + ], + [ + 0.0002350543363718316, + 6.660447979811579e-06, + 0.00027359905652701855 + ], + [ + 0.0003172083816025406, + 6.653787750110496e-06, + 0.0002733254514168948 + ], + [ + 0.00032253251993097365, + 3.3417785743949935e-05, + 0.0002754508168436587 + ], + [ + 0.00046950066462159157, + 3.338436727062799e-05, + 0.00031008778023533523 + ], + [ + 0.0004999467637389898, + 0.0002439170639263466, + 0.00034867285285145044 + ], + [ + 0.0005020120879635215, + 0.00024367314472328871, + 0.0003483241889625788 + ], + [ + 0.0005148426280356944, + 0.0002434294729027897, + 0.0003479758743196726 + ] + ], + "fc2.weight": [ + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + [ + 0.00010129342990694568, + 0.0, + 1.0461079000378959e-05, + 0.0006197427865117788, + 0.0, + 4.159923628321849e-05 + ], + [ + 0.00014263761113397777, + 1.687108976966556e-07, + 1.0503702469577547e-05, + 0.001138443243689835, + 4.6300644385155465e-07, + 4.231764251017012e-05 + ], + [ + 0.00016082510410342366, + 1.3946860235591885e-05, + 1.100241752283182e-05, + 0.0011418850626796484, + 3.7832276120752795e-06, + 4.283335874788463e-05 + ], + [ + 0.00027712041628547013, + 1.3932913134340197e-05, + 1.0991415365424473e-05, + 0.001231654081493616, + 3.7794443414895795e-06, + 4.279052518540993e-05 + ], + [ + 0.00037285761209204793, + 1.3918980585003737e-05, + 1.0980424121953547e-05, + 0.0033274253364652395, + 3.775664936256362e-06, + 4.274773527868092e-05 + ], + [ + 0.00042740744538605213, + 0.00030239386251196265, + 1.861324199126102e-05, + 0.0033410994801670313, + 7.588868902530521e-05, + 4.461579374037683e-05 + ], + [ + 0.0005188234499655664, + 0.0003020914737135172, + 1.8760143575491384e-05, + 0.003762914100661874, + 7.581280078738928e-05, + 4.469786654226482e-05 + ], + [ + 0.0005690802354365587, + 0.00030184240313246846, + 2.7685715394909494e-05, + 0.005124280694872141, + 9.958815644495189e-05, + 8.291398989968002e-05 + ], + [ + 0.0005691052647307515, + 0.0003015405673068017, + 2.7658030376187526e-05, + 0.005462900269776583, + 9.948857041308656e-05, + 8.28310803626664e-05 + ], + [ + 0.0005827377317473292, + 0.00030123902251943946, + 2.763037264230661e-05, + 0.005508867092430592, + 9.938908624462783e-05, + 8.274825086118653e-05 + ] + ], + "fc2.bias": [ + [ + 0.0, + 0.0 + ], + [ + 0.0005453223711811006, + 0.0029798434115946293 + ], + [ + 0.0007428119424730539, + 0.005536017008125782 + ], + [ + 0.000818950473330915, + 0.0055628749541938305 + ], + [ + 0.0012522918405011296, + 0.005910141859203577 + ], + [ + 0.0014072192134335637, + 0.009947707876563072 + ], + [ + 0.0014194580726325512, + 0.010069223120808601 + ], + [ + 0.0018266342813149095, + 0.011885503306984901 + ], + [ + 0.0021298895590007305, + 0.017083121463656425 + ], + [ + 0.0021283868700265884, + 0.017742576077580452 + ], + [ + 0.0021789735183119774, + 0.017724964767694473 + ] + ] + } +} \ No newline at end of file diff --git a/onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_single_weight_mode_0.json b/onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_single_weight_mode_0.json new file mode 100644 index 0000000000..b80a2a60d6 --- /dev/null +++ b/onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_single_weight_mode_0.json @@ -0,0 +1,362 @@ +{ + "Parameters": { + "fc1.weight": [ + [ + -0.18330414593219757, + 0.6739549040794373, + 0.31170889735221863, + 0.42830976843833923, + -0.3957911729812622, + 0.07424858212471008 + ], + [ + -0.1823023110628128, + 0.6729481816291809, + 0.31270575523376465, + 0.4273054897785187, + -0.3947872221469879, + 0.07324783504009247 + ], + [ + -0.1813177466392517, + 0.6725299954414368, + 0.31352442502975464, + 0.42674997448921204, + -0.395175576210022, + 0.07235248386859894 + ], + [ + -0.18059653043746948, + 0.6721751093864441, + 0.3133564889431, + 0.426154226064682, + -0.3952518403530121, + 0.07177812606096268 + ], + [ + -0.1797766238451004, + 0.6717073917388916, + 0.31389161944389343, + 0.42545583844184875, + -0.3956265151500702, + 0.0718231126666069 + ], + [ + -0.1797962188720703, + 0.671058714389801, + 0.31389155983924866, + 0.42466309666633606, + -0.395219624042511, + 0.07233688980340958 + ], + [ + -0.1798933446407318, + 0.6702919602394104, + 0.3141656517982483, + 0.42445075511932373, + -0.3947013020515442, + 0.07293850928544998 + ], + [ + -0.17997199296951294, + 0.6695549488067627, + 0.314453125, + 0.42422589659690857, + -0.3940390348434448, + 0.07336582988500595 + ], + [ + -0.17965377867221832, + 0.6687509417533875, + 0.3146941363811493, + 0.42395132780075073, + -0.3932723104953766, + 0.07322361320257187 + ], + [ + -0.17922039330005646, + 0.6680362820625305, + 0.31512850522994995, + 0.42367929220199585, + -0.39243340492248535, + 0.07309210300445557 + ], + [ + -0.17886574566364288, + 0.667279064655304, + 0.3152349293231964, + 0.4234037399291992, + -0.39180412888526917, + 0.07277241349220276 + ] + ] + }, + "Gradients": { + "fc1.weight": [ + [ + -0.18660534918308258, + 1.0501877069473267, + -0.06538727134466171, + 0.7892400622367859, + -0.06989894062280655, + 0.08311288058757782 + ], + [ + -0.14019422233104706, + -0.33581918478012085, + -0.015272594057023525, + -0.11933345347642899, + 0.15028853714466095, + 0.3035297095775604 + ], + [ + 0.014209908433258533, + 0.05260481685400009, + 0.09717129915952682, + 0.23888230323791504, + -0.05665421858429909, + -0.053932324051856995 + ], + [ + -0.31329306960105896, + 0.38529160618782043, + -0.5335826277732849, + 0.41899508237838745, + 0.11755916476249695, + -0.2919110357761383 + ], + [ + 0.5206083655357361, + 0.9029364585876465, + 0.4515741765499115, + 1.2402161359786987, + -0.5781408548355103, + -0.9771140217781067 + ], + [ + 0.10178368538618088, + 1.272549033164978, + -0.43885383009910583, + -1.229077935218811, + -0.24205324053764343, + -0.4146330654621124 + ], + [ + -0.007725818548351526, + 0.2862418293952942, + -0.08212745189666748, + 0.15442971885204315, + -0.46201223134994507, + 0.2059977948665619 + ], + [ + -0.7134706377983093, + 0.9803712964057922, + 0.02016977034509182, + 0.31701749563217163, + -0.6400585174560547, + 1.4056847095489502 + ], + [ + -0.3544383645057678, + 0.0006996551528573036, + -0.5025327205657959, + 0.11614283919334412, + -0.7705468535423279, + 0.019379474222660065 + ], + [ + 0.06727457791566849, + 0.7126177549362183, + 0.5345063805580139, + 0.13658584654331207, + 0.2790890336036682, + 0.8989311456680298 + ] + ] + }, + "Momentum1s": { + "fc1.weight": [ + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + [ + -0.018660536035895348, + 0.10501877218484879, + -0.006538727320730686, + 0.07892400771379471, + -0.0069898939691483974, + 0.008311288431286812 + ], + [ + -0.030813904479146004, + 0.06093497574329376, + -0.007412114180624485, + 0.0590982623398304, + 0.008737949654459953, + 0.03783313184976578 + ], + [ + -0.026311522349715233, + 0.060101959854364395, + 0.0030462276190519333, + 0.07707666605710983, + 0.0021987329237163067, + 0.02865658327937126 + ], + [ + -0.05500967800617218, + 0.09262092411518097, + -0.050616659224033356, + 0.11126850545406342, + 0.01373477652668953, + -0.0034001797903329134 + ], + [ + 0.002552127931267023, + 0.17365248501300812, + -0.00039757348713465035, + 0.2241632640361786, + -0.0454527884721756, + -0.10077156871557236 + ], + [ + 0.012475283816456795, + 0.2835421562194824, + -0.04424320161342621, + 0.07883913069963455, + -0.06511283665895462, + -0.13215771317481995 + ], + [ + 0.010455173440277576, + 0.28381210565567017, + -0.04803162440657616, + 0.08639819175004959, + -0.10480277240276337, + -0.0983421579003334 + ], + [ + -0.06193741038441658, + 0.35346800088882446, + -0.041211485862731934, + 0.10946012288331985, + -0.1583283543586731, + 0.052060529589653015 + ], + [ + -0.0911875069141388, + 0.31819117069244385, + -0.08734361082315445, + 0.11012839525938034, + -0.21955019235610962, + 0.04879242181777954 + ], + [ + -0.07534129917621613, + 0.3576337993144989, + -0.025158612057566643, + 0.11277413368225098, + -0.16968625783920288, + 0.1338062882423401 + ] + ] + }, + "Momentum2s": { + "fc1.weight": [ + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + [ + 3.4821558074327186e-05, + 0.001102894195355475, + 4.275495484762359e-06, + 0.0006228999118320644, + 4.885862381343031e-06, + 6.90775095790741e-06 + ], + [ + 5.444115595309995e-05, + 0.0012145658256486058, + 4.504472144617466e-06, + 0.0006365174776874483, + 2.7467622203403153e-05, + 9.903113095788285e-05 + ], + [ + 5.458863961393945e-05, + 0.0012161185732111335, + 1.3942229998065159e-05, + 0.000692945730406791, + 3.0649855034425855e-05, + 0.00010184079292230308 + ], + [ + 0.00015268661081790924, + 0.0013633521739393473, + 0.00029863871168345213, + 0.000867809634655714, + 4.44393626821693e-05, + 0.0001869510015239939 + ], + [ + 0.0004235670203343034, + 0.0021772831678390503, + 0.000502259295899421, + 0.0024050779175013304, + 0.00037864179466851056, + 0.0011415159096941352 + ], + [ + 0.0004335033881943673, + 0.003794487100094557, + 0.0006943496991880238, + 0.0039133052341639996, + 0.0004368529189378023, + 0.0013122950913384557 + ], + [ + 0.000433129578595981, + 0.0038726271595805883, + 0.0007004002691246569, + 0.003933240193873644, + 0.0006498713628388941, + 0.0013534179888665676 + ], + [ + 0.0009417368564754725, + 0.004829882178455591, + 0.0007001066696830094, + 0.004029807168990374, + 0.0010588964214548469, + 0.00332801416516304 + ], + [ + 0.0010664216242730618, + 0.004825052805244923, + 0.0009519456652924418, + 0.004039266612380743, + 0.0016515799798071384, + 0.0033250616397708654 + ], + [ + 0.0010698811383917928, + 0.005328051745891571, + 0.0012366907903924584, + 0.004053883254528046, + 0.001727819093503058, + 0.004129813984036446 + ] + ] + } +} \ No newline at end of file diff --git a/onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_single_weight_mode_1.json b/onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_single_weight_mode_1.json new file mode 100644 index 0000000000..b277311d80 --- /dev/null +++ b/onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_single_weight_mode_1.json @@ -0,0 +1,362 @@ +{ + "Parameters": { + "fc1.weight": [ + [ + -0.18330414593219757, + 0.6739549040794373, + 0.31170889735221863, + 0.42830976843833923, + -0.3957911729812622, + 0.07424858212471008 + ], + [ + -0.18230432271957397, + 0.6729549169540405, + 0.31270840764045715, + 0.4273098111152649, + -0.3947916328907013, + 0.0732489600777626 + ], + [ + -0.18132172524929047, + 0.672543466091156, + 0.3135298192501068, + 0.42675861716270447, + -0.39518389105796814, + 0.07235442847013474 + ], + [ + -0.18060243129730225, + 0.672195315361023, + 0.31336504220962524, + 0.42616716027259827, + -0.3952641189098358, + 0.07178084552288055 + ], + [ + -0.1797843873500824, + 0.6717343330383301, + 0.313903272151947, + 0.42547306418418884, + -0.3956426680088043, + 0.07182653248310089 + ], + [ + -0.17980578541755676, + 0.6710923910140991, + 0.31390631198883057, + 0.4246846139431, + -0.3952397406101227, + 0.07234101742506027 + ], + [ + -0.17990471422672272, + 0.6703324317932129, + 0.3141835033893585, + 0.4244765043258667, + -0.39472541213035583, + 0.07294334471225739 + ], + [ + -0.1799851655960083, + 0.6696021556854248, + 0.31447410583496094, + 0.4242558777332306, + -0.39406710863113403, + 0.07337138056755066 + ], + [ + -0.17966875433921814, + 0.6688048243522644, + 0.3147182762622833, + 0.4239855408668518, + -0.39330434799194336, + 0.07322990149259567 + ], + [ + -0.17923718690872192, + 0.6680968403816223, + 0.3151557743549347, + 0.4237177073955536, + -0.3924693763256073, + 0.07309911400079727 + ], + [ + -0.178884357213974, + 0.6673462986946106, + 0.31526532769203186, + 0.4234463572502136, + -0.3918440341949463, + 0.07278015464544296 + ] + ] + }, + "Gradients": { + "fc1.weight": [ + [ + -0.18660534918308258, + 1.0501877069473267, + -0.06538727134466171, + 0.7892400622367859, + -0.06989894062280655, + 0.08311288058757782 + ], + [ + -0.14019371569156647, + -0.3358178436756134, + -0.015272140502929688, + -0.11933225393295288, + 0.15028847754001617, + 0.3035295903682709 + ], + [ + 0.014210130088031292, + 0.05260540917515755, + 0.09717151522636414, + 0.23888282477855682, + -0.05665425583720207, + -0.05393238365650177 + ], + [ + -0.313301146030426, + 0.385299414396286, + -0.5335837602615356, + 0.4189963638782501, + 0.11755432933568954, + -0.29190653562545776 + ], + [ + 0.5206307172775269, + 0.9029856324195862, + 0.4515964090824127, + 1.2402637004852295, + -0.5781452655792236, + -0.9771224856376648 + ], + [ + 0.10178495943546295, + 1.2725839614868164, + -0.43884608149528503, + -1.229053020477295, + -0.2420593947172165, + -0.41463127732276917 + ], + [ + -0.007730415090918541, + 0.28626471757888794, + -0.08212656527757645, + 0.15444251894950867, + -0.4620162546634674, + 0.2060021609067917 + ], + [ + -0.7135099768638611, + 0.9804390668869019, + 0.020160868763923645, + 0.3170405924320221, + -0.6400777697563171, + 1.4057106971740723 + ], + [ + -0.35445013642311096, + 0.0007170308963395655, + -0.5025206804275513, + 0.11615195125341415, + -0.7705658674240112, + 0.019383220002055168 + ], + [ + 0.06728999316692352, + 0.7126797437667847, + 0.5345433354377747, + 0.13664238154888153, + 0.27906644344329834, + 0.898922860622406 + ] + ] + }, + "Momentum1s": { + "fc1.weight": [ + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + [ + -0.018660536035895348, + 0.10501877218484879, + -0.006538727320730686, + 0.07892400771379471, + -0.0069898939691483974, + 0.008311288431286812 + ], + [ + -0.030813854187726974, + 0.06093510985374451, + -0.007412068545818329, + 0.05909838154911995, + 0.008737944066524506, + 0.03783312067389488 + ], + [ + -0.02631145343184471, + 0.06010213866829872, + 0.003046290250495076, + 0.07707682996988297, + 0.0021987236104905605, + 0.028656570240855217 + ], + [ + -0.05501042306423187, + 0.09262186288833618, + -0.05061671510338783, + 0.1112687811255455, + 0.013734283857047558, + -0.003399740904569626 + ], + [ + 0.002553692553192377, + 0.17365823686122894, + -0.00039540237048640847, + 0.22416827082633972, + -0.045453671365976334, + -0.10077201575040817 + ], + [ + 0.012476819567382336, + 0.28355079889297485, + -0.044240470975637436, + 0.07884613424539566, + -0.06511424481868744, + -0.13215793669223785 + ], + [ + 0.010456095449626446, + 0.2838221788406372, + -0.04802908003330231, + 0.08640576899051666, + -0.10480444133281708, + -0.0983419269323349 + ], + [ + -0.0619405135512352, + 0.3534838855266571, + -0.04121008515357971, + 0.10946924984455109, + -0.1583317667245865, + 0.052063338458538055 + ], + [ + -0.09119147807359695, + 0.3182072043418884, + -0.08734114468097687, + 0.11013751477003098, + -0.21955516934394836, + 0.04879532381892204 + ], + [ + -0.07534332573413849, + 0.3576544523239136, + -0.025152696296572685, + 0.11278799921274185, + -0.1696930080652237, + 0.13380807638168335 + ] + ] + }, + "Momentum2s": { + "fc1.weight": [ + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + [ + 3.4821558074327186e-05, + 0.001102894195355475, + 4.275495484762359e-06, + 0.0006228999118320644, + 4.885862381343031e-06, + 6.90775095790741e-06 + ], + [ + 5.4441014071926475e-05, + 0.0012145648943260312, + 4.504458047449589e-06, + 0.0006365172448568046, + 2.7467604013509117e-05, + 9.903105819830671e-05 + ], + [ + 5.458850137074478e-05, + 0.0012161176418885589, + 1.3942257282906212e-05, + 0.000692945730406791, + 3.0649840482510626e-05, + 0.00010184072743868455 + ], + [ + 0.00015269152936525643, + 0.0013633571797981858, + 0.00029863996314816177, + 0.0008678107406012714, + 4.443821308086626e-05, + 0.0001869483239715919 + ], + [ + 0.0004235951928421855, + 0.002177376998588443, + 0.0005022806581109762, + 0.002405197126790881, + 0.0003786457236856222, + 0.0011415297631174326 + ], + [ + 0.00043353179353289306, + 0.0037946696393191814, + 0.0006943642511032522, + 0.003913363441824913, + 0.000436859845649451, + 0.0013123073149472475 + ], + [ + 0.00043315801303833723, + 0.003872822504490614, + 0.0007004146464169025, + 0.003933302592486143, + 0.0006498820148408413, + 0.0013534319587051868 + ], + [ + 0.0009418213739991188, + 0.004830210469663143, + 0.0007001206977292895, + 0.004029884003102779, + 0.0010589316952973604, + 0.0033281012438237667 + ], + [ + 0.0010665145236998796, + 0.004825380630791187, + 0.0009519476443529129, + 0.004039345309138298, + 0.0016516445903107524, + 0.0033251489512622356 + ], + [ + 0.0010699760168790817, + 0.005328468047082424, + 0.0012367323506623507, + 0.004053977318108082, + 0.0017278711311519146, + 0.004129886161535978 + ] + ] + } +} \ No newline at end of file diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc index bfd7d2e01f..21a1a2b681 100644 --- a/orttraining/orttraining/core/graph/training_op_defs.cc +++ b/orttraining/orttraining/core/graph/training_op_defs.cc @@ -153,14 +153,13 @@ bool SCELossInternalFunBuilder( const FunctionBodyBuildContext& ctx, const OpSchema& schema, FunctionProto& functionProto) { - bool hasWeight = ctx.hasInput(2); bool hasIgnoreIndex = ctx.hasInput(3); FunctionBuilder builder(functionProto); - builder // - .Const("Shape3D", std::vector({0, 0, -1})) // + builder + .Const("Shape3D", std::vector({0, 0, -1})) .Add(R"( X_NCD = Reshape (scores, Shape3D) X_NDC = Transpose (X_NCD) @@ -179,11 +178,10 @@ bool SCELossInternalFunBuilder( builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 (X_Log, labels, weights, ignore_index)"); else builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 (X_Log, labels, weights)"); + else if (hasIgnoreIndex) + builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 (X_Log, labels, , ignore_index)"); else - if (hasIgnoreIndex) - builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 (X_Log, labels, , ignore_index)"); - else - builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 (X_Log, labels)"); + builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 (X_Log, labels)"); schema.BuildFunction(functionProto); return true; @@ -380,10 +378,10 @@ bool BuildNllLossInternalFunctionHelper( if (opset_version > 12) body.push_back( - {{"const_one_64"}, - "Constant", - {}, - {MakeAttribute("value", ToDimensionOneTensor(int64_t(1)))}}); + {{"const_one_64"}, + "Constant", + {}, + {MakeAttribute("value", ToDimensionOneTensor(int64_t(1)))}}); body.push_back( {{"expanded_target"}, @@ -1265,6 +1263,130 @@ void RegisterTrainingOpSchemas() { {"tensor(bool)"}, "Constrain types to boolean tensors."); + /** + * AdamWOptimizer operator, taking multiple parameters as inputs (seq). + * Ideally, a group of parameters sharing same learning rate (or other meta data) can use one single AdamWOptimizer. + * Implementation-wise, this bring opportunities for achieving better performance. + * + * The differences with AdamOptimizer: + * > Different inputs. + * + * AdamWOptimizer can accept multiple parameters and other states related to them as inputs (seq). + * This make multi-tensor-apply applicable to the GPU implementation. Existing LambOptimizer has similar + * capability, while it is using many fixed-length optional vardaric inputs, which is not a clean op definition. + * + * AdamOptimizer takes one single parameter and its other states. + * + * > Different computation. + * + * Despite of normal adam computation, for better perf in ORTTrainer, AdamOptimizer also fused gradient norm + * clipping in its implementation. This sometimes makes it hard to align the optimizer with other frameworks during + * model onboarding, on the other hand, the fusion did not bring very significant gains actually. + * + * AdamWOptimizer has simplified definitions, excludes inputs/attributes not related to optimizer computations. + * + * AdamWOptimizer is recommended for new usage, AdamOptimizer is left as it is to support existing ORTTrainer + * solutions. + */ + ONNX_CONTRIB_OPERATOR_SCHEMA(AdamWOptimizer) + .SetDomain(kMSDomain) + .SinceVersion(1) + .Input(0, "lr", "The learning rate.", "T1") + .Input(1, "step", "The update count of weights. It should be a scalar.", "T2") + .Input(2, "weights", "Sequence of weights to optimize.", "S_WEIGHT") + .Input(3, "gradients", "Sequence of gradients computed in this iteration.", "S_GRAD") + .Input(4, "momentums_1", "Sequence of exponentially averaged historical gradients.", "S_MOMENT") + .Input(5, "momentums_2", "Sequence of exponentially averaged historical squared gradients.", "S_MOMENT") + .Input(6, "update_signal", + "This signal indicates if weight updates are skipped, applicable to gradient infinity check" + " in mixed precision training. ", + "T_BOOL", OpSchema::Optional) + .Output(0, "updated_flag", "Whether gradient is applied or not.", "T2") + .Output(1, "updated_weights", "Sequence of weights after optimize.", "S_WEIGHT", OpSchema::Optional) + .Output(2, "updated_momentums_1", "Sequence of momentum_1 after optimize.", "S_MOMENT", OpSchema::Optional) + .Output(3, "updated_momentums_2", "Sequence of momentum_2 after optimize.", "S_MOMENT", OpSchema::Optional) + .Attr( + "alpha", + "Coefficient of previously accumulated gradient in running average.", + AttributeProto::FLOAT, + 0.9f) + .Attr( + "beta", + "Coefficient of previously accumulated squared-gradient in running average.", + AttributeProto::FLOAT, + 0.999f) + .Attr( + "epsilon", + "Small scalar to avoid dividing by zero.", + AttributeProto::FLOAT, + 1e-8f) + .Attr( + "weight_decay", + "weight decay coefficient.", + AttributeProto::FLOAT, + 1e-2f) + .Attr( + "correct_bias", + "Whether or not to correct bias, enabled by default.", + AttributeProto::INT, + static_cast(1)) + .Attr( + "adam_mode", + "Modes for applying bias correction and weight decay (default 0) " + "0 : Weight decay is applied before weight is updated." + " Computation aligned with Torch AdamW. In this mode, " + " correct_bias should be 1 to keep aligned with PyTorch." + "1 : Weight decay is applied after weight is updated." + " Computation is aligned with Huggingface AdamW.", + AttributeProto::INT, + static_cast(0)) + .TypeConstraint( + "T1", + {"tensor(float)"}, + "Constrain learning rate to float") + .TypeConstraint( + "T2", + {"tensor(int64)"}, + "Constrain step count to 64-bit integer") + .TypeConstraint( + "S_WEIGHT", + {"seq(tensor(float16))", "seq(tensor(float))", "seq(tensor(double))"}, + "Constrain weights' types.") + .TypeConstraint( + "S_GRAD", + {"seq(tensor(float16))", "seq(tensor(float))", "seq(tensor(double))"}, + "Constrain gradients' types.") + .TypeConstraint( + "S_MOMENT", + {"seq(tensor(float16))", "seq(tensor(float))", "seq(tensor(double))"}, + "Constrain momentums' types.") + .TypeConstraint( + "T_BOOL", + {"tensor(bool)"}, + "Constrain types to boolean tensors.") + .TypeAndShapeInferenceFunction([](InferenceContext& ctx) { + size_t num_of_outputs = ctx.getNumOutputs(); + std::unordered_map output_to_input_index_map{{0, 1}, {1, 2}, {2, 4}, {3, 5}}; + assert(output_to_input_index_map.size() >= num_of_outputs); + + size_t sequence_source_input_index = 0; // Be noted: 0 is invalid for sequence source input index + for (size_t output_index = 0; output_index < num_of_outputs; ++output_index) { + auto& input_index = output_to_input_index_map[output_index]; + propagateElemTypeFromInputToOutput(ctx, input_index, output_index); + + // All 3 sequence inputs/outputs should have same shapes, searched for the first available shape + // and use it to infer output shapes. + if (output_index > 0 && sequence_source_input_index == 0 && hasInputShape(ctx, input_index)) { + sequence_source_input_index = input_index; + } + } + + for (size_t output_index = 1; sequence_source_input_index > 1 && output_index < num_of_outputs; + ++output_index) { + propagateShapeFromInputToOutput(ctx, sequence_source_input_index, output_index); + } + }); + ONNX_CONTRIB_OPERATOR_SCHEMA_ELSEWHERE(LambOptimizer, RegisterLambOpSchema); ONNX_CONTRIB_OPERATOR_SCHEMA(InPlaceAccumulator) @@ -1666,9 +1788,9 @@ Example 4: propagateShapeFromInputToOutput(ctx, 1, 0); }) .SetContextDependentFunctionBodyBuilder( - [](const FunctionBodyBuildContext& ctx, const OpSchema& schema, FunctionProto& functionProto) { - return SCELossGradFunBuilder (true, ctx, schema, functionProto); - }) + [](const FunctionBodyBuildContext& ctx, const OpSchema& schema, FunctionProto& functionProto) { + return SCELossGradFunBuilder(true, ctx, schema, functionProto); + }) .SetDoc(R"DOC(SoftmaxCrossEntropyLossGrad)DOC"); ONNX_CONTRIB_OPERATOR_SCHEMA(ReduceSumTraining) @@ -3659,7 +3781,6 @@ Return true if all elements are true and false otherwise. .SetContextDependentFunctionBodyBuilder(SCELossInternalFunBuilder) .SetDoc(R"DOC(SoftmaxCrossEntropyLossInternal)DOC"); - ONNX_CONTRIB_OPERATOR_SCHEMA(SoftmaxCrossEntropyLossInternalGrad) .SetDomain(kMSDomain) .SinceVersion(1) @@ -3685,9 +3806,9 @@ Return true if all elements are true and false otherwise. propagateShapeFromInputToOutput(ctx, 1, 0); }) .SetContextDependentFunctionBodyBuilder( - [](const FunctionBodyBuildContext& ctx, const OpSchema& schema, FunctionProto& functionProto) { - return SCELossGradFunBuilder (false, ctx, schema, functionProto); - }) + [](const FunctionBodyBuildContext& ctx, const OpSchema& schema, FunctionProto& functionProto) { + return SCELossGradFunBuilder(false, ctx, schema, functionProto); + }) .SetDoc(R"DOC(SoftmaxCrossEntropyLossInternalGrad)DOC"); ONNX_CONTRIB_OPERATOR_SCHEMA(NegativeLogLikelihoodLossInternal) @@ -3716,7 +3837,7 @@ Return true if all elements are true and false otherwise. .TypeAndShapeInferenceFunction([](InferenceContext& ctx) { propagateElemTypeFromInputToOutput(ctx, 0, 0); }) .SetDoc(R"DOC(NegativeLogLikelihoodLossInternal)DOC"); - ONNX_CONTRIB_OPERATOR_SCHEMA(NegativeLogLikelihoodLossInternal2) + ONNX_CONTRIB_OPERATOR_SCHEMA(NegativeLogLikelihoodLossInternal2) .SetDomain(kMSDomain) .SinceVersion(1) .Attr("reduction", reduction_doc, AttributeProto::STRING, std::string("mean")) diff --git a/orttraining/orttraining/test/training_ops/cuda/optimizer/adamw_test.cc b/orttraining/orttraining/test/training_ops/cuda/optimizer/adamw_test.cc new file mode 100644 index 0000000000..d02e25be93 --- /dev/null +++ b/orttraining/orttraining/test/training_ops/cuda/optimizer/adamw_test.cc @@ -0,0 +1,321 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "gtest/gtest.h" + +#include "nlohmann/json.hpp" +#include "test/providers/provider_test_utils.h" +#include "orttraining/test/training_ops/cuda/optimizer/common.h" + +namespace onnxruntime { +namespace test { +namespace optimizer { + +#if USE_CUDA +using json = nlohmann::json; +namespace { + +// key: weight name; value: multiple steps of weight/grad/momentums values. +typedef std::unordered_map>> WeightDictType; + +constexpr const char* kParamName = "Parameters"; +constexpr const char* kGradientName = "Gradients"; +constexpr const char* kMomentum1Name = "Momentum1s"; +constexpr const char* kMomentum2Name = "Momentum2s"; + +// key: name of data, e.g. one of kParamName, kGradientName, etc. +// value: weight dicts. +typedef std::unordered_map TestDataDictType; + +#define ADAM_TEST_DATA_FOLDER ORT_TSTR("testdata/test_data_generation/adamw_test/") + +void TorchAdamWSingleWeightTestLoop10Steps(bool use_baseline_inputs_for_each_iteration, bool* update_signal = nullptr) { + size_t total_step = 10; + float lr = 1e-03f; + + std::pair weight_tolerance{1e-4f, 1e-5f}; // rtol, atol + std::pair momentum1_tolerance{1e-3f, 1e-6f}; + std::pair momentum2_tolerance{1e-2f, 1e-7f}; + + if (!use_baseline_inputs_for_each_iteration) { + // Loose the tolerance as all states are maintained (without reloading from baseline) across different steps. + momentum2_tolerance.first = 1e-3f; + momentum2_tolerance.second = 1e-6f; + } + + auto data_uri = ADAM_TEST_DATA_FOLDER "adamw_test_single_weight_mode_0.json"; + std::ifstream in{data_uri}; + + TestDataDictType test_data; + const json j = json::parse(in); + j.get_to(test_data); + + // 11 steps of weight values before applying optimization. + WeightDictType& named_weights = test_data[kParamName]; + + // 10 steps of gradient values used to apply optimization. + WeightDictType& named_gradients = test_data[kGradientName]; + + // 11 steps of momentum1 values before applying optimization. + WeightDictType& named_momentum1s = test_data[kMomentum1Name]; + + // 11 steps of momentum2 values before applying optimization. + WeightDictType& named_momentum2s = test_data[kMomentum2Name]; + + ASSERT_EQ(named_weights.size(), 1); + ASSERT_EQ(named_gradients.size(), 1); + ASSERT_EQ(named_momentum1s.size(), 1); + ASSERT_EQ(named_momentum2s.size(), 1); + + ASSERT_EQ(named_weights["fc1.weight"].size(), total_step + 1); + ASSERT_EQ(named_gradients["fc1.weight"].size(), total_step); + ASSERT_EQ(named_momentum1s["fc1.weight"].size(), total_step + 1); + ASSERT_EQ(named_momentum2s["fc1.weight"].size(), total_step + 1); + + std::unordered_map weight_name_shape_mapping = + {{"fc1.weight", {2, 3}}}; + + AdamWTestLoop(use_baseline_inputs_for_each_iteration, total_step, lr, + static_cast(0.9f), // alpha + static_cast(0.999f), // beta + static_cast(1e-8f), // epsilon + static_cast(1e-2f), // weight_decay + static_cast(0), // adam_mode + static_cast(1), // correct_bias + named_weights, named_gradients, + named_momentum1s, named_momentum2s, + weight_name_shape_mapping, + weight_tolerance, + momentum1_tolerance, + momentum2_tolerance, + update_signal); +} + +TEST(AdamWTest, TorchAdamWSingleWeightTest_Loop10Steps) { + TorchAdamWSingleWeightTestLoop10Steps(true); +} + +TEST(AdamWTest, TorchAdamWSingleWeightStrictTest_Loop10Steps) { + TorchAdamWSingleWeightTestLoop10Steps(false); +} + +TEST(AdamWTest, TorchAdamWSingleWeightNoUpdateTest_Loop10Steps) { + bool update_signal = false; + TorchAdamWSingleWeightTestLoop10Steps(true, &update_signal); +} + +void TorchAdamWMultipleWeightsTestLoop10Steps(bool use_baseline_inputs_for_each_iteration, + bool* update_signal = nullptr) { + size_t total_step = 10; + float lr = 1e-03f; + + std::pair weight_tolerance{1e-4f, 1e-5f}; // rtol, atol + std::pair momentum1_tolerance{1e-3f, 1e-6f}; + std::pair momentum2_tolerance{1e-2f, 1e-7f}; + + if (!use_baseline_inputs_for_each_iteration) { + // Loose the tolerance as all states are maintained (without reloading from baseline) across different steps. + momentum2_tolerance.first = 1e-3f; + momentum2_tolerance.second = 1e-6f; + } + + auto data_uri = ADAM_TEST_DATA_FOLDER "adamw_test_multiple_weights_mode_0.json"; + std::ifstream in{data_uri}; + + TestDataDictType test_data; + const json j = json::parse(in); + j.get_to(test_data); + + // 11 steps of weight values before applying optimization. + WeightDictType& named_weights = test_data[kParamName]; + + // 10 steps of gradient values used to apply optimization. + WeightDictType& named_gradients = test_data[kGradientName]; + + // 11 steps of momentum1 values before applying optimization. + WeightDictType& named_momentum1s = test_data[kMomentum1Name]; + + // 11 steps of momentum2 values before applying optimization. + WeightDictType& named_momentum2s = test_data[kMomentum2Name]; + + ASSERT_EQ(named_weights.size(), 4); + ASSERT_EQ(named_gradients.size(), 4); + ASSERT_EQ(named_momentum1s.size(), 4); + ASSERT_EQ(named_momentum2s.size(), 4); + + ASSERT_EQ(named_weights["fc1.weight"].size(), total_step + 1); + ASSERT_EQ(named_gradients["fc1.weight"].size(), total_step); + ASSERT_EQ(named_momentum1s["fc1.weight"].size(), total_step + 1); + ASSERT_EQ(named_momentum2s["fc1.weight"].size(), total_step + 1); + + std::unordered_map weight_name_shape_mapping = + {{"fc1.weight", {2, 3}}, {"fc1.bias", {3}}, {"fc2.weight", {3, 2}}, {"fc2.bias", {2}}}; + + AdamWTestLoop(use_baseline_inputs_for_each_iteration, total_step, lr, + static_cast(0.9f), // alpha + static_cast(0.999f), // beta + static_cast(1e-8f), // epsilon + static_cast(1e-2f), // weight_decay + static_cast(0), // adam_mode + static_cast(1), // correct_bias + named_weights, named_gradients, + named_momentum1s, named_momentum2s, + weight_name_shape_mapping, + weight_tolerance, + momentum1_tolerance, + momentum2_tolerance, + update_signal); +} + +TEST(AdamWTest, TorchAdamWMultipleWeightsTest_Loop10Steps) { + TorchAdamWMultipleWeightsTestLoop10Steps(true); +} + +TEST(AdamWTest, TorchAdamWMultipleWeightsStrictTest_Loop10Steps) { + TorchAdamWMultipleWeightsTestLoop10Steps(false); +} + +TEST(AdamWTest, TorchAdamWMultipleWeightsNoUpdateTest_Loop10Steps) { + bool update_signal = false; + TorchAdamWMultipleWeightsTestLoop10Steps(true, &update_signal); +} + +void HFAdamWSingleWeightTestLoop10Steps(bool use_baseline_inputs_for_each_iteration) { + size_t total_step = 10; + float lr = 1e-03f; + + std::pair weight_tolerance{1e-4f, 1e-5f}; // rtol, atol + std::pair momentum1_tolerance{1e-3f, 1e-6f}; + std::pair momentum2_tolerance{1e-2f, 1e-7f}; + + auto data_uri = ADAM_TEST_DATA_FOLDER "adamw_test_single_weight_mode_1.json"; + std::ifstream in{data_uri}; + + TestDataDictType test_data; + const json j = json::parse(in); + j.get_to(test_data); + + // 11 steps of weight values before applying optimization. + WeightDictType& named_weights = test_data[kParamName]; + + // 10 steps of gradient values used to apply optimization. + WeightDictType& named_gradients = test_data[kGradientName]; + + // 11 steps of momentum1 values before applying optimization. + WeightDictType& named_momentum1s = test_data[kMomentum1Name]; + + // 11 steps of momentum2 values before applying optimization. + WeightDictType& named_momentum2s = test_data[kMomentum2Name]; + + ASSERT_EQ(named_weights.size(), 1); + ASSERT_EQ(named_gradients.size(), 1); + ASSERT_EQ(named_momentum1s.size(), 1); + ASSERT_EQ(named_momentum2s.size(), 1); + + ASSERT_EQ(named_weights["fc1.weight"].size(), total_step + 1); + ASSERT_EQ(named_gradients["fc1.weight"].size(), total_step); + ASSERT_EQ(named_momentum1s["fc1.weight"].size(), total_step + 1); + ASSERT_EQ(named_momentum2s["fc1.weight"].size(), total_step + 1); + + std::unordered_map weight_name_shape_mapping = + {{"fc1.weight", {2, 3}}}; + + AdamWTestLoop(use_baseline_inputs_for_each_iteration, total_step, lr, + static_cast(0.9f), // alpha + static_cast(0.999f), // beta + static_cast(1e-6f), // epsilon + static_cast(0.0f), // weight_decay + static_cast(1), // adam_mode + static_cast(1), // correct_bias + named_weights, named_gradients, + named_momentum1s, named_momentum2s, + weight_name_shape_mapping, + weight_tolerance, + momentum1_tolerance, + momentum2_tolerance); +} + +TEST(AdamWTest, HFAdamWSingleWeightTest_Loop10Steps) { + HFAdamWSingleWeightTestLoop10Steps(false); +} + +TEST(AdamWTest, HFAdamWSingleWeightStrictTest_Loop10Steps) { + HFAdamWSingleWeightTestLoop10Steps(true); +} + +void HFAdamWMultipleWeightsTestLoop10Steps( + bool use_baseline_inputs_for_each_iteration) { + size_t total_step = 10; + float lr = 1e-03f; + + std::pair weight_tolerance{1e-4f, 1e-5f}; // rtol, atol + std::pair momentum1_tolerance{1e-3f, 1e-6f}; + std::pair momentum2_tolerance{1e-2f, 1e-7f}; + + if (!use_baseline_inputs_for_each_iteration) { + // Loose the tolerance as all states are maintained (without reloading from baseline) across different steps. + momentum2_tolerance.first = 1e-3f; + momentum2_tolerance.second = 1e-6f; + } + auto data_uri = ADAM_TEST_DATA_FOLDER "adamw_test_multiple_weights_mode_1.json"; + std::ifstream in{data_uri}; + + TestDataDictType test_data; + const json j = json::parse(in); + j.get_to(test_data); + + // 11 steps of weight values before applying optimization. + WeightDictType& named_weights = test_data[kParamName]; + + // 10 steps of gradient values used to apply optimization. + WeightDictType& named_gradients = test_data[kGradientName]; + + // 11 steps of momentum1 values before applying optimization. + WeightDictType& named_momentum1s = test_data[kMomentum1Name]; + + // 11 steps of momentum2 values before applying optimization. + WeightDictType& named_momentum2s = test_data[kMomentum2Name]; + + ASSERT_EQ(named_weights.size(), 4); + ASSERT_EQ(named_gradients.size(), 4); + ASSERT_EQ(named_momentum1s.size(), 4); + ASSERT_EQ(named_momentum2s.size(), 4); + + ASSERT_EQ(named_weights["fc1.weight"].size(), total_step + 1); + ASSERT_EQ(named_gradients["fc1.weight"].size(), total_step); + ASSERT_EQ(named_momentum1s["fc1.weight"].size(), total_step + 1); + ASSERT_EQ(named_momentum2s["fc1.weight"].size(), total_step + 1); + + std::unordered_map weight_name_shape_mapping = + {{"fc1.weight", {2, 3}}, {"fc1.bias", {3}}, {"fc2.weight", {3, 2}}, {"fc2.bias", {2}}}; + + AdamWTestLoop(use_baseline_inputs_for_each_iteration, total_step, lr, + static_cast(0.9f), // alpha + static_cast(0.999f), // beta + static_cast(1e-6f), // epsilon + static_cast(0.0f), // weight_decay + static_cast(1), // adam_mode + static_cast(1), // correct_bias + named_weights, named_gradients, + named_momentum1s, named_momentum2s, + weight_name_shape_mapping, + weight_tolerance, + momentum1_tolerance, + momentum2_tolerance); +} + +TEST(AdamWTest, HFAdamWMultipleWeightsTest_Loop10Steps) { + HFAdamWMultipleWeightsTestLoop10Steps(false); +} + +TEST(AdamWTest, HFAdamWMultipleWeightsStrictTest_Loop10Steps) { + HFAdamWMultipleWeightsTestLoop10Steps(true); +} + +} // namespace + +#endif // USE_CUDA + +} // namespace optimizer +} // namespace test +} // namespace onnxruntime diff --git a/orttraining/orttraining/test/training_ops/cuda/optimizer/common.cc b/orttraining/orttraining/test/training_ops/cuda/optimizer/common.cc new file mode 100644 index 0000000000..cbc465f337 --- /dev/null +++ b/orttraining/orttraining/test/training_ops/cuda/optimizer/common.cc @@ -0,0 +1,243 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "test/providers/provider_test_utils.h" +#include "orttraining/test/training_ops/cuda/optimizer/common.h" + +namespace onnxruntime { +namespace test { +namespace optimizer { + +TensorInfo::TensorInfo(const VectorInt64& shapes, const std::vector& values) { + shapes_ = shapes; + fp32_values_ = values; + + size_t total_size = 1; + for (size_t i = 0; i < shapes_.size(); ++i) { + total_size *= shapes_[i]; + } + + EXPECT_TRUE(fp32_values_.size() == total_size) << "Number of elements mismatch between shapes and values." + << "fp32_values_.size():" << fp32_values_.size() + << ", total_size: " << total_size; +} + +template +AdamTestInputOutput::AdamTestInputOutput( + float lr, + int64_t step, + const std::vector& weight_tensor_infos, + const std::vector& gradient_tensor_infos, + const std::vector& momentum_1_tensor_infos, + const std::vector& momentum_2_tensor_infos, + const std::vector& updated_weight_tensor_infos, + const std::vector& updated_momentum_1_tensor_infos, + const std::vector& updated_momentum_2_tensor_infos) { + lr_vector.push_back(lr); + step_vector.push_back(step); + + // Input Sequence tensors. + + for (const TensorInfo& ti : weight_tensor_infos) { + weight_seq_tensors_.AddTensor(ti.Shapes(), ti.Values()); + } + + for (const TensorInfo& ti : gradient_tensor_infos) { + gradient_seq_tensors_.AddTensor(ti.Shapes(), ti.Values()); + } + + for (const TensorInfo& ti : momentum_1_tensor_infos) { + momentum_1_seq_tensors_.AddTensor(ti.Shapes(), ti.Values()); + } + + for (const TensorInfo& ti : momentum_2_tensor_infos) { + momentum_2_seq_tensors_.AddTensor(ti.Shapes(), ti.Values()); + } + + // Update sequence tensors. + + for (const TensorInfo& ti : updated_weight_tensor_infos) { + updated_weight_seq_tensors_.AddTensor(ti.Shapes(), ti.Values()); + } + + for (const TensorInfo& ti : updated_momentum_1_tensor_infos) { + updated_momentum_1_seq_tensors_.AddTensor(ti.Shapes(), ti.Values()); + } + + for (const TensorInfo& ti : updated_momentum_2_tensor_infos) { + updated_momentum_2_seq_tensors_.AddTensor(ti.Shapes(), ti.Values()); + } +} + +void GetPerStepInput( + const std::unordered_map& weight_name_shape_mapping, + std::unordered_map>>& named_weights, + std::unordered_map>>& named_momentum1s, + std::unordered_map>>& named_momentum2s, + size_t step, + std::unordered_map>& weights_to_train, + std::unordered_map>& momentum1_to_train, + std::unordered_map>& momentum2_to_train) { + weights_to_train.clear(); + momentum1_to_train.clear(); + momentum2_to_train.clear(); + + for (auto it = weight_name_shape_mapping.begin(); it != weight_name_shape_mapping.end(); ++it) { + const std::string& weight_name = it->first; + + ASSERT_TRUE(weights_to_train.find(weight_name) == weights_to_train.end()); + weights_to_train.insert({weight_name, named_weights[weight_name][step]}); + + ASSERT_TRUE(momentum1_to_train.find(weight_name) == momentum1_to_train.end()); + momentum1_to_train.insert({weight_name, named_momentum1s[weight_name][step]}); + + ASSERT_TRUE(momentum2_to_train.find(weight_name) == momentum2_to_train.end()); + momentum2_to_train.insert({weight_name, named_momentum2s[weight_name][step]}); + } +} + +void AdamWTestLoop( + bool use_baseline_inputs_for_each_iteration, size_t total_step, float lr, + float alpha, float beta, float epsilon, float weight_decay, int64_t adam_mode, int64_t correct_bias, + std::unordered_map>>& named_weights, + std::unordered_map>>& named_gradients, + std::unordered_map>>& named_momentums_1, + std::unordered_map>>& named_momentums_2, + std::unordered_map& weight_name_shape_mapping, + std::pair weight_tolerance, + std::pair momentum_1_tolerance, + std::pair momentum_2_tolerance, + bool* update_signal) { + std::vector ordered_weight_names; + for (auto it = weight_name_shape_mapping.begin(); it != weight_name_shape_mapping.end(); ++it) { + const std::string& weight_name = it->first; + ASSERT_TRUE( + std::find(ordered_weight_names.begin(), ordered_weight_names.end(), weight_name) == ordered_weight_names.end()); + ordered_weight_names.push_back(weight_name); + } + + std::unordered_map> weights_to_train; + std::unordered_map> momentum1_to_train; + std::unordered_map> momentum2_to_train; + GetPerStepInput(weight_name_shape_mapping, named_weights, named_momentums_1, named_momentums_2, + 0, weights_to_train, momentum1_to_train, momentum2_to_train); + + for (size_t step = 0; step < 1; ++step) { + OpTester test("AdamWOptimizer", 1, onnxruntime::kMSDomain); + + // Update the steps for each param group update. + // Both torch and HF increase training step before applying gradients. + // The test alignes with them. + int64_t increased_update_count = step + 1; + + // Weights/momentums before applying optimization. + std::vector weight_tensor_infos; + std::vector momentum1_tensor_infos; + std::vector momentum2_tensor_infos; + std::vector gradient_tensor_infos; + + // Updated weights/momentums values for validation. + std::vector updated_weight_tensor_infos; + std::vector updated_momentum1_tensor_infos; + std::vector updated_momentum2_tensor_infos; + for (auto& weight_name : ordered_weight_names) { + VectorInt64 weight_shape = weight_name_shape_mapping[weight_name]; + weight_tensor_infos.emplace_back(TensorInfo(weight_shape, weights_to_train[weight_name])); + momentum1_tensor_infos.emplace_back(TensorInfo(weight_shape, momentum1_to_train[weight_name])); + momentum2_tensor_infos.emplace_back(TensorInfo(weight_shape, momentum2_to_train[weight_name])); + gradient_tensor_infos.emplace_back(TensorInfo(weight_shape, named_gradients[weight_name][step])); + + updated_weight_tensor_infos.emplace_back(TensorInfo(weight_shape, named_weights[weight_name][step + 1])); + updated_momentum1_tensor_infos.emplace_back(TensorInfo(weight_shape, named_momentums_1[weight_name][step + 1])); + updated_momentum2_tensor_infos.emplace_back(TensorInfo(weight_shape, named_momentums_2[weight_name][step + 1])); + } + + AdamTestInputOutput data( + lr, increased_update_count, weight_tensor_infos, gradient_tensor_infos, momentum1_tensor_infos, + momentum2_tensor_infos, updated_weight_tensor_infos, updated_momentum1_tensor_infos, + updated_momentum2_tensor_infos); + + test.AddAttribute("alpha", alpha); + test.AddAttribute("beta", beta); + test.AddAttribute("epsilon", epsilon); + test.AddAttribute("weight_decay", weight_decay); + test.AddAttribute("adam_mode", adam_mode); + test.AddAttribute("correct_bias", correct_bias); + + // Add test inputs. + test.AddInput("lr", {}, data.lr_vector); + test.AddInput("step", {}, data.step_vector); + test.AddSeqInput("weights", data.WeightSeq()); + test.AddSeqInput("gradients", data.GradientSeq()); + test.AddSeqInput("momentums_1", data.Momentum_1_Seq()); + test.AddSeqInput("momentums_2", data.Momentum_2_Seq()); + if (update_signal != nullptr) { + test.AddInput("update_signal", {}, {*update_signal}); + } + + // Add test outputs as baseline. + if (update_signal == nullptr || *update_signal) { + test.AddOutput("updated_flag", {}, {1}); + test.AddSeqOutput("updated_weights", data.UpdatedWeightSeq(), weight_tolerance.first, weight_tolerance.second); + test.AddSeqOutput("updated_momentums_1", data.UpdatedMomentum_1_Seq(), momentum_1_tolerance.first, + momentum_1_tolerance.second); + test.AddSeqOutput("updated_momentums_2", data.UpdatedMomentum_2_Seq(), momentum_2_tolerance.first, + momentum_2_tolerance.second); + + } else { + // No update happens. + test.AddOutput("updated_flag", {}, {0}); + test.AddSeqOutput("updated_weights", data.WeightSeq(), weight_tolerance.first, weight_tolerance.second); + test.AddSeqOutput("updated_momentums_1", data.Momentum_1_Seq(), momentum_1_tolerance.first, + momentum_1_tolerance.second); + test.AddSeqOutput("updated_momentums_2", data.Momentum_2_Seq(), momentum_2_tolerance.first, + momentum_2_tolerance.second); + } + + test.Run(); + + if (use_baseline_inputs_for_each_iteration) { + GetPerStepInput(weight_name_shape_mapping, named_weights, named_momentums_1, named_momentums_2, + step + 1, weights_to_train, momentum1_to_train, momentum2_to_train); + } else { + std::vector outputs = test.GetFetches(); + ASSERT_EQ(outputs.size(), 4); + + const TensorSeq& updated_seq_weight = outputs[1].Get(); + const TensorSeq& updated_seq_momentum1 = outputs[2].Get(); + const TensorSeq& updated_seq_momentum2 = outputs[3].Get(); + + size_t weight_index = 0; + for (auto& weight_name : ordered_weight_names) { + ASSERT_TRUE(weights_to_train.find(weight_name) != weights_to_train.end()); + ASSERT_TRUE(momentum1_to_train.find(weight_name) != momentum1_to_train.end()); + ASSERT_TRUE(momentum2_to_train.find(weight_name) != momentum2_to_train.end()); + + const float* updated_weight_buffer = updated_seq_weight.Get(weight_index).Data(); + std::copy(updated_weight_buffer, updated_weight_buffer + weights_to_train[weight_name].size(), + weights_to_train[weight_name].begin()); + + const float* updated_momentum1_buffer = updated_seq_momentum1.Get(weight_index).Data(); + std::copy(updated_momentum1_buffer, updated_momentum1_buffer + momentum1_to_train[weight_name].size(), + momentum1_to_train[weight_name].begin()); + + const float* updated_momentum2_buffer = updated_seq_momentum2.Get(weight_index).Data(); + std::copy(updated_momentum2_buffer, updated_momentum2_buffer + momentum2_to_train[weight_name].size(), + momentum2_to_train[weight_name].begin()); + + weight_index += 1; + } + } + } +} +} // namespace optimizer +} // namespace test +} // namespace onnxruntime diff --git a/orttraining/orttraining/test/training_ops/cuda/optimizer/common.h b/orttraining/orttraining/test/training_ops/cuda/optimizer/common.h new file mode 100644 index 0000000000..ab797ed189 --- /dev/null +++ b/orttraining/orttraining/test/training_ops/cuda/optimizer/common.h @@ -0,0 +1,116 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "test/providers/provider_test_utils.h" + +namespace onnxruntime { +namespace test { +namespace optimizer { + +struct TensorInfo { + TensorInfo(const VectorInt64& shapes, const std::vector& values); + + template + std::vector Values() const { + if (std::is_same::value) { + std::vector fp16_values; + fp16_values.reserve(fp32_values_.size()); + ConvertFloatToMLFloat16(fp32_values_.data(), + reinterpret_cast(fp16_values.data()), + static_cast(fp32_values_.size())); + return fp16_values; + } else if (std::is_same::value) { + return std::vector(fp32_values_); + } else { + ORT_THROW("Not supported data type."); + } + } + + VectorInt64 Shapes() const { + return shapes_; + } + + VectorInt64 shapes_; + std::vector fp32_values_; +}; + +template +struct AdamTestInputOutput { + AdamTestInputOutput( + float lr, + int64_t step, + const std::vector& weight_tensor_infos, + const std::vector& gradient_tensor_infos, + const std::vector& momentum_1_tensor_infos, + const std::vector& momentum_2_tensor_infos, + const std::vector& updated_weight_tensor_infos, + const std::vector& updated_momentum_1_tensor_infos, + const std::vector& updated_momentum_2_tensor_infos); + + SeqTensors& WeightSeq() { + return weight_seq_tensors_; + } + + SeqTensors& GradientSeq() { + return gradient_seq_tensors_; + } + + SeqTensors& Momentum_1_Seq() { + return momentum_1_seq_tensors_; + } + + SeqTensors& Momentum_2_Seq() { + return momentum_2_seq_tensors_; + } + + SeqTensors& UpdatedWeightSeq() { + return updated_weight_seq_tensors_; + } + + SeqTensors& UpdatedMomentum_1_Seq() { + return updated_momentum_1_seq_tensors_; + } + + SeqTensors& UpdatedMomentum_2_Seq() { + return updated_momentum_2_seq_tensors_; + } + + std::vector lr_vector; + std::vector step_vector; + + private: + SeqTensors weight_seq_tensors_; + SeqTensors gradient_seq_tensors_; + SeqTensors momentum_1_seq_tensors_; + SeqTensors momentum_2_seq_tensors_; + + SeqTensors updated_weight_seq_tensors_; + SeqTensors updated_momentum_1_seq_tensors_; + SeqTensors updated_momentum_2_seq_tensors_; +}; + +void AdamWTestLoop( + bool use_baseline_inputs_for_each_iteration, size_t total_step, float lr, + float alpha, float beta, float epsilon, float weight_decay, int64_t adam_mode, int64_t correct_bias, + std::unordered_map>>& named_weights, + std::unordered_map>>& named_gradients, + std::unordered_map>>& named_momentums_1, + std::unordered_map>>& named_momentums_2, + std::unordered_map& weight_name_shape_mapping, + std::pair weight_tolerance, + std::pair momentum_1_tolerance, + std::pair momentum_2_tolerance, + bool* update_signal = nullptr); + +} // namespace optimizer +} // namespace test +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cpu/optimizer/common.h b/orttraining/orttraining/training_ops/cpu/optimizer/common.h index b08f75c8dd..8bb2584f86 100644 --- a/orttraining/orttraining/training_ops/cpu/optimizer/common.h +++ b/orttraining/orttraining/training_ops/cpu/optimizer/common.h @@ -20,4 +20,5 @@ TC compute_bias_correction_coefficient( } } // namespace contrib -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime + diff --git a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc index 769f7d73dc..6faf6edd40 100644 --- a/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc +++ b/orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc @@ -31,6 +31,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_int64_t_float_MLFloat16_MLFloat16_float_MLFloat16, AdamOptimizer); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_MLFloat16_MLFloat16_MLFloat16_MLFloat16, AdamOptimizer); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_MLFloat16_MLFloat16_float_MLFloat16, AdamOptimizer); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, AdamWOptimizer); // Lamb class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_float_float_float_float_MLFloat16, LambOptimizer); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_float_MLFloat16_float_MLFloat16_MLFloat16, LambOptimizer); @@ -256,6 +257,7 @@ Status RegisterCudaTrainingKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, // Lamb BuildKernelCreateInfo, diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/adam_impl.cu b/orttraining/orttraining/training_ops/cuda/optimizer/adam_impl.cu index 95ee9d2dae..47c0daa157 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/adam_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/optimizer/adam_impl.cu @@ -275,3 +275,4 @@ SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, BFloat16, BFloat16, float, } // namespace cuda } // namespace onnxruntime + diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw.cc b/orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw.cc new file mode 100644 index 0000000000..3470aac034 --- /dev/null +++ b/orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw.cc @@ -0,0 +1,158 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include + +#include "orttraining/training_ops/cuda/optimizer/adamw/adamw.h" +#include "orttraining/training_ops/cuda/optimizer/adamw/adamw_impl.h" + +namespace onnxruntime { +namespace cuda { + +ONNX_OPERATOR_KERNEL_EX( + AdamWOptimizer, + kMSDomain, + 1, + kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .InputMemoryType(OrtMemTypeCPUInput, 0) + .InputMemoryType(OrtMemTypeCPUInput, 1) + .InputMemoryType(OrtMemTypeCPUInput, 6) + .OutputMemoryType(OrtMemTypeCPUOutput, 0) + .Alias(2, 1) /* Return updated weights in-place */ + .Alias(4, 2) /* Return updated moment-1 in-place */ + .Alias(5, 3) /* Return updated moment-2 in-place */ + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) + .TypeConstraint("S_WEIGHT", DataTypeImpl::AllFixedSizeSequenceTensorTypes()) + .TypeConstraint("S_GRAD", DataTypeImpl::AllFixedSizeSequenceTensorTypes()) + .TypeConstraint("S_MOMENT", DataTypeImpl::AllFixedSizeSequenceTensorTypes()), + AdamWOptimizer); + +AllocatorPtr CreateAllocatorPtr(OpKernelContext* ctx) { + AllocatorPtr alloc; + ORT_ENFORCE(ctx->GetTempSpaceAllocator(&alloc).IsOK(), + "AdamWOptimizer GPU: Unable to get an allocator."); + return alloc; +} + +AllocatorPtr& GetAllocatorPtr(OpKernelContext* ctx) { + static AllocatorPtr alloc = CreateAllocatorPtr(ctx); + return alloc; +} + +Status GenerateOutputs(OpKernelContext* ctx, cudaStream_t stream, + const TensorSeq* values, TensorSeq* updated_values, + size_t number_of_values) { + // Return if the output edge is not fetched. + if (updated_values == nullptr) { + return Status::OK(); + } + + bool is_same_buffer = const_cast(values) == updated_values; + if (!is_same_buffer) { + updated_values->SetType(values->DataType()); + updated_values->Reserve(number_of_values); + for (size_t input_idx = 0; input_idx < number_of_values; ++input_idx) { + const Tensor& source_tensor = values->Get(input_idx); + std::unique_ptr target_tensor = Tensor::Create(source_tensor.DataType(), + source_tensor.Shape(), GetAllocatorPtr(ctx)); + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(target_tensor->MutableDataRaw(), + source_tensor.DataRaw(), + source_tensor.SizeInBytes(), + cudaMemcpyDeviceToDevice, stream)); + updated_values->Add(std::move(*target_tensor)); // Add will check for type consistency + } + } + + return Status::OK(); +} + +Status AdamWOptimizer::ComputeInternal(OpKernelContext* ctx) const { + const Tensor* learning_rate = ctx->Input(0); + const Tensor* step = ctx->Input(1); + const TensorSeq* weights = ctx->Input(2); + const TensorSeq* gradients = ctx->Input(3); + const TensorSeq* momentums_1 = ctx->Input(4); + const TensorSeq* momentums_2 = ctx->Input(5); + + size_t num_of_weights = weights->Size(); + size_t num_of_gradients = gradients->Size(); + size_t num_of_momentums_1 = momentums_1->Size(); + size_t num_of_momentums_2 = momentums_2->Size(); + + // Check the number of weights, gradients, momentums matchs. + ORT_ENFORCE(num_of_weights > 0, "Invalid count of tensors in Seq."); + ORT_ENFORCE(num_of_weights == num_of_gradients, "Number of weights and gradients mismatch."); + ORT_ENFORCE(num_of_gradients == num_of_momentums_1, "Number of gradients and momentums_1 mismatch."); + ORT_ENFORCE(num_of_momentums_1 == num_of_momentums_2, "Number of momentums_1 and momentums_2 mismatch."); + + std::vector tensor_sizes(num_of_weights); + std::vector> grouped_tensor_pointers(num_of_weights); + + for (size_t i = 0; i < num_of_weights; ++i) { + const Tensor& weight_tensor = weights->Get(i); + const Tensor& gradient_tensor = gradients->Get(i); + const Tensor& momentum_1_tensor = momentums_1->Get(i); + const Tensor& momentum_2_tensor = momentums_2->Get(i); + + // Check the weight/gradient/momentums at the same index should have same shape. + ORT_ENFORCE(weight_tensor.Shape() == gradient_tensor.Shape(), + "Shape of weight and gradient mismatch, weight index:", i); + ORT_ENFORCE(gradient_tensor.Shape() == momentum_1_tensor.Shape(), + "Shape of gradient and momentum_1 mismatch, weight index:", i); + ORT_ENFORCE(momentum_1_tensor.Shape() == momentum_2_tensor.Shape(), + "Shape of momentum_1 and momentum_2 mismatch, weight index:", i); + + // Currently we only support float data types. + ORT_ENFORCE(weight_tensor.IsDataType() && + gradient_tensor.IsDataType() && + momentum_1_tensor.IsDataType() && + momentum_2_tensor.IsDataType(), + "Only float data type support for Adam cuda kernel."); + + tensor_sizes[i] = static_cast(weight_tensor.Shape().Size()); + + grouped_tensor_pointers[i] = { + const_cast(weight_tensor.Data()), + const_cast(gradient_tensor.Data()), + const_cast(momentum_1_tensor.Data()), + const_cast(momentum_2_tensor.Data())}; + } + + Tensor* updated_flag = ctx->Output(0, step->Shape()); + TensorSeq* updated_weights = ctx->Output(1); + TensorSeq* updated_momentums_1 = ctx->Output(2); + TensorSeq* updated_momentums_2 = ctx->Output(3); + + int64_t* updated_flag_ptr = updated_flag->template MutableData(); + + // Currently placed on CPU, need revisit when we had mixed precision training requirement. + const Tensor* update_signal = ctx->Input(6); + if (update_signal == nullptr || *update_signal->template Data()) { + typedef typename ToCudaType::MappedType CudaT_FLOAT; + typedef AdamWMTAFunctor TFunctor; + TFunctor functor; + + const float* lr_ptr = learning_rate->template Data(); + const int64_t* step_ptr = step->template Data(); + ORT_ENFORCE(lr_ptr && step_ptr); + + launch_multi_tensor_functor( + Stream(), MTA_ADAMW_CHUNK_SIZE, tensor_sizes, grouped_tensor_pointers, functor, + alpha_, beta_, epsilon_, *lr_ptr, weight_decay_, adam_mode_, correct_bias_, *step_ptr); + *updated_flag_ptr = 1; + } else { + *updated_flag_ptr = 0; + } + + ORT_RETURN_IF_ERROR(GenerateOutputs(ctx, Stream(), weights, updated_weights, num_of_weights)); + ORT_RETURN_IF_ERROR(GenerateOutputs(ctx, Stream(), momentums_1, updated_momentums_1, num_of_weights)); + ORT_RETURN_IF_ERROR(GenerateOutputs(ctx, Stream(), momentums_2, updated_momentums_2, num_of_weights)); + + return Status::OK(); +} + +} // namespace cuda +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw.h b/orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw.h new file mode 100644 index 0000000000..dc849870b6 --- /dev/null +++ b/orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw.h @@ -0,0 +1,45 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include + +#include "core/common/common.h" +#include "core/providers/cuda/cuda_kernel.h" + +namespace onnxruntime { +namespace cuda { + +class AdamWOptimizer final : public CudaKernel { + public: + AdamWOptimizer(const OpKernelInfo& info) : CudaKernel(info) { + info.GetAttrOrDefault("alpha", &alpha_, 0.9f); + info.GetAttrOrDefault("beta", &beta_, 0.999f); + info.GetAttrOrDefault("epsilon", &epsilon_, 1e-8f); + + info.GetAttrOrDefault("weight_decay", &weight_decay_, 0.f); + info.GetAttrOrDefault("adam_mode", &adam_mode_, static_cast(0)); + info.GetAttrOrDefault("correct_bias", &correct_bias_, static_cast(1)); + + ORT_ENFORCE(adam_mode_ == 0 || adam_mode_ == 1, "The value of adam_mode is invalid."); + ORT_ENFORCE(correct_bias_ == 0 || correct_bias_ == 1, "The value of correct_bias is invalid."); + + // To have torch adamw equivalence, correct_bias must be 1 for adam_mode=0. + ORT_ENFORCE(adam_mode_ != 0 || correct_bias_ == 1, "The correct_bias should be 1 for adam_mode = 0."); + } + + Status ComputeInternal(OpKernelContext* context) const override; + + private: + float alpha_; + float beta_; + float epsilon_; + + float weight_decay_; + int64_t adam_mode_; + int64_t correct_bias_; +}; + +} // namespace cuda +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw_impl.cu b/orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw_impl.cu new file mode 100644 index 0000000000..5314cc78ec --- /dev/null +++ b/orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw_impl.cu @@ -0,0 +1,222 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include + +#include "core/providers/cuda/cuda_common.h" +#include "core/providers/cuda/cu_inc/common.cuh" + +#include "orttraining/training_ops/cuda/optimizer/adamw/adamw_impl.h" +#include "orttraining/training_ops/cuda/optimizer/common.cuh" +#include "orttraining/training_ops/cuda/optimizer/common.h" + +namespace onnxruntime { +namespace cuda { + +template +__device__ void PrepareMTAData( + const ChunkGroup& chunks, + const int& block_idx, + T_WEIGHT*& weight_chunk_ptr, + T_GRAD*& grad_chunk_ptr, + T_MOMENTUM*& momentum_1_chunk_ptr, + T_MOMENTUM*& momentum_2_chunk_ptr, + int& chunk_size) { + const int tensor_idx = chunks.block_index_to_tensor_group_index[block_idx]; + const int tensor_size = chunks.tensor_sizes[tensor_idx]; + T_WEIGHT* weight_tensor_ptr = static_cast(chunks.tensor_ptrs[0][tensor_idx]); + T_GRAD* grad_tensor_ptr = static_cast(chunks.tensor_ptrs[1][tensor_idx]); + T_MOMENTUM* momentum_1_tensor_ptr = static_cast(chunks.tensor_ptrs[2][tensor_idx]); + T_MOMENTUM* momentum_2_tensor_ptr = static_cast(chunks.tensor_ptrs[3][tensor_idx]); + const int chunk_start_idx = chunks.block_index_to_chunk_start_index[block_idx]; + // chunk_size is chunks.chunk_size if the loaded chunk is full. Otherwise (this + // chunk is the last one in the source tensor), the actual size is determined + // by the bound of the source tensor. + chunk_size = min(tensor_size, chunk_start_idx + chunks.chunk_size) - chunk_start_idx; + + weight_chunk_ptr = weight_tensor_ptr + chunk_start_idx; + grad_chunk_ptr = grad_tensor_ptr + chunk_start_idx; + momentum_1_chunk_ptr = momentum_1_tensor_ptr + chunk_start_idx; + momentum_2_chunk_ptr = momentum_2_tensor_ptr + chunk_start_idx; +} + +// Torch Adam equivalence. +template +__global__ void AdamWComputeMode0( + ChunkGroup chunks, + const float alpha, + const float beta, + const float epsilon, + const float lr, + const float alpha_correction, + const float beta_correction, + const float decay) { + const int block_idx = blockIdx.x; + + T_WEIGHT* weight_chunk_ptr; + T_GRAD* grad_chunk_ptr; + T_MOMENTUM* momentum_1_chunk_ptr; + T_MOMENTUM* momentum_2_chunk_ptr; + + // TODO(pengwa): unroll this one for better perf. + int chunk_size; + + PrepareMTAData(chunks, block_idx, weight_chunk_ptr, grad_chunk_ptr, + momentum_1_chunk_ptr, momentum_2_chunk_ptr, chunk_size); + +#pragma unroll(4) + for (int i = threadIdx.x; i < chunk_size; i += blockDim.x) { + float w = static_cast(weight_chunk_ptr[i]); + float g = static_cast(grad_chunk_ptr[i]); + float m1 = static_cast(momentum_1_chunk_ptr[i]); + float m2 = static_cast(momentum_2_chunk_ptr[i]); + + // Perform weight decay. + w = w - (w * lr * decay); + + // Compute exponentially-averaged historical gradient. + m1 = alpha * m1 + (1.f - alpha) * g; + + // Compute exponentially-averaged historical squared gradient. + m2 = beta * m2 + (1.f - beta) * g * g; + + // Compute the new weight. + const float denom = (_Sqrt(m2) / _Sqrt(beta_correction)) + epsilon; + w = w - (lr * m1) / (alpha_correction * denom); + + // Update the new weight and momentums. + weight_chunk_ptr[i] = static_cast(w); + momentum_1_chunk_ptr[i] = static_cast(m1); + momentum_2_chunk_ptr[i] = static_cast(m2); + } +} + +// Huggingface AdamW equivalence. +template +__global__ void AdamWComputeMode1( + ChunkGroup chunks, + const float alpha, + const float beta, + const float epsilon, + const float lr, + const float lr_corrected, + const float decay) { + const int block_idx = blockIdx.x; + + T_WEIGHT* weight_chunk_ptr; + T_GRAD* grad_chunk_ptr; + T_MOMENTUM* momentum_1_chunk_ptr; + T_MOMENTUM* momentum_2_chunk_ptr; + int chunk_size; + + PrepareMTAData(chunks, block_idx, weight_chunk_ptr, grad_chunk_ptr, + momentum_1_chunk_ptr, momentum_2_chunk_ptr, chunk_size); + +#pragma unroll(4) + for (int i = threadIdx.x; i < chunk_size; i += blockDim.x) { + float w = static_cast(weight_chunk_ptr[i]); + float g = static_cast(grad_chunk_ptr[i]); + float m1 = static_cast(momentum_1_chunk_ptr[i]); + float m2 = static_cast(momentum_2_chunk_ptr[i]); + + // Compute exponentially-averaged historical gradient. + m1 = alpha * m1 + (1.f - alpha) * g; + + // Compute exponentially-averaged historical squared gradient. + m2 = beta * m2 + (1.f - beta) * g * g; + + float denom = _Sqrt(m2) + epsilon; + w = w - (lr_corrected * m1 / denom); + + // Perform weight decay. + w = w - (lr * decay * w); + + // Update the new weight and momentums. + weight_chunk_ptr[i] = static_cast(w); + momentum_1_chunk_ptr[i] = static_cast(m1); + momentum_2_chunk_ptr[i] = static_cast(m2); + } +} + +template +void AdamWMTAFunctor::operator()( + cudaStream_t stream, + ChunkGroup chunks, + const float alpha, + const float beta, + const float epsilon, + const float lr, + const float decay, + const int64_t adam_mode, + const int64_t correct_bias, + const int64_t update_count) { + const int block_count = chunks.chunk_count; + const int thread_count = ChunkGroup::thread_count_per_block; + + float alpha_correction = 1.f, beta_correction = 1.f; + float lr_corrected = lr; + if (correct_bias == 1) { + // Notes: + // > there is a minor difference compared with Apex's implementation, + // which uses double storing corrections before casting to float passing to kernels. + // > std::pow(float, int) return double since C++11, so we cast back to float. + alpha_correction = 1.f - static_cast(std::pow(alpha, update_count)); + beta_correction = 1.f - static_cast(std::pow(beta, update_count)); + lr_corrected *= std::sqrt(beta_correction) / alpha_correction; + } + + // Currently two kinds of AdamW supported: + // Mode 0: Pytorch https://pytorch.org/docs/stable/_modules/torch/optim/adamw.html#AdamW, + // bias correction is applied on m and v individually, + // weight decay is applied before weight is updated. + // Mode 1: Huggingface https://github.com/huggingface/transformers/blob/d91841315aab55cf1347f4eb59332858525fad0f/ + // src/transformers/optimization.py, + // bias correction is applied on learning rate, then use lr_corrected for subsequent computations. + // weight decay is applied after weight is updated. + if (adam_mode == 0) { + AdamWComputeMode0<<>>( + chunks, alpha, beta, epsilon, lr, alpha_correction, beta_correction, decay); + } else if (adam_mode == 1) { + AdamWComputeMode1<<>>( + chunks, alpha, beta, epsilon, lr, lr_corrected, decay); + } else { + ORT_THROW("Unsupported Adamw optimizer mode."); + } +} + +#define INSTANTIATE_ADAMMTA_FUNCTOR(T_WEIGHT, T_GRAD, T_MOMENTUM) \ + template void AdamWMTAFunctor::operator()( \ + cudaStream_t stream, \ + ChunkGroup chunks, \ + const float alpha, \ + const float beta, \ + const float epsilon, \ + const float lr, \ + const float decay, \ + const int64_t adam_mode, \ + const int64_t correct_bias, \ + const int64_t update_count); \ + \ + template __global__ void AdamWComputeMode0( \ + ChunkGroup chunks, \ + const float alpha, \ + const float beta, \ + const float epsilon, \ + const float lr, \ + const float alpha_correction, \ + const float beta_correction, \ + const float decay); \ + \ + template __global__ void AdamWComputeMode1( \ + ChunkGroup chunks, \ + const float alpha, \ + const float beta, \ + const float epsilon, \ + const float lr, \ + const float lr_corrected, \ + const float decay); + +INSTANTIATE_ADAMMTA_FUNCTOR(float, float, float); + +} // namespace cuda +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw_impl.h b/orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw_impl.h new file mode 100644 index 0000000000..f67a23033b --- /dev/null +++ b/orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw_impl.h @@ -0,0 +1,29 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/cuda/multi_tensor/common.cuh" + +namespace onnxruntime { +namespace cuda { + +#define MTA_ADAMW_GROUP_SIZE 4 +#define MTA_ADAMW_CHUNK_SIZE 2048 * 32 + +template +struct AdamWMTAFunctor { + void operator()(cudaStream_t stream, + ChunkGroup chunks, + const float alpha, + const float beta, + const float epsilon, + const float lr, + const float decay, + const int64_t adam_mode, + const int64_t correct_bias, + const int64_t update_count); +}; + +} // namespace cuda +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/lamb_impl.cu b/orttraining/orttraining/training_ops/cuda/optimizer/lamb_impl.cu index 683bb23910..420a8cd616 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/lamb_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/optimizer/lamb_impl.cu @@ -418,9 +418,15 @@ __global__ void LambMultiTensorUpdateImpl( const T2* r_norm = reinterpret_cast(chunk_group.tensor_ptrs[1][group_index]); const T2* w = reinterpret_cast(chunk_group.tensor_ptrs[2][group_index]) + chunk_start; const T3* d = reinterpret_cast(chunk_group.tensor_ptrs[3][group_index]) + chunk_start; - T2* w_new = chunk_group.tensor_ptrs[4][group_index] != nullptr ? reinterpret_cast(chunk_group.tensor_ptrs[4][group_index]) + chunk_start : nullptr; + T2* w_new = + chunk_group.tensor_ptrs[4][group_index] != nullptr + ? reinterpret_cast(chunk_group.tensor_ptrs[4][group_index]) + chunk_start + : nullptr; T3* g_new = chunk_group.tensor_ptrs[5][group_index] != nullptr ? reinterpret_cast(chunk_group.tensor_ptrs[5][group_index]) + chunk_start : nullptr; - T_MIXED_PRECISION_FP* w_mixed_precision_new = chunk_group.tensor_ptrs[6][group_index] != nullptr ? reinterpret_cast(chunk_group.tensor_ptrs[6][group_index]) + chunk_start : nullptr; + T_MIXED_PRECISION_FP* w_mixed_precision_new = + chunk_group.tensor_ptrs[6][group_index] != nullptr + ? reinterpret_cast(chunk_group.tensor_ptrs[6][group_index]) + chunk_start + : nullptr; for (int i = threadIdx.x; i < chunk_size && i + chunk_start < tensor_size; i += blockDim.x) { _LambUpdateRule( @@ -589,7 +595,9 @@ __launch_bounds__(ChunkGroup<4>::thread_count_per_block) } } -CudaKernel::CudaAsyncBuffer compute_tensor_range_and_lock(ChunkGroup<4> chunk_group, const CudaKernel& kernel) { +CudaKernel::CudaAsyncBuffer compute_tensor_range_and_lock( + ChunkGroup<4> chunk_group, + const CudaKernel& kernel) { const int num_blocks = chunk_group.chunk_count; // sync_range_and_lock is a struct consisting of (start_block, num_blocks, lock) for each tensor @@ -609,7 +617,12 @@ CudaKernel::CudaAsyncBuffer compute_tensor_rang } template -void LambMultiTensorReductionFunctor::operator()(cudaStream_t stream, ChunkGroup<4> chunk_group, const CudaKernel& kernel, void* reduction_buffer, size_t reduction_buffer_size) { +void LambMultiTensorReductionFunctor::operator()( + cudaStream_t stream, + ChunkGroup<4> chunk_group, + const CudaKernel& kernel, + void* reduction_buffer, + size_t reduction_buffer_size) { // thread count per block. constexpr int thread_count = ChunkGroup<4>::thread_count_per_block; // shared memory's size per block. @@ -629,12 +642,15 @@ void LambMultiTensorReductionFunctor::operator() TOut2* d_buffer = reinterpret_cast(w_buffer + num_blocks); auto sync_range_and_lock = compute_tensor_range_and_lock(chunk_group, kernel); - LambMultiTensorReductionImpl<<>>( - chunk_group, w_buffer, d_buffer, sync_range_and_lock.GpuPtr()); + LambMultiTensorReductionImpl + <<>>( + chunk_group, w_buffer, d_buffer, sync_range_and_lock.GpuPtr()); } -#define INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(TIn1, TIn2, TOut1, TOut2, TBuf) \ - template void LambMultiTensorReductionFunctor::operator()(cudaStream_t stream, ChunkGroup<4> chunk_group, const CudaKernel& kernel, void* reduction_buffer, size_t reduction_buffer_size); +#define INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(TIn1, TIn2, TOut1, TOut2, TBuf) \ + template void LambMultiTensorReductionFunctor::operator()( \ + cudaStream_t stream, ChunkGroup<4> chunk_group, const CudaKernel& kernel, void* reduction_buffer, \ + size_t reduction_buffer_size); INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(float, float, float, float, float) INSTANTIATE_LAMB_MULTI_TENSOR_REDUCTION_FUNCTOR(double, double, double, double, double)