saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

MTA AdamWOptimizer (#11506)

Browse source 
* 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
This commit is contained in:
pengwa 2022-05-27 19:52:04 +08:00 committed by GitHub
parent 02724c54ff
commit 44f7b1bf2c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
23 changed files with 4519 additions and 48 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
onnxruntime/core/providers/cuda/multi_tensor/common.cuh
View file

@ -81,6 +81,7 @@ void launch_multi_tensor_functor(
std::vector<std::vector<void*>>& 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<size_t>(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<size_t>(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) {

18
onnxruntime/core/providers/shared_library/provider_interfaces.h
View file

@ -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)

2
onnxruntime/core/providers/shared_library/provider_wrappedtypes.h
View file

@ -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); }

2
onnxruntime/core/session/provider_bridge_ort.cc
View file

@ -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 {

6
onnxruntime/test/providers/cpu/sequence/concat_from_sequence_op_test.cc
View file

@ -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<int64_t>("axis", 2);
test.AddAttribute<int64_t>("new_axis", 0); // concat mode
SeqTensors<int64_t> 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<int64_t>("axis", 0);
test.AddAttribute<int64_t>("new_axis", 0); // concat mode
SeqTensors<int64_t> input;

55
onnxruntime/test/providers/provider_test_utils.h
View file

@ -3,7 +3,9 @@
#pragma once
#include <utility>
#include <variant>
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include <gsl/gsl>
@ -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 <typename T>
void AddSeqOutput(const char* name, const SeqTensors<T>& seq_tensors) {
AddSeqData<T>(output_data_, name, &seq_tensors);
void AddSeqOutput(const char* name, const SeqTensors<T>& seq_tensors,
float rel_error = 0.0f, float abs_error = 0.0f) {
AddSeqData<T>(output_data_, name, &seq_tensors, false, rel_error, abs_error);
}
#if !defined(DISABLE_OPTIONAL_TYPE)
@ -450,8 +453,9 @@ class OpTester {
template <typename T>
void AddOptionalTypeSeqOutput(const char* name,
const SeqTensors<T>* seq_tensors) {
AddSeqData<T>(output_data_, name, seq_tensors, true);
const SeqTensors<T>* seq_tensors,
float rel_error = 0.0f, float abs_error = 0.0f) {
AddSeqData<T>(output_data_, name, seq_tensors, true, rel_error, abs_error);
}
#endif
@ -945,7 +949,8 @@ class OpTester {
template <typename T>
void AddSeqData(std::vector<Data>& data, const char* name,
const SeqTensors<T>* 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<TensorSeq> ptr;
SequenceTensorTypeProto<T> sequence_tensor_proto;
if (seq_tensors) {
auto num_tensors = seq_tensors->tensors.size();
std::vector<Tensor> tensors;
tensors.resize(num_tensors);
auto elem_type = DataTypeImpl::GetType<T>();
for (size_t i = 0; i < num_tensors; ++i) {
TensorShape shape{seq_tensors->tensors[i].shape};
auto values_count = static_cast<int64_t>(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<T>());
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<TensorSeq>(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<T> 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<float>(), optional<float>()));
optional<float> rel;
optional<float> 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<int64_t> GetDimsForProto(gsl::span<const int64_t> dims);

192
onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_data_generator.py vendored Normal file
View file

@ -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()

1117
onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_multiple_weights_mode_0.json vendored Normal file
View file

File diff suppressed because it is too large Load diff

1117
onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_multiple_weights_mode_1.json vendored Normal file
View file

File diff suppressed because it is too large Load diff

362
onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_single_weight_mode_0.json vendored Normal file
View file

@ -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
]
]
}
}

362
onnxruntime/test/testdata/test_data_generation/adamw_test/adamw_test_single_weight_mode_1.json vendored Normal file
View file

@ -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
]
]
}
}

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

@ -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<int64_t>({0, 0, -1})) //
builder
.Const("Shape3D", std::vector<int64_t>({0, 0, -1}))
.Add(R"(
X_NCD = Reshape (scores, Shape3D)
X_NDC = Transpose <perm = [0, 2, 1]> (X_NCD)
@ -179,11 +178,10 @@ bool SCELossInternalFunBuilder(
builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 <reduction : string = @reduction> (X_Log, labels, weights, ignore_index)");
else
builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 <reduction : string = @reduction> (X_Log, labels, weights)");
else if (hasIgnoreIndex)
builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 <reduction : string = @reduction> (X_Log, labels, , ignore_index)");
else
if (hasIgnoreIndex)
builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 <reduction : string = @reduction> (X_Log, labels, , ignore_index)");
else
builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 <reduction : string = @reduction> (X_Log, labels)");
builder.Add("output = com.microsoft.NegativeLogLikelihoodLossInternal2 <reduction : string = @reduction> (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<tensor>).
* 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<tensor>).
* 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<int64_t>(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<int64_t>(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<size_t, size_t> 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"))

321
orttraining/orttraining/test/training_ops/cuda/optimizer/adamw_test.cc Normal file
View file

@ -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<std::string, std::vector<std::vector<float>>> 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<std::string, WeightDictType> 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<float, float> weight_tolerance{1e-4f, 1e-5f}; // rtol, atol
std::pair<float, float> momentum1_tolerance{1e-3f, 1e-6f};
std::pair<float, float> 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<TestDataDictType>(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<std::string, VectorInt64> weight_name_shape_mapping =
{{"fc1.weight", {2, 3}}};
AdamWTestLoop(use_baseline_inputs_for_each_iteration, total_step, lr,
static_cast<float>(0.9f), // alpha
static_cast<float>(0.999f), // beta
static_cast<float>(1e-8f), // epsilon
static_cast<float>(1e-2f), // weight_decay
static_cast<int64_t>(0), // adam_mode
static_cast<int64_t>(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<float, float> weight_tolerance{1e-4f, 1e-5f}; // rtol, atol
std::pair<float, float> momentum1_tolerance{1e-3f, 1e-6f};
std::pair<float, float> 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<TestDataDictType>(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<std::string, VectorInt64> 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<float>(0.9f), // alpha
static_cast<float>(0.999f), // beta
static_cast<float>(1e-8f), // epsilon
static_cast<float>(1e-2f), // weight_decay
static_cast<int64_t>(0), // adam_mode
static_cast<int64_t>(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<float, float> weight_tolerance{1e-4f, 1e-5f}; // rtol, atol
std::pair<float, float> momentum1_tolerance{1e-3f, 1e-6f};
std::pair<float, float> 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<TestDataDictType>(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<std::string, VectorInt64> weight_name_shape_mapping =
{{"fc1.weight", {2, 3}}};
AdamWTestLoop(use_baseline_inputs_for_each_iteration, total_step, lr,
static_cast<float>(0.9f), // alpha
static_cast<float>(0.999f), // beta
static_cast<float>(1e-6f), // epsilon
static_cast<float>(0.0f), // weight_decay
static_cast<int64_t>(1), // adam_mode
static_cast<int64_t>(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<float, float> weight_tolerance{1e-4f, 1e-5f}; // rtol, atol
std::pair<float, float> momentum1_tolerance{1e-3f, 1e-6f};
std::pair<float, float> 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<TestDataDictType>(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<std::string, VectorInt64> 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<float>(0.9f), // alpha
static_cast<float>(0.999f), // beta
static_cast<float>(1e-6f), // epsilon
static_cast<float>(0.0f), // weight_decay
static_cast<int64_t>(1), // adam_mode
static_cast<int64_t>(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

243
orttraining/orttraining/test/training_ops/cuda/optimizer/common.cc Normal file
View file

@ -0,0 +1,243 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <algorithm>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
#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<float>& 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 <typename T>
AdamTestInputOutput<T>::AdamTestInputOutput(
float lr,
int64_t step,
const std::vector<TensorInfo>& weight_tensor_infos,
const std::vector<TensorInfo>& gradient_tensor_infos,
const std::vector<TensorInfo>& momentum_1_tensor_infos,
const std::vector<TensorInfo>& momentum_2_tensor_infos,
const std::vector<TensorInfo>& updated_weight_tensor_infos,
const std::vector<TensorInfo>& updated_momentum_1_tensor_infos,
const std::vector<TensorInfo>& 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<T>());
}
for (const TensorInfo& ti : gradient_tensor_infos) {
gradient_seq_tensors_.AddTensor(ti.Shapes(), ti.Values<T>());
}
for (const TensorInfo& ti : momentum_1_tensor_infos) {
momentum_1_seq_tensors_.AddTensor(ti.Shapes(), ti.Values<T>());
}
for (const TensorInfo& ti : momentum_2_tensor_infos) {
momentum_2_seq_tensors_.AddTensor(ti.Shapes(), ti.Values<T>());
}
// Update sequence tensors.
for (const TensorInfo& ti : updated_weight_tensor_infos) {
updated_weight_seq_tensors_.AddTensor(ti.Shapes(), ti.Values<T>());
}
for (const TensorInfo& ti : updated_momentum_1_tensor_infos) {
updated_momentum_1_seq_tensors_.AddTensor(ti.Shapes(), ti.Values<T>());
}
for (const TensorInfo& ti : updated_momentum_2_tensor_infos) {
updated_momentum_2_seq_tensors_.AddTensor(ti.Shapes(), ti.Values<T>());
}
}
void GetPerStepInput(
const std::unordered_map<std::string, VectorInt64>& weight_name_shape_mapping,
std::unordered_map<std::string, std::vector<std::vector<float>>>& named_weights,
std::unordered_map<std::string, std::vector<std::vector<float>>>& named_momentum1s,
std::unordered_map<std::string, std::vector<std::vector<float>>>& named_momentum2s,
size_t step,
std::unordered_map<std::string, std::vector<float>>& weights_to_train,
std::unordered_map<std::string, std::vector<float>>& momentum1_to_train,
std::unordered_map<std::string, std::vector<float>>& 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<std::string, std::vector<std::vector<float>>>& named_weights,
std::unordered_map<std::string, std::vector<std::vector<float>>>& named_gradients,
std::unordered_map<std::string, std::vector<std::vector<float>>>& named_momentums_1,
std::unordered_map<std::string, std::vector<std::vector<float>>>& named_momentums_2,
std::unordered_map<std::string, VectorInt64>& weight_name_shape_mapping,
std::pair<float, float> weight_tolerance,
std::pair<float, float> momentum_1_tolerance,
std::pair<float, float> momentum_2_tolerance,
bool* update_signal) {
std::vector<std::string> 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<std::string, std::vector<float>> weights_to_train;
std::unordered_map<std::string, std::vector<float>> momentum1_to_train;
std::unordered_map<std::string, std::vector<float>> 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<TensorInfo> weight_tensor_infos;
std::vector<TensorInfo> momentum1_tensor_infos;
std::vector<TensorInfo> momentum2_tensor_infos;
std::vector<TensorInfo> gradient_tensor_infos;
// Updated weights/momentums values for validation.
std::vector<TensorInfo> updated_weight_tensor_infos;
std::vector<TensorInfo> updated_momentum1_tensor_infos;
std::vector<TensorInfo> 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<float> 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<float>("lr", {}, data.lr_vector);
test.AddInput<int64_t>("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<bool>("update_signal", {}, {*update_signal});
}
// Add test outputs as baseline.
if (update_signal == nullptr || *update_signal) {
test.AddOutput<int64_t>("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<int64_t>("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<OrtValue> outputs = test.GetFetches();
ASSERT_EQ(outputs.size(), 4);
const TensorSeq& updated_seq_weight = outputs[1].Get<TensorSeq>();
const TensorSeq& updated_seq_momentum1 = outputs[2].Get<TensorSeq>();
const TensorSeq& updated_seq_momentum2 = outputs[3].Get<TensorSeq>();
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<float>();
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<float>();
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<float>();
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

116
orttraining/orttraining/test/training_ops/cuda/optimizer/common.h Normal file
View file

@ -0,0 +1,116 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
#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<float>& values);
template <typename OutT>
std::vector<OutT> Values() const {
if (std::is_same<OutT, MLFloat16>::value) {
std::vector<OutT> fp16_values;
fp16_values.reserve(fp32_values_.size());
ConvertFloatToMLFloat16(fp32_values_.data(),
reinterpret_cast<onnxruntime::MLFloat16*>(fp16_values.data()),
static_cast<int>(fp32_values_.size()));
return fp16_values;
} else if (std::is_same<OutT, float>::value) {
return std::vector<OutT>(fp32_values_);
} else {
ORT_THROW("Not supported data type.");
}
}
VectorInt64 Shapes() const {
return shapes_;
}
VectorInt64 shapes_;
std::vector<float> fp32_values_;
};
template <typename T>
struct AdamTestInputOutput {
AdamTestInputOutput(
float lr,
int64_t step,
const std::vector<TensorInfo>& weight_tensor_infos,
const std::vector<TensorInfo>& gradient_tensor_infos,
const std::vector<TensorInfo>& momentum_1_tensor_infos,
const std::vector<TensorInfo>& momentum_2_tensor_infos,
const std::vector<TensorInfo>& updated_weight_tensor_infos,
const std::vector<TensorInfo>& updated_momentum_1_tensor_infos,
const std::vector<TensorInfo>& updated_momentum_2_tensor_infos);
SeqTensors<T>& WeightSeq() {
return weight_seq_tensors_;
}
SeqTensors<T>& GradientSeq() {
return gradient_seq_tensors_;
}
SeqTensors<T>& Momentum_1_Seq() {
return momentum_1_seq_tensors_;
}
SeqTensors<T>& Momentum_2_Seq() {
return momentum_2_seq_tensors_;
}
SeqTensors<T>& UpdatedWeightSeq() {
return updated_weight_seq_tensors_;
}
SeqTensors<T>& UpdatedMomentum_1_Seq() {
return updated_momentum_1_seq_tensors_;
}
SeqTensors<T>& UpdatedMomentum_2_Seq() {
return updated_momentum_2_seq_tensors_;
}
std::vector<float> lr_vector;
std::vector<int64_t> step_vector;
private:
SeqTensors<T> weight_seq_tensors_;
SeqTensors<T> gradient_seq_tensors_;
SeqTensors<T> momentum_1_seq_tensors_;
SeqTensors<T> momentum_2_seq_tensors_;
SeqTensors<T> updated_weight_seq_tensors_;
SeqTensors<T> updated_momentum_1_seq_tensors_;
SeqTensors<T> 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<std::string, std::vector<std::vector<float>>>& named_weights,
std::unordered_map<std::string, std::vector<std::vector<float>>>& named_gradients,
std::unordered_map<std::string, std::vector<std::vector<float>>>& named_momentums_1,
std::unordered_map<std::string, std::vector<std::vector<float>>>& named_momentums_2,
std::unordered_map<std::string, VectorInt64>& weight_name_shape_mapping,
std::pair<float, float> weight_tolerance,
std::pair<float, float> momentum_1_tolerance,
std::pair<float, float> momentum_2_tolerance,
bool* update_signal = nullptr);
} // namespace optimizer
} // namespace test
} // namespace onnxruntime

3
orttraining/orttraining/training_ops/cpu/optimizer/common.h
View file

@ -20,4 +20,5 @@ TC compute_bias_correction_coefficient(
}
} // namespace contrib
} // namespace onnxruntime
} // namespace onnxruntime

2
orttraining/orttraining/training_ops/cuda/cuda_training_kernels.cc
View file

@ -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<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16_int64_t_float_MLFloat16_MLFloat16_float_MLFloat16, AdamOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_MLFloat16_MLFloat16_MLFloat16_MLFloat16, AdamOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_int64_t_float_MLFloat16_MLFloat16_float_MLFloat16, AdamOptimizer)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, AdamWOptimizer)>,
// Lamb
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float_float_float_float_float_MLFloat16, LambOptimizer)>,

1
orttraining/orttraining/training_ops/cuda/optimizer/adam_impl.cu
View file

@ -275,3 +275,4 @@ SPECIALIZED_AdamOptimizerImpl(float, int64_t, float, BFloat16, BFloat16, float,
} // namespace cuda
} // namespace onnxruntime

158
orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw.cc Normal file
View file

@ -0,0 +1,158 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <memory>
#include <utility>
#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<float>())
.TypeConstraint("T2", DataTypeImpl::GetTensorType<int64_t>())
.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<TensorSeq*>(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<Tensor> 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<Tensor>(0);
const Tensor* step = ctx->Input<Tensor>(1);
const TensorSeq* weights = ctx->Input<TensorSeq>(2);
const TensorSeq* gradients = ctx->Input<TensorSeq>(3);
const TensorSeq* momentums_1 = ctx->Input<TensorSeq>(4);
const TensorSeq* momentums_2 = ctx->Input<TensorSeq>(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<Tensor>.");
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<int> tensor_sizes(num_of_weights);
std::vector<std::vector<void*>> 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<float>() &&
gradient_tensor.IsDataType<float>() &&
momentum_1_tensor.IsDataType<float>() &&
momentum_2_tensor.IsDataType<float>(),
"Only float data type support for Adam cuda kernel.");
tensor_sizes[i] = static_cast<int>(weight_tensor.Shape().Size());
grouped_tensor_pointers[i] = {
const_cast<float*>(weight_tensor.Data<float>()),
const_cast<float*>(gradient_tensor.Data<float>()),
const_cast<float*>(momentum_1_tensor.Data<float>()),
const_cast<float*>(momentum_2_tensor.Data<float>())};
}
Tensor* updated_flag = ctx->Output(0, step->Shape());
TensorSeq* updated_weights = ctx->Output<TensorSeq>(1);
TensorSeq* updated_momentums_1 = ctx->Output<TensorSeq>(2);
TensorSeq* updated_momentums_2 = ctx->Output<TensorSeq>(3);
int64_t* updated_flag_ptr = updated_flag->template MutableData<int64_t>();
// Currently placed on CPU, need revisit when we had mixed precision training requirement.
const Tensor* update_signal = ctx->Input<Tensor>(6);
if (update_signal == nullptr || *update_signal->template Data<bool>()) {
typedef typename ToCudaType<float>::MappedType CudaT_FLOAT;
typedef AdamWMTAFunctor<CudaT_FLOAT, CudaT_FLOAT, CudaT_FLOAT> TFunctor;
TFunctor functor;
const float* lr_ptr = learning_rate->template Data<float>();
const int64_t* step_ptr = step->template Data<int64_t>();
ORT_ENFORCE(lr_ptr && step_ptr);
launch_multi_tensor_functor<MTA_ADAMW_GROUP_SIZE, TFunctor>(
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

45
orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw.h Normal file
View file

@ -0,0 +1,45 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <vector>
#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<int64_t>(0));
info.GetAttrOrDefault("correct_bias", &correct_bias_, static_cast<int64_t>(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

222
orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw_impl.cu Normal file
View file

@ -0,0 +1,222 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <algorithm>
#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 <typename T_WEIGHT, typename T_GRAD, typename T_MOMENTUM>
__device__ void PrepareMTAData(
const ChunkGroup<MTA_ADAMW_GROUP_SIZE>& 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<T_WEIGHT*>(chunks.tensor_ptrs[0][tensor_idx]);
T_GRAD* grad_tensor_ptr = static_cast<T_GRAD*>(chunks.tensor_ptrs[1][tensor_idx]);
T_MOMENTUM* momentum_1_tensor_ptr = static_cast<T_MOMENTUM*>(chunks.tensor_ptrs[2][tensor_idx]);
T_MOMENTUM* momentum_2_tensor_ptr = static_cast<T_MOMENTUM*>(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 <typename T_WEIGHT, typename T_GRAD, typename T_MOMENTUM>
__global__ void AdamWComputeMode0(
ChunkGroup<MTA_ADAMW_GROUP_SIZE> 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<float>(weight_chunk_ptr[i]);
float g = static_cast<float>(grad_chunk_ptr[i]);
float m1 = static_cast<float>(momentum_1_chunk_ptr[i]);
float m2 = static_cast<float>(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<T_WEIGHT>(w);
momentum_1_chunk_ptr[i] = static_cast<T_MOMENTUM>(m1);
momentum_2_chunk_ptr[i] = static_cast<T_MOMENTUM>(m2);
}
}
// Huggingface AdamW equivalence.
template <typename T_WEIGHT, typename T_GRAD, typename T_MOMENTUM>
__global__ void AdamWComputeMode1(
ChunkGroup<MTA_ADAMW_GROUP_SIZE> 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<float>(weight_chunk_ptr[i]);
float g = static_cast<float>(grad_chunk_ptr[i]);
float m1 = static_cast<float>(momentum_1_chunk_ptr[i]);
float m2 = static_cast<float>(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<T_WEIGHT>(w);
momentum_1_chunk_ptr[i] = static_cast<T_MOMENTUM>(m1);
momentum_2_chunk_ptr[i] = static_cast<T_MOMENTUM>(m2);
}
}
template <typename T_WEIGHT, typename T_GRAD, typename T_MOMENTUM>
void AdamWMTAFunctor<T_WEIGHT, T_GRAD, T_MOMENTUM>::operator()(
cudaStream_t stream,
ChunkGroup<MTA_ADAMW_GROUP_SIZE> 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<MTA_ADAMW_GROUP_SIZE>::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<float>(std::pow(alpha, update_count));
beta_correction = 1.f - static_cast<float>(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<T_WEIGHT, T_GRAD, T_MOMENTUM><<<block_count, thread_count, 0, stream>>>(
chunks, alpha, beta, epsilon, lr, alpha_correction, beta_correction, decay);
} else if (adam_mode == 1) {
AdamWComputeMode1<T_WEIGHT, T_GRAD, T_MOMENTUM><<<block_count, thread_count, 0, stream>>>(
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<T_WEIGHT, T_GRAD, T_MOMENTUM>::operator()( \
cudaStream_t stream, \
ChunkGroup<MTA_ADAMW_GROUP_SIZE> 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<T_WEIGHT, T_GRAD, T_MOMENTUM>( \
ChunkGroup<MTA_ADAMW_GROUP_SIZE> 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<T_WEIGHT, T_GRAD, T_MOMENTUM>( \
ChunkGroup<MTA_ADAMW_GROUP_SIZE> 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

29
orttraining/orttraining/training_ops/cuda/optimizer/adamw/adamw_impl.h Normal file
View file

@ -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 <typename T_WEIGHT, typename T_GRAD, typename T_MOMENTUM>
struct AdamWMTAFunctor {
void operator()(cudaStream_t stream,
ChunkGroup<MTA_ADAMW_GROUP_SIZE> 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

32
orttraining/orttraining/training_ops/cuda/optimizer/lamb_impl.cu
View file

@ -418,9 +418,15 @@ __global__ void LambMultiTensorUpdateImpl(
const T2* r_norm = reinterpret_cast<const T2*>(chunk_group.tensor_ptrs[1][group_index]);
const T2* w = reinterpret_cast<const T2*>(chunk_group.tensor_ptrs[2][group_index]) + chunk_start;
const T3* d = reinterpret_cast<const T3*>(chunk_group.tensor_ptrs[3][group_index]) + chunk_start;
T2* w_new = chunk_group.tensor_ptrs[4][group_index] != nullptr ? reinterpret_cast<T2*>(chunk_group.tensor_ptrs[4][group_index]) + chunk_start : nullptr;
T2* w_new =
chunk_group.tensor_ptrs[4][group_index] != nullptr
? reinterpret_cast<T2*>(chunk_group.tensor_ptrs[4][group_index]) + chunk_start
: nullptr;
T3* g_new = chunk_group.tensor_ptrs[5][group_index] != nullptr ? reinterpret_cast<T3*>(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<T_MIXED_PRECISION_FP*>(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<T_MIXED_PRECISION_FP*>(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<LambMultiTensorSyncRangeAndLock> compute_tensor_range_and_lock(ChunkGroup<4> chunk_group, const CudaKernel& kernel) {
CudaKernel::CudaAsyncBuffer<LambMultiTensorSyncRangeAndLock> 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<LambMultiTensorSyncRangeAndLock> compute_tensor_rang
}
template <typename TIn1, typename TIn2, typename TOut1, typename TOut2, typename TBuf>
void LambMultiTensorReductionFunctor<TIn1, TIn2, TOut1, TOut2, TBuf>::operator()(cudaStream_t stream, ChunkGroup<4> chunk_group, const CudaKernel& kernel, void* reduction_buffer, size_t reduction_buffer_size) {
void LambMultiTensorReductionFunctor<TIn1, TIn2, TOut1, TOut2, TBuf>::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<TIn1, TIn2, TOut1, TOut2, TBuf>::operator()
TOut2* d_buffer = reinterpret_cast<TOut2*>(w_buffer + num_blocks);
auto sync_range_and_lock = compute_tensor_range_and_lock(chunk_group, kernel);
LambMultiTensorReductionImpl<TIn1, TIn2, TOut1, TOut2, TBuf><<<chunk_group.chunk_count, thread_count, shared_memory_size, stream>>>(
chunk_group, w_buffer, d_buffer, sync_range_and_lock.GpuPtr());
LambMultiTensorReductionImpl<TIn1, TIn2, TOut1, TOut2, TBuf>
<<<chunk_group.chunk_count, thread_count, shared_memory_size, stream>>>(
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<TIn1, TIn2, TOut1, TOut2, TBuf>::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<TIn1, TIn2, TOut1, TOut2, TBuf>::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)