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

Fuse ScaledSum and its backward BatchScale (#16517)

Browse source 
### Fuse ScaledSum and its backward BatchScale

For deberta models, there is a pattern

a / scalar_0 + b / scalar_1 + c / scalar_2

We can fuse this into ScaledSum operator, taking 2(or 3) inputs, and
2(or 3) attributes scalar, generating one output.

For the backward, the gradient of a, b and c will be computed with
BatchScale.

### Benchmark on 8x32GV100

```bash
torchrun --nproc_per_node=8 examples/onnxruntime/training/language-modeling/run_mlm.py  --model_name_or_path microsoft/deberta-v3-large --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1  --num_train_epochs 10 --do_train  --overwrite_output_dir --output_dir ./outputs/ --seed 1137 --fp16 --report_to none --optim adamw_ort_fused  --max_steps 400 --logging_steps 1 --use_module_with_loss --deepspeed aml_ds_config_zero_1.json --per_device_train_batch_size 10
```

#### Main Branch

```
Total overhead: 127954ms where export takes 116489ms.
  epoch                    =      14.29
  train_loss               =     4.9803
  train_runtime            = 0:10:27.29
  train_samples            =       2223
  train_samples_per_second =     51.013
  train_steps_per_second   =      0.638


throughput per GPU = 14.29* 2223/ (627.29 - 127.954) / 8 (gpu) = 7.952 samples/second
```

#### This PR

```
Total overhead: 128761ms where export takes 118510ms.
***** train metrics *****
  epoch                    =      14.29
  train_loss               =     4.6144
  train_runtime            = 0:10:04.31
  train_samples            =       2223
  train_samples_per_second =     52.953
  train_steps_per_second   =      0.662

throughput per GPU = 14.29*2223 / (604.31 - 128.761) / 8 = 8.350 samples/second
```

5.x% performance gains.
This commit is contained in:
pengwa 2023-08-31 14:55:27 +08:00 committed by GitHub
parent c11ed065ba
commit 58af36b49a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
20 changed files with 1595 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

53
orttraining/orttraining/core/graph/gradient_builder.cc
View file

@ -6,6 +6,7 @@
#include <cmath>
#include <numeric>
#include <list>
#include <vector>
#include "onnx/defs/attr_proto_util.h"
#include "onnx/defs/tensor_proto_util.h"
@ -2087,5 +2088,57 @@ IMPLEMENT_GRADIENT_BUILDER(GetConvTransposeGradient) {
SrcNodeAttributes())};
}
IMPLEMENT_GRADIENT_BUILDER(GetScaledSumGradient) {
int input_count = GetSrcNodeInputSize();
auto attributes = SrcNodeAttributes();
float scale_0 = attributes.at("scale_0").f();
float scale_1 = attributes.at("scale_1").f();
if (input_count == 2) {
if (scale_0 == scale_1) {
// Specialized branch to avoid duplicated data write.
NodeDef scale_node = ConstantScalarNode(scale_0, Name("Scale"), IElemType(0));
return std::vector<NodeDef>{
scale_node,
NodeDef(OpDef{"Mul"},
{GO(0), scale_node.output_args[0]},
{GI(0)}),
NodeDef(OpDef{"Identity"},
{GI(0)},
{GI(1)})};
} else {
return std::vector<NodeDef>{
NodeDef(OpDef{"BatchScale", kMSDomain, 1},
{GO(0)},
{GI(0), GI(1)},
SrcNodeAttributes())};
}
} else if (input_count == 3) {
float scale_2 = attributes.at("scale_2").f();
if (scale_0 == scale_1 && scale_1 == scale_2) {
// Specialized branch to avoid duplicated data write.
NodeDef scale_node = ConstantScalarNode(scale_0, Name("Scale"), IElemType(0));
return std::vector<NodeDef>{
scale_node,
NodeDef(OpDef{"Mul"},
{GO(0), scale_node.output_args[0]},
{GI(0)}),
NodeDef(OpDef{"Identity"},
{GI(0)},
{GI(1)}),
NodeDef(OpDef{"Identity"},
{GI(0)},
{GI(2)})};
} else {
return std::vector<NodeDef>{
NodeDef(OpDef{"BatchScale", kMSDomain, 1},
{GO(0)},
{GI(0), GI(1), GI(2)},
SrcNodeAttributes())};
}
}
ORT_THROW("ScaledSum gradient builder does not support ", input_count, " inputs");
}
} // namespace training
} // namespace onnxruntime

1
orttraining/orttraining/core/graph/gradient_builder.h
View file

@ -85,6 +85,7 @@ DECLARE_GRADIENT_BUILDER(GetScatterElementsGradient)
DECLARE_GRADIENT_BUILDER(GetTriluGradient)
DECLARE_GRADIENT_BUILDER(GetFakeQuantGradient)
DECLARE_GRADIENT_BUILDER(GetLSTMGradient)
DECLARE_GRADIENT_BUILDER(GetScaledSumGradient)
DECLARE_GRADIENT_BUILDER(GetGRUGradient)
DECLARE_GRADIENT_BUILDER(GetReciprocalGradient)
DECLARE_GRADIENT_BUILDER(GetLeakyReluGradient)

1
orttraining/orttraining/core/graph/gradient_builder_registry.cc
View file

@ -117,6 +117,7 @@ void GradientBuilderRegistry::RegisterGradientBuilders() {
REGISTER_GRADIENT_BUILDER("Trilu", GetTriluGradient);
REGISTER_GRADIENT_BUILDER("FakeQuant", GetFakeQuantGradient);
REGISTER_GRADIENT_BUILDER("LSTMTraining", GetLSTMGradient);
REGISTER_GRADIENT_BUILDER("ScaledSum", GetScaledSumGradient);
REGISTER_GRADIENT_BUILDER("GRUTraining", GetGRUGradient);
REGISTER_GRADIENT_BUILDER("Reciprocal", GetReciprocalGradient);
REGISTER_GRADIENT_BUILDER("LeakyRelu", GetLeakyReluGradient);

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

@ -4104,7 +4104,8 @@ Return true if all elements are true and false otherwise.
ORT_ENFORCE(inferred_input_type->value_case() == TypeProto::kTensorType,
"PythonOpGrad's ", i, "-th input type must be a tensor.");
ORT_ENFORCE(inferred_input_type->tensor_type().elem_type() == input_tensor_types_proto->ints().at(i - 1),
"PythonOpGrad's ", i, "-th input type must be ", input_tensor_types_proto->ints().at(i - 1));
"PythonOpGrad's ", i, "-th input type must be ", input_tensor_types_proto->ints().at(i - 1),
", but inferred to be ", inferred_input_type->tensor_type().elem_type());
}
// Load expected output types.
@ -4515,6 +4516,61 @@ Return true if all elements are true and false otherwise.
updateOutputShape(ctx, 4, {sequence_length, num_directions, batch_size, hidden_size_x4});
});
ONNX_CONTRIB_OPERATOR_SCHEMA(ScaledSum)
.SetDomain(kMSDomain)
.SinceVersion(1)
.SetDoc(
"Compute scaled sum of multiple tensors in same shape (no broadcasting)."
"Formula: output = (input_0 * scale_0) + (input_1 * scale_1) + (input_2 * scale_2)")
.Attr("scale_0", "Scale for input_0.", AttributeProto::FLOAT)
.Attr("scale_1", "Scale for input_1.", AttributeProto::FLOAT)
.Attr("scale_2", "(Optional) Scale for input_2.", AttributeProto::FLOAT, OPTIONAL_VALUE)
.Input(0, "input_0", "input tensor", "T")
.Input(1, "input_1", "input tensor", "T")
.Input(2, "input_2", "input tensor", "T", OpSchema::Optional)
.Output(0, "output", "output tensor", "T")
.TypeConstraint(
"T",
{"tensor(float16)", "tensor(float)", "tensor(double)"},
"Constrain input types to float tensors.")
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
if (ctx.getNumInputs() == 3 && nullptr == ctx.getAttribute("scale_2"))
fail_shape_inference("Input count must be equal with scale count.");
propagateShapeAndTypeFromFirstInput(ctx);
});
ONNX_CONTRIB_OPERATOR_SCHEMA(BatchScale)
.SetDomain(kMSDomain)
.SinceVersion(1)
.SetDoc(
"Compute scaled input into outputs with different scaling factors (no broadcasting)."
"Formula:"
" output_0 = input * scale_0"
" output_1 = input * scale_1"
" output_2 = input * scale_2")
.Attr("scale_0", "Scale for input_0.", AttributeProto::FLOAT)
.Attr("scale_1", "Scale for input_1.", AttributeProto::FLOAT)
.Attr("scale_2", "(Optional) Scale for input_2.", AttributeProto::FLOAT, OPTIONAL_VALUE)
.Input(0, "input", "input tensor", "T")
.Output(0, "output_0", "output tensor", "T")
.Output(1, "output_1", "output tensor", "T")
.Output(2, "output_2", "output tensor", "T", OpSchema::Optional)
.TypeConstraint(
"T",
{"tensor(float16)", "tensor(float)", "tensor(double)"},
"Constrain input types to float tensors.")
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
if (ctx.getNumOutputs() == 3 && nullptr == ctx.getAttribute("scale_2"))
fail_shape_inference("Output count must be equal with scale count.");
for (size_t i = 0; i < ctx.getNumOutputs(); ++i) {
propagateElemTypeFromInputToOutput(ctx, 0, i);
if (hasInputShape(ctx, 0)) {
propagateShapeFromInputToOutput(ctx, 0, i);
}
}
});
ONNX_CONTRIB_OPERATOR_SCHEMA(LSTMGrad)
.SetDomain(kMSDomain)
.SinceVersion(1)

7
orttraining/orttraining/core/optimizer/graph_transformer_utils.cc
View file

@ -60,11 +60,12 @@
#include "orttraining/core/optimizer/lstm_replacement.h"
#include "orttraining/core/optimizer/transformer_layer_recompute.h"
#include "orttraining/core/optimizer/qdq_fusion.h"
#include "orttraining/core/optimizer/scaled_sum_fusion.h"
#include "orttraining/core/optimizer/shape_optimizer.h"
#include "orttraining/core/optimizer/transformer_layer_recompute.h"
#include "core/optimizer/pre_shape_node_elimination.h"
#include "core/optimizer/compute_optimizer/upstream_gather.h"
#include "core/optimizer/compute_optimizer/upstream_reshape.h"
#include "core/optimizer/pre_shape_node_elimination.h"
#include "orttraining/core/optimizer/compute_optimizer/padding_elimination.h"
#include "orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h"
@ -136,10 +137,12 @@ std::vector<std::unique_ptr<GraphTransformer>> GeneratePreTrainingTransformers(
transformers.emplace_back(std::make_unique<QDQFusion>(compatible_eps));
#if defined(USE_CUDA) || defined(USE_ROCM)
// We are supposed to use execution provider as indicator, but here we don't have access to the registered EP at this point
// We are supposed to use the execution provider as an indicator,
// but here we don't have access to the registered EP at this point
// as the session is not initialized yet. So using macro for now.
transformers.emplace_back(std::make_unique<BiasGeluFusion>(compatible_eps));
transformers.emplace_back(std::make_unique<IsInfReduceSumFusion>(compatible_eps));
transformers.emplace_back(std::make_unique<ScaledSumFusion>(compatible_eps));
#endif
if (config.enable_gelu_approximation) {

262
orttraining/orttraining/core/optimizer/scaled_sum_fusion.cc Normal file
View file

@ -0,0 +1,262 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "orttraining/core/optimizer/scaled_sum_fusion.h"
#include <onnx/defs/attr_proto_util.h>
#include "core/graph/graph_utils.h"
#include "core/optimizer/initializer.h"
#include "core/optimizer/utils.h"
namespace onnxruntime {
namespace {
// Supports limited data types.
static constexpr std::array supported_data_types{
ONNX_NAMESPACE::TensorProto_DataType_FLOAT16,
ONNX_NAMESPACE::TensorProto_DataType_FLOAT,
};
bool IsSupportedDataType(int32_t data_type) {
return std::find(supported_data_types.cbegin(), supported_data_types.cend(), data_type) !=
supported_data_types.cend();
}
bool IsShapeEqual(const ONNX_NAMESPACE::TensorShapeProto* lhs_shape,
const ONNX_NAMESPACE::TensorShapeProto* rhs_shape) {
ORT_ENFORCE(lhs_shape != nullptr && rhs_shape != nullptr);
if (lhs_shape->dim_size() != rhs_shape->dim_size()) {
return false;
}
for (int i = 0; i < lhs_shape->dim_size(); ++i) {
if (lhs_shape->dim(i).has_dim_value() && rhs_shape->dim(i).has_dim_value()) {
if (lhs_shape->dim(i).dim_value() != rhs_shape->dim(i).dim_value()) {
return false;
}
} else if (lhs_shape->dim(i).has_dim_param() && rhs_shape->dim(i).has_dim_param()) {
if (lhs_shape->dim(i).dim_param() != rhs_shape->dim(i).dim_param()) {
return false;
}
} else {
return false;
}
}
return true;
}
bool IsScaleOperator(Graph& graph, Node& node,
const ONNX_NAMESPACE::TensorShapeProto* output_shape,
float& scale_value) {
if (graph_utils::IsSupportedOptypeVersionAndDomain(node, "Div", {7, 13, 14})) {
// If node is Div, check:
// 1. The first input has the same shape as the given output_shape.
// 2. The second input is a constant initializer (containing a scalar or 1-D 1-element tensor).
bool first_input_check = (node.InputDefs()[0]->Shape() &&
IsShapeEqual(node.InputDefs()[0]->Shape(), output_shape));
if (first_input_check) {
const Node* div_input_2 = graph_utils::GetInputNode(node, 1);
auto div_input_2_shape = node.InputDefs()[1]->Shape();
bool second_input_check = div_input_2 == nullptr && div_input_2_shape &&
graph_utils::IsConstantInitializer(graph, node.InputDefs()[1]->Name(), false) &&
(div_input_2_shape->dim_size() == 0 // scalar
|| (div_input_2_shape->dim_size() == 1 &&
div_input_2_shape->dim(0).has_dim_value() &&
div_input_2_shape->dim(0).dim_value() == 1) /* 1d with 1 element */);
if (second_input_check) {
const ONNX_NAMESPACE::TensorProto* tensor_proto = nullptr;
if (!graph.GetInitializedTensor(node.InputDefs()[1]->Name(), tensor_proto)) {
return false;
}
Initializer init_const{*tensor_proto, graph.ModelPath()};
const auto data_type = tensor_proto->data_type();
if (data_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) {
const MLFloat16* val = init_const.data<MLFloat16>();
scale_value = 1.0f / math::halfToFloat(val[0].val);
} else if (data_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) {
scale_value = 1.0f / *init_const.data<float>();
} else {
return false;
}
return true;
}
}
}
return false;
}
} // namespace
Status ScaledSumFusion::ApplyImpl(Graph& graph, bool& modified, int /*graph_level*/,
const logging::Logger& logger) const {
GraphViewer graph_viewer(graph);
[[maybe_unused]] size_t handled_scaled_sum_count = 0; // For summary
const auto& order = graph_viewer.GetNodesInTopologicalOrder();
for (const auto index : order) {
auto* node_ptr = graph.GetNode(index);
if (!node_ptr)
// node was removed.
continue;
auto& node = *node_ptr;
// Find an Add that takes two inputs from other nodes' outputs (instead of any graph inputs or initializers).
// We also don't allow Add is generating graph outputs.
if (!graph_utils::IsSupportedOptypeVersionAndDomain(node, "Add", {6, 7, 13, 14}) ||
node.GetInputEdgesCount() != 2 /* two input MUST come from other nodes' outputs */ ||
graph.IsOutput(node.OutputDefs()[0]) ||
!graph_utils::IsSupportedProvider(node, GetCompatibleExecutionProviders())) {
continue;
}
const ONNX_NAMESPACE::TensorShapeProto* output_shape = node.OutputDefs()[0]->Shape();
const ONNX_NAMESPACE::TypeProto* output_type = node.OutputDefs()[0]->TypeAsProto();
if (!output_shape || !output_type) {
continue;
}
int elem_type = output_type->tensor_type().elem_type();
if (!IsSupportedDataType(elem_type)) {
continue;
}
InlinedVector<NodeArg*> data_input_args;
InlinedVector<float> scales;
data_input_args.reserve(3);
scales.reserve(3);
InlinedVector<std::reference_wrapper<Node>> nodes_to_remove;
// Be noted: it is possible the two input nodes are from the same node.
const Node* add_input_0 = graph_utils::GetInputNode(node, 0);
const Node* add_input_1 = graph_utils::GetInputNode(node, 1);
if (add_input_0 == nullptr || add_input_1 == nullptr) {
continue;
}
// Check the two inputs nodes of Add, if they are scaled operators, add them to the node list to remove.
auto check_add_input = [&graph, &output_shape, &nodes_to_remove,
&data_input_args, &scales](Node* add_input_node) -> bool {
float scale_value = 1.0f;
if (!IsScaleOperator(graph, *add_input_node, output_shape, scale_value)) {
return false;
}
// If node is not in nodes_to_remove, add it.
auto it = std::find_if(nodes_to_remove.begin(), nodes_to_remove.end(),
[&add_input_node](std::reference_wrapper<Node> n) {
return ((Node&)n).Index() == add_input_node->Index();
});
if (it == nodes_to_remove.end()) {
nodes_to_remove.push_back(*add_input_node);
}
data_input_args.push_back(add_input_node->MutableInputDefs()[0]);
scales.push_back(scale_value);
return true;
};
Node* add_input_node_0 = graph.GetNode(add_input_0->Index());
Node* add_input_node_1 = graph.GetNode(add_input_1->Index());
if (!check_add_input(add_input_node_0) || !check_add_input(add_input_node_1)) {
continue;
}
Node* last_node = &node;
// Handle three inputs only when Add node has one single consumer; and be noted we already check earlier
// the output is not in graph outputs.
if (node.GetOutputEdgesCount() == 1) {
Node& output_node = *graph.GetNode(node.OutputEdgesBegin()->GetNode().Index());
int output_node_port = node.OutputEdgesBegin()->GetDstArgIndex();
// Find the next Add node that use the output of current Add node as one of its inputs.
if (graph_utils::IsSupportedOptypeVersionAndDomain(output_node, "Add", {6, 7, 13, 14}) &&
!graph.IsOutput(output_node.OutputDefs()[0]) /* this Add cannot generate graph output */
) {
int the_other_input_port = 1 - output_node_port;
NodeArg* the_other_input_arg = output_node.MutableInputDefs()[the_other_input_port];
const Node* the_other_input_node = graph.GetProducerNode(the_other_input_arg->Name());
Node* mutable_the_other_input_node = the_other_input_node
? graph.GetNode(the_other_input_node->Index())
: nullptr;
bool the_other_node_output_edge_check = mutable_the_other_input_node == nullptr ||
mutable_the_other_input_node->GetOutputEdgesCount() == 1;
// Also make sure the other input arg has Shape equal to output_shape, we don't want to
// handle broadcast cases now.
if (the_other_node_output_edge_check &&
the_other_input_arg->Shape() && IsShapeEqual(the_other_input_arg->Shape(), output_shape)) {
last_node = &output_node;
nodes_to_remove.push_back(node);
float scale_value = 1.0f;
if (mutable_the_other_input_node && IsScaleOperator(graph, *mutable_the_other_input_node,
output_shape, scale_value)) {
data_input_args.push_back(mutable_the_other_input_node->MutableInputDefs()[0]);
nodes_to_remove.push_back(*mutable_the_other_input_node);
scales.push_back(scale_value);
} else {
// The other input is 1). a constant initializer or graph input, OR 2). it is not a scale operator:
// then we only add node arg into data input args, NOT need add any mode into nodes_to_remove.
data_input_args.push_back(mutable_the_other_input_node->MutableInputDefs()[0]);
scales.push_back(scale_value);
}
}
}
}
if (data_input_args.size() != scales.size() || data_input_args.size() < 2) {
continue;
}
auto type_info = *output_type;
InlinedVector<NodeArg*> output_args{&graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("ScaledSum"), &type_info)};
Node& scaled_sum_node = graph.AddNode(graph.GenerateNodeName("ScaledSum"),
"ScaledSum",
"FusedScaledSum",
data_input_args,
output_args,
nullptr,
kMSDomain);
ORT_ENFORCE(graph.SetOpSchemaFromRegistryForNode(scaled_sum_node),
"Failed to set op schema for " + scaled_sum_node.Name());
scaled_sum_node.SetExecutionProviderType(last_node->GetExecutionProviderType());
for (size_t scale_index = 0; scale_index < scales.size(); ++scale_index) {
scaled_sum_node.AddAttribute("scale_" + std::to_string(scale_index), scales[scale_index]);
}
graph_utils::ReplaceDownstreamNodeInput(graph, *last_node, 0, scaled_sum_node, 0);
// Firstly remove the node itself.
graph_utils::RemoveNodeOutputEdges(graph, *last_node);
graph.RemoveNode(last_node->Index());
// Then remove the parent nodes that may not be used by other nodes.
for (auto it = nodes_to_remove.rbegin(); it != nodes_to_remove.rend(); ++it) {
Node& n = *it;
if (n.GetOutputEdgesCount() != 0) {
continue;
}
graph_utils::RemoveNodeOutputEdges(graph, n);
graph.RemoveNode(n.Index());
}
modified = true;
handled_scaled_sum_count += 1;
}
LOGS(logger, INFO) << "Total fused ScaledSum node count: " << handled_scaled_sum_count;
return Status::OK();
}
} // namespace onnxruntime

47
orttraining/orttraining/core/optimizer/scaled_sum_fusion.h Normal file
View file

@ -0,0 +1,47 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/optimizer/graph_transformer.h"
namespace onnxruntime {
/*
Fuse continuous Add without broadcasting into ScaledSum.
Here is the pattern to find and fuse:
input_0 scale_0 input_1 scale_1
\ / \ /
Div Div
\ /
\ /
input_2 Add
\ /
Add
|
scale_0 and scale_1
> 1). MUST be scalar or single element 1D tensors,
> 2). and MUST be constant initializers.
==>
input_0 input_1 input_2
\ | /
ScaledSum
(attribute: scale_0=1/scale_0, scale_1=1/scale_1, scale_2=1)
|
**/
class ScaledSumFusion : public GraphTransformer {
public:
explicit ScaledSumFusion(const InlinedHashSet<std::string_view>& compatible_execution_providers = {}) noexcept
: GraphTransformer("ScaledSumFusion", compatible_execution_providers) {
}
Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override;
};
} // namespace onnxruntime

63
orttraining/orttraining/test/gradient/gradient_ops_test.cc
View file

@ -3025,6 +3025,69 @@ TEST(GradientCheckerTest, PadAndUnflattenGrad) {
x_datas, {}, true, false, &execution_providers));
EXPECT_IS_TINY(max_error);
}
TEST(GradientCheckerTest, ScaledSumGrad) {
// Two inputs.
{
float max_error;
GradientChecker<float, float, float> gradient_checker;
OpDef op_def{"ScaledSum", kMSDomain, 1};
TensorInfo x_info({4, 3});
TensorInfo y_info({4, 3});
std::vector<std::vector<float>> x_datas = {
{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11},
{0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f, 0.9f, 1.0f, 1.1f, 1.2f},
};
TensorInfo output0_info({4, 3}, true);
std::vector<ONNX_NAMESPACE::AttributeProto> attributes = {};
attributes.push_back(MakeAttribute("scale_0", static_cast<float>(0.5)));
attributes.push_back(MakeAttribute("scale_1", static_cast<float>(0.3)));
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
#ifdef USE_CUDA
execution_providers.emplace_back(DefaultCudaExecutionProvider());
#elif USE_ROCM
execution_providers.emplace_back(DefaultRocmExecutionProvider());
#endif
ASSERT_STATUS_OK(gradient_checker.ComputeGradientError(op_def, {x_info, y_info},
{output0_info}, &max_error,
x_datas, attributes, true, false, &execution_providers));
EXPECT_IS_TINY(max_error);
}
// Three inputs.
{
float max_error;
GradientChecker<float, float, float> gradient_checker;
OpDef op_def{"ScaledSum", kMSDomain, 1};
TensorInfo x_info({4, 3});
TensorInfo y_info({4, 3});
TensorInfo z_info({4, 3});
std::vector<std::vector<float>> x_datas = {
{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11},
{0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f, 0.9f, 1.0f, 1.1f, 1.2f},
{0.01f, 0.02f, 0.03f, 0.04f, 0.05f, 0.06f, -0.07f, -0.08f, -0.09f, -0.10f, -0.11f, -0.12f},
};
TensorInfo output0_info({4, 3}, true);
std::vector<ONNX_NAMESPACE::AttributeProto> attributes = {};
attributes.push_back(MakeAttribute("scale_0", static_cast<float>(0.2)));
attributes.push_back(MakeAttribute("scale_1", static_cast<float>(0.3)));
attributes.push_back(MakeAttribute("scale_2", static_cast<float>(0.5)));
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
#ifdef USE_CUDA
execution_providers.emplace_back(DefaultCudaExecutionProvider());
#elif USE_ROCM
execution_providers.emplace_back(DefaultRocmExecutionProvider());
#endif
ASSERT_STATUS_OK(gradient_checker.ComputeGradientError(op_def, {x_info, y_info, z_info},
{output0_info}, &max_error,
x_datas, attributes, true, false, &execution_providers));
EXPECT_IS_TINY(max_error);
}
}
#endif
TEST(GradientCheckerTest, ReciprocalGrad) {

281
orttraining/orttraining/test/optimizer/graph_transform_test.cc
View file

@ -26,8 +26,9 @@
#include "orttraining/core/session/training_session.h"
#include "orttraining/core/optimizer/loss_rewriter.h"
#include "orttraining/core/optimizer/bias_softmax_dropout_fusion.h"
#include "orttraining/core/optimizer/sce_loss_grad_bias_fusion.h"
#include "orttraining/core/optimizer/qdq_fusion.h"
#include "orttraining/core/optimizer/scaled_sum_fusion.h"
#include "orttraining/core/optimizer/sce_loss_grad_bias_fusion.h"
#include "orttraining/core/optimizer/lstm_replacement.h"
#include "orttraining/core/optimizer/gru_replacement.h"
#ifdef ENABLE_TRITON
@ -1302,6 +1303,284 @@ TEST_F(GraphTransformationTests, MegatronBARTSelfAttentionPartitionCorrectnessTe
// end of USE_CUDA
#endif
/*
Test graph as below.
graph input [1, 1, 256, 256] (float) scalar_0 graph input [1, 1, 256, 256] (float)
\ / /
Div Div -- scalar_1
[1, 1, 256, 256] (float) scalar_3 \ /
\ / Add
Div /
\ /
\ /
Add
|
Identity
|
graph out [1, 1, 256, 256] (float)
*/
TEST_F(GraphTransformationTests, ScaledSumFusionThreeInputs) {
auto pre_graph_checker = [](Graph& graph) -> Status {
auto op_count_pre = CountOpsInGraph(graph);
TEST_RETURN_IF_NOT(op_count_pre.size() == 3U);
TEST_RETURN_IF_NOT(op_count_pre["Div"] == 3);
TEST_RETURN_IF_NOT(op_count_pre["Add"] == 2);
TEST_RETURN_IF_NOT(op_count_pre["Identity"] == 1);
TEST_RETURN_IF_NOT(graph.GetAllInitializedTensors().size() == 3U);
return Status::OK();
};
auto post_graph_checker = [](Graph& graph) {
auto op_count = CountOpsInGraph(graph);
TEST_RETURN_IF_NOT(op_count.size() == 2U);
TEST_RETURN_IF_NOT(op_count["com.microsoft.ScaledSum"] == 1);
TEST_RETURN_IF_NOT(op_count["Identity"] == 1);
for (auto& node : graph.Nodes()) {
if (node.OpType() == "ScaledSum") {
TEST_RETURN_IF_NOT(node.InputDefs().size() == 3U);
auto& attrs = node.GetAttributes();
TEST_RETURN_IF_NOT(attrs.find("scale_0") != attrs.end());
TEST_RETURN_IF_NOT(attrs.find("scale_1") != attrs.end());
TEST_RETURN_IF_NOT(attrs.find("scale_2") != attrs.end());
TEST_RETURN_IF_NOT(1.0f / 0.5f == attrs.at("scale_0").f());
TEST_RETURN_IF_NOT(1.0f / 0.3f == attrs.at("scale_1").f());
TEST_RETURN_IF_NOT(1.0f / 0.2f == attrs.at("scale_2").f());
}
}
return Status::OK();
};
InlinedVector<bool> switch_orders{false, true};
for (bool switch_order : switch_orders) {
auto build_test_case = [switch_order](ModelTestBuilder& builder) {
auto* input_0_arg = builder.MakeInput<float>({{1, 1, 256, 256}});
auto* input_1_arg = builder.MakeInput<float>({{1, 1, 256, 256}});
auto* input_2_arg = builder.MakeInput<float>({{1, 1, 256, 256}});
auto* scalar_0_arg = builder.MakeScalarInitializer<float>(0.5f);
auto* scalar_1_arg = builder.MakeScalarInitializer<float>(0.3f);
auto* scalar_2_arg = builder.MakeScalarInitializer<float>(0.2f);
auto* div0_out = builder.MakeIntermediate();
auto* div1_out = builder.MakeIntermediate();
auto* div2_out = builder.MakeIntermediate();
builder.AddNode("Div", {input_0_arg, scalar_0_arg}, {div0_out});
builder.AddNode("Div", {input_1_arg, scalar_1_arg}, {div1_out});
auto* add1_out = builder.MakeIntermediate();
builder.AddNode("Add", {div0_out, div1_out}, {add1_out});
builder.AddNode("Div", {input_2_arg, scalar_2_arg}, {div2_out});
auto* add2_out = builder.MakeIntermediate();
if (switch_order) {
builder.AddNode("Add", {div2_out, add1_out}, {add2_out});
} else {
builder.AddNode("Add", {add1_out, div2_out}, {add2_out});
}
auto* graph_out = builder.MakeOutput();
builder.AddNode("Identity", {add2_out}, {graph_out});
};
const std::vector<int> opsets{12, 13, 14, 15};
for (auto& opset_version : opsets) {
std::unique_ptr<GraphTransformer> transformer = std::make_unique<ScaledSumFusion>();
ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, opset_version, *logger_, std::move(transformer),
TransformerLevel::Level1,
1, pre_graph_checker, post_graph_checker));
}
}
}
/*
Test graph as below.
graph input [1, 1, 256, 256] (float) scalar_0 graph input [1, 1, 256, 256] (float)
\ / |
Div Div -- scalar_1
[1, 1, 256, 256] (float) scalar_3 \ /
\ / Add
Sub /
\ /
\ /
Add
|
Identity
|
graph out [1, 1, 256, 256] (float)
*/
TEST_F(GraphTransformationTests, ScaledSumFusionThreeInputs_LastAddNotHaveScaleInput) {
auto pre_graph_checker = [](Graph& graph) -> Status {
auto op_count_pre = CountOpsInGraph(graph);
TEST_RETURN_IF_NOT(op_count_pre.size() == 4U);
TEST_RETURN_IF_NOT(op_count_pre["Div"] == 2);
TEST_RETURN_IF_NOT(op_count_pre["Add"] == 2);
TEST_RETURN_IF_NOT(op_count_pre["Identity"] == 1);
TEST_RETURN_IF_NOT(op_count_pre["Sub"] == 1);
TEST_RETURN_IF_NOT(graph.GetAllInitializedTensors().size() == 3U);
return Status::OK();
};
auto post_graph_checker = [](Graph& graph) {
auto op_count = CountOpsInGraph(graph);
TEST_RETURN_IF_NOT(op_count.size() == 3U);
TEST_RETURN_IF_NOT(op_count["com.microsoft.ScaledSum"] == 1);
TEST_RETURN_IF_NOT(op_count["Identity"] == 1);
TEST_RETURN_IF_NOT(op_count["Sub"] == 1);
for (auto& node : graph.Nodes()) {
if (node.OpType() == "ScaledSum") {
TEST_RETURN_IF_NOT(node.InputDefs().size() == 3U);
auto& attrs = node.GetAttributes();
TEST_RETURN_IF_NOT(attrs.find("scale_0") != attrs.end());
TEST_RETURN_IF_NOT(attrs.find("scale_1") != attrs.end());
TEST_RETURN_IF_NOT(attrs.find("scale_2") != attrs.end());
TEST_RETURN_IF_NOT(1.0f / 0.5f == attrs.at("scale_0").f());
TEST_RETURN_IF_NOT(1.0f / 0.3f == attrs.at("scale_1").f());
TEST_RETURN_IF_NOT(1.0f == attrs.at("scale_2").f());
}
}
return Status::OK();
};
InlinedVector<bool> switch_orders{false, true};
for (bool switch_order : switch_orders) {
auto build_test_case = [switch_order](ModelTestBuilder& builder) {
auto* input_0_arg = builder.MakeInput<float>({{1, 1, 256, 256}});
auto* input_1_arg = builder.MakeInput<float>({{1, 1, 256, 256}});
auto* input_2_arg = builder.MakeInput<float>({{1, 1, 256, 256}});
auto* scalar_0_arg = builder.MakeScalarInitializer<float>(0.5f);
auto* scalar_1_arg = builder.MakeScalarInitializer<float>(0.3f);
auto* scalar_2_arg = builder.MakeScalarInitializer<float>(0.2f);
auto* div0_out = builder.MakeIntermediate();
auto* div1_out = builder.MakeIntermediate();
auto* sub0_out = builder.MakeIntermediate();
builder.AddNode("Div", {input_0_arg, scalar_0_arg}, {div0_out});
builder.AddNode("Div", {input_1_arg, scalar_1_arg}, {div1_out});
auto* add1_out = builder.MakeIntermediate();
builder.AddNode("Add", {div0_out, div1_out}, {add1_out});
builder.AddNode("Sub", {input_2_arg, scalar_2_arg}, {sub0_out});
auto* add2_out = builder.MakeIntermediate();
if (switch_order) {
builder.AddNode("Add", {sub0_out, add1_out}, {add2_out});
} else {
builder.AddNode("Add", {add1_out, sub0_out}, {add2_out});
}
auto* graph_out = builder.MakeOutput();
builder.AddNode("Identity", {add2_out}, {graph_out});
};
const std::vector<int> opsets{12, 13, 14, 15};
for (auto& opset_version : opsets) {
std::unique_ptr<GraphTransformer> transformer = std::make_unique<ScaledSumFusion>();
ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, opset_version, *logger_, std::move(transformer),
TransformerLevel::Level1,
1, pre_graph_checker, post_graph_checker));
}
}
}
/*
Test graph as below.
graph input [1, 1, 256, 256] (float) scalar_0 graph input [1, 1, 256, 256] (float)
\ / /
Div Div -- scalar_1
[1, 1, 256, 256] (float) scalar_3 \ /
\ / Add
Div / \
\ / Identity
\ / |
Add graph out [1, 1, 256, 256] (float)
|
Identity
|
graph out [1, 1, 256, 256] (float)
*/
TEST_F(GraphTransformationTests, ScaledSumFusionTwoInputs) {
auto pre_graph_checker = [](Graph& graph) -> Status {
auto op_count_pre = CountOpsInGraph(graph);
TEST_RETURN_IF_NOT(op_count_pre.size() == 3U);
TEST_RETURN_IF_NOT(op_count_pre["Div"] == 3);
TEST_RETURN_IF_NOT(op_count_pre["Add"] == 2);
TEST_RETURN_IF_NOT(op_count_pre["Identity"] == 2);
TEST_RETURN_IF_NOT(graph.GetAllInitializedTensors().size() == 3U);
return Status::OK();
};
auto post_graph_checker = [](Graph& graph) {
auto op_count = CountOpsInGraph(graph);
TEST_RETURN_IF_NOT(op_count.size() == 4U);
TEST_RETURN_IF_NOT(op_count["Div"] == 1);
TEST_RETURN_IF_NOT(op_count["Add"] == 1);
TEST_RETURN_IF_NOT(op_count["com.microsoft.ScaledSum"] == 1);
TEST_RETURN_IF_NOT(op_count["Identity"] == 2);
for (auto& node : graph.Nodes()) {
if (node.OpType() == "ScaledSum") {
TEST_RETURN_IF_NOT(node.InputDefs().size() == 2U);
auto& attrs = node.GetAttributes();
TEST_RETURN_IF_NOT(attrs.find("scale_0") != attrs.end());
TEST_RETURN_IF_NOT(attrs.find("scale_1") != attrs.end());
TEST_RETURN_IF_NOT(attrs.find("scale_2") == attrs.end());
TEST_RETURN_IF_NOT(1.0f / 0.5f == attrs.at("scale_0").f());
TEST_RETURN_IF_NOT(1.0f / 0.3f == attrs.at("scale_1").f());
}
}
return Status::OK();
};
InlinedVector<bool> switch_orders{false, true};
for (bool switch_order : switch_orders) {
auto build_test_case = [switch_order](ModelTestBuilder& builder) {
auto* input_0_arg = builder.MakeInput<float>({{1, 1, 256, 256}});
auto* input_1_arg = builder.MakeInput<float>({{1, 1, 256, 256}});
auto* input_2_arg = builder.MakeInput<float>({{1, 1, 256, 256}});
auto* scalar_0_arg = builder.MakeScalarInitializer<float>(0.5f);
auto* scalar_1_arg = builder.MakeScalarInitializer<float>(0.3f);
auto* scalar_2_arg = builder.MakeScalarInitializer<float>(0.2f);
auto* div0_out = builder.MakeIntermediate();
auto* div1_out = builder.MakeIntermediate();
auto* div2_out = builder.MakeIntermediate();
builder.AddNode("Div", {input_0_arg, scalar_0_arg}, {div0_out});
builder.AddNode("Div", {input_1_arg, scalar_1_arg}, {div1_out});
auto* add1_out = builder.MakeIntermediate();
builder.AddNode("Add", {div0_out, div1_out}, {add1_out});
builder.AddNode("Div", {input_2_arg, scalar_2_arg}, {div2_out});
auto* add2_out = builder.MakeIntermediate();
if (switch_order) {
builder.AddNode("Add", {div2_out, add1_out}, {add2_out});
} else {
builder.AddNode("Add", {add1_out, div2_out}, {add2_out});
}
auto* graph_out = builder.MakeOutput();
builder.AddNode("Identity", {add2_out}, {graph_out});
auto* graph_output2 = builder.MakeOutput();
builder.AddNode("Identity", {add1_out}, {graph_output2});
};
const std::vector<int> opsets{12, 13, 14, 15};
for (auto& opset_version : opsets) {
std::unique_ptr<GraphTransformer> transformer = std::make_unique<ScaledSumFusion>();
ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, opset_version, *logger_, std::move(transformer),
TransformerLevel::Level1,
1, pre_graph_checker, post_graph_checker));
}
}
}
// end of DISABLE_CONTRIB_OPS
#endif

120
orttraining/orttraining/test/training_ops/cuda/batch_scale_test.cc Normal file
View file

@ -0,0 +1,120 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#if defined(USE_CUDA) || defined(USE_ROCM)
#include "test/common/tensor_op_test_utils.h"
#include "test/providers/provider_test_utils.h"
namespace onnxruntime {
namespace test {
static void PrepareInputAndOutputData(const std::vector<float>& input,
const std::vector<float>& scales,
std::vector<std::vector<float>>& outputs) {
for (size_t i = 0; i < outputs.size(); ++i) {
outputs.at(i).resize(input.size());
}
for (size_t i = 0; i < input.size(); ++i) {
outputs[0][i] = input[i] * scales[0];
outputs[1][i] = input[i] * scales[1];
if (outputs.size() == 3)
outputs[2][i] = input[i] * scales[2];
}
}
template <typename T>
static void RunBatchScaleOpTester(const std::vector<T>& input,
const std::vector<float>& scales,
const std::vector<std::vector<T>>& outputs,
const std::vector<int64_t>& shape) {
ORT_ENFORCE(scales.size() == outputs.size(), "scales and outputs should have the same size.");
OpTester test("BatchScale", 1, onnxruntime::kMSDomain);
test.AddInput<T>("input", shape, input);
test.AddOutput<T>("output_0", shape, outputs[0]);
test.AddOutput<T>("output_1", shape, outputs[1]);
if (outputs.size() == 3) {
test.AddOutput<T>("output_2", shape, outputs[2]);
}
test.AddAttribute<float>("scale_0", scales[0]);
test.AddAttribute<float>("scale_1", scales[1]);
if (scales.size() == 3) {
test.AddAttribute<float>("scale_2", scales[2]);
}
// Exclude CPU EP since it is not implemented yet.
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kCpuExecutionProvider});
}
static void RunBatchScaleTestWithFloatAndMLFloat16(const std::vector<float>& input,
const std::vector<float>& scales,
const std::vector<int64_t>& shape) {
std::vector<std::vector<float>> outputs;
outputs.resize(scales.size());
PrepareInputAndOutputData(input, scales, outputs);
RunBatchScaleOpTester(input, scales, outputs, shape);
std::vector<MLFloat16> input_half;
input_half.resize(input.size());
ConvertFloatToMLFloat16(input.data(), input_half.data(), static_cast<int>(input.size()));
std::vector<std::vector<MLFloat16>> outputs_half;
outputs_half.resize(scales.size());
for (size_t i = 0; i < outputs.size(); ++i) {
outputs_half[i].resize(outputs[i].size());
ConvertFloatToMLFloat16(outputs[i].data(), outputs_half[i].data(), static_cast<int>(outputs[i].size()));
}
RunBatchScaleOpTester(input_half, scales, outputs_half, shape);
}
TEST(BatchScaleTest, SmallTensor1D) {
std::vector<float> input = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.f};
float scale_0 = 0.25f;
float scale_1 = 0.25f;
float scale_2 = 0.5f;
std::vector<int64_t> shape{static_cast<int64_t>(input.size())};
RunBatchScaleTestWithFloatAndMLFloat16(input, {scale_0, scale_1, scale_2}, shape);
RunBatchScaleTestWithFloatAndMLFloat16(input, {scale_0, scale_1}, shape);
}
TEST(BatchScaleTest, SmallTensorVectorized1D) {
std::vector<float> input = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.f, 7.0f, 8.0f,
9.0f, 10.0f, 11.0f, 12.f, 13.0f, 14.0f, 15.0f, 16.0f,
17.0f, 18.0f, 19.0f, 20.0f, 21.0f, 22.f, 23.0f, 24.0f,
25.0f, 26.0f, 27.0f, 28.0f, 29.0f, 30.f, 31.0f, 32.0f};
float scale_0 = 0.25f;
float scale_1 = 0.25f;
float scale_2 = 0.5f;
std::vector<int64_t> shape{static_cast<int64_t>(input.size())};
RunBatchScaleTestWithFloatAndMLFloat16(input, {scale_0, scale_1, scale_2}, shape);
RunBatchScaleTestWithFloatAndMLFloat16(input, {scale_0, scale_1}, shape);
}
TEST(BatchScaleTest, SmallTensor2D) {
std::vector<float> input = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.f, 7.f, 8.f, 9.f};
float scale_0 = 0.25f;
float scale_1 = 0.25f;
float scale_2 = 0.5f;
std::vector<int64_t> shape{3, 3};
RunBatchScaleTestWithFloatAndMLFloat16(input, {scale_0, scale_1, scale_2}, shape);
RunBatchScaleTestWithFloatAndMLFloat16(input, {scale_0, scale_1}, shape);
}
TEST(BatchScaleTest, SmallTensorVectorized2D) {
std::vector<float> input = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.f, 7.0f, 8.0f,
9.0f, 10.0f, 11.0f, 12.f, 13.0f, 14.0f, 15.0f, 16.0f,
17.0f, 18.0f, 19.0f, 20.0f, 21.0f, 22.f, 23.0f, 24.0f,
25.0f, 26.0f, 27.0f, 28.0f, 29.0f, 30.f, 31.0f, 32.0f};
float scale_0 = 0.25f;
float scale_1 = 0.25f;
float scale_2 = 0.5f;
std::vector<int64_t> shape{4, 8};
RunBatchScaleTestWithFloatAndMLFloat16(input, {scale_0, scale_1, scale_2}, shape);
RunBatchScaleTestWithFloatAndMLFloat16(input, {scale_0, scale_1}, shape);
}
} // namespace test
} // namespace onnxruntime
#endif

129
orttraining/orttraining/test/training_ops/cuda/scaled_sum_test.cc Normal file
View file

@ -0,0 +1,129 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#if defined(USE_CUDA) || defined(USE_ROCM)
#include "test/common/tensor_op_test_utils.h"
#include "test/providers/provider_test_utils.h"
namespace onnxruntime {
namespace test {
static void PrepareInputAndOutputData(const std::vector<std::vector<float>>& input,
const std::vector<float>& scales,
std::vector<float>& output) {
output.resize(input[0].size());
size_t scale_size = scales.size();
for (size_t i = 0; i < input[0].size(); ++i) {
output[i] = input[0][i] * scales[0] + input[1][i] * scales[1] + (scale_size == 3 ? input[2][i] * scales[2] : 0.0f);
}
}
template <typename T>
static void RunScaledSumOpTester(const std::vector<std::vector<T>>& inputs,
const std::vector<float>& scales,
const std::vector<T>& output,
const std::vector<int64_t>& shape) {
OpTester test("ScaledSum", 1, onnxruntime::kMSDomain);
test.AddInput<T>("input0", shape, inputs[0]);
test.AddInput<T>("input1", shape, inputs[1]);
if (scales.size() == 3) {
test.AddInput<T>("input2", shape, inputs[2]);
}
test.AddOutput<T>("output", shape, output);
test.AddAttribute<float>("scale_0", scales[0]);
test.AddAttribute<float>("scale_1", scales[1]);
if (scales.size() == 3) {
test.AddAttribute<float>("scale_2", scales[2]);
}
// Exclude CPU EP since it is not implemented yet.
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kCpuExecutionProvider});
}
static void RunScaledSumWithFloatAndMLFloat16(const std::vector<std::vector<float>>& inputs,
const std::vector<float>& scales,
const std::vector<int64_t>& shape) {
std::vector<float> output;
PrepareInputAndOutputData(inputs, scales, output);
RunScaledSumOpTester(inputs, scales, output, shape);
std::vector<std::vector<MLFloat16>> inputs_half;
inputs_half.resize(inputs.size());
for (size_t i = 0; i < inputs.size(); ++i) {
inputs_half[i].resize(inputs[i].size());
ConvertFloatToMLFloat16(inputs[i].data(), inputs_half[i].data(), static_cast<int>(inputs[i].size()));
}
std::vector<MLFloat16> output_half;
output_half.resize(output.size());
ConvertFloatToMLFloat16(output.data(), output_half.data(), static_cast<int>(output.size()));
RunScaledSumOpTester(inputs_half, scales, output_half, shape);
}
TEST(ScaledSumTest, SmallTensor1D) {
std::vector<std::vector<float>> inputs = {{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.f},
{0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f},
{0.01f, 0.02f, 0.03f, 0.04f, 0.05f, 0.06f}};
float scale_0 = 0.25f;
float scale_1 = 0.25f;
float scale_2 = 0.5f;
std::vector<int64_t> shape{static_cast<int64_t>(inputs[0].size())};
RunScaledSumWithFloatAndMLFloat16(inputs, {scale_0, scale_1, scale_2}, shape);
RunScaledSumWithFloatAndMLFloat16(inputs, {scale_0, scale_1}, shape);
} // namespace test
TEST(ScaledSumTest, SmallTensorVectorized1D) {
std::vector<float> input = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.f, 7.0f, 8.0f,
9.0f, 10.0f, 11.0f, 12.f, 13.0f, 14.0f, 15.0f, 16.0f,
17.0f, 18.0f, 19.0f, 20.0f, 21.0f, 22.f, 23.0f, 24.0f,
25.0f, 26.0f, 27.0f, 28.0f, 29.0f, 30.f, 31.0f, 32.0f};
std::vector<std::vector<float>> inputs{input, input, input};
float scale_0 = 0.25f;
float scale_1 = 0.25f;
float scale_2 = 0.5f;
std::vector<int64_t> shape{static_cast<int64_t>(input.size())};
RunScaledSumWithFloatAndMLFloat16(inputs, {scale_0, scale_1, scale_2}, shape);
RunScaledSumWithFloatAndMLFloat16(inputs, {scale_0, scale_1}, shape);
}
TEST(ScaledSumTest, SmallTensor2D) {
std::vector<float> input = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.f, 7.f, 8.f, 9.f};
std::vector<std::vector<float>> inputs{input, input, input};
float scale_0 = 0.25f;
float scale_1 = 0.25f;
float scale_2 = 0.5f;
std::vector<int64_t> shape{3, 3};
RunScaledSumWithFloatAndMLFloat16(inputs, {scale_0, scale_1, scale_2}, shape);
RunScaledSumWithFloatAndMLFloat16(inputs, {scale_0, scale_1}, shape);
}
TEST(ScaledSumTest, SmallTensorVectorized2D) {
std::vector<float> input = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.f, 7.0f, 8.0f,
9.0f, 10.0f, 11.0f, 12.f, 13.0f, 14.0f, 15.0f, 16.0f,
17.0f, 18.0f, 19.0f, 20.0f, 21.0f, 22.f, 23.0f, 24.0f,
25.0f, 26.0f, 27.0f, 28.0f, 29.0f, 30.f, 31.0f, 32.0f};
std::vector<std::vector<float>> inputs{input, input, input};
float scale_0 = 0.25f;
float scale_1 = 0.25f;
float scale_2 = 0.5f;
std::vector<int64_t> shape{4, 8};
RunScaledSumWithFloatAndMLFloat16(inputs, {scale_0, scale_1, scale_2}, shape);
RunScaledSumWithFloatAndMLFloat16(inputs, {scale_0, scale_1}, shape);
}
} // namespace test
} // namespace onnxruntime
#endif

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

@ -204,7 +204,9 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, Inpl
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, FakeQuant);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, FakeQuantGrad);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BatchScale);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, PadAndUnflatten);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, ScaledSum);
// the kernels within the following ifdef are not included in a build with
// --enable_training_ops but without --enable_training
@ -455,7 +457,9 @@ Status RegisterCudaTrainingKernels(KernelRegistry& kernel_registry) {
kCudaExecutionProvider, kMSDomain, 1, float, FakeQuant)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(
kCudaExecutionProvider, kMSDomain, 1, float, FakeQuantGrad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BatchScale)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, PadAndUnflatten)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, ScaledSum)>,
// the kernels within the following ifdef are not included in a build with
// --enable_training_ops but without --enable_training
#ifdef ENABLE_TRAINING

68
orttraining/orttraining/training_ops/cuda/math/batch_scale.cc Normal file
View file

@ -0,0 +1,68 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <vector>
#include "orttraining/training_ops/cuda/math/batch_scale.h"
#include "orttraining/training_ops/cuda/math/batch_scale_impl.h"
namespace onnxruntime {
namespace cuda {
ONNX_OPERATOR_KERNEL_EX(
BatchScale,
kMSDomain,
1,
kCudaExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", BuildKernelDefConstraints<MLFloat16, float, double, BFloat16>()),
BatchScale);
// Put implementation in the anonymous namespace to avoid name collision in the global namespace.
namespace {
template <typename T>
struct BatchScaleFunctor {
void operator()(cudaStream_t stream,
int64_t input_element_count,
const Tensor* input_tensor,
const std::vector<float>& scales,
const std::vector<Tensor*>& output_tensors) const {
typedef typename ToCudaType<T>::MappedType CudaT;
std::vector<CudaT*> output_data_ptrs;
output_data_ptrs.reserve(output_tensors.size());
for (Tensor* output_tensor : output_tensors) {
output_data_ptrs.push_back(reinterpret_cast<CudaT*>(output_tensor->MutableData<T>()));
}
BatchScaleImpl<CudaT>(stream, input_element_count, reinterpret_cast<const CudaT*>(input_tensor->Data<T>()),
scales, output_data_ptrs);
}
};
} // namespace
Status BatchScale::ComputeInternal(OpKernelContext* context) const {
const Tensor* input_tensor = context->Input<Tensor>(0);
size_t output_count = scale2_.has_value() ? 3 : 2;
const auto& input_tensor_shape = input_tensor->Shape();
std::vector<Tensor*> output_tensors;
output_tensors.reserve(output_count);
for (size_t i = 0; i < output_count; ++i) {
output_tensors.push_back(context->Output(static_cast<int>(i), input_tensor_shape));
}
std::vector<float> scales{scale0_, scale1_};
if (output_count == 3) {
scales.push_back(scale2_.value());
}
utils::MLTypeCallDispatcher<float, MLFloat16, double, BFloat16> t_disp(input_tensor->GetElementType());
t_disp.Invoke<BatchScaleFunctor>(Stream(context), input_tensor_shape.Size(),
input_tensor, scales, output_tensors);
return Status::OK();
}
} // namespace cuda
} // namespace onnxruntime

34
orttraining/orttraining/training_ops/cuda/math/batch_scale.h Normal file
View file

@ -0,0 +1,34 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <optional>
#include "core/providers/cuda/cuda_kernel.h"
namespace onnxruntime {
namespace cuda {
class BatchScale final : public CudaKernel {
public:
BatchScale(const OpKernelInfo& info) : CudaKernel(info) {
ORT_ENFORCE(info.GetAttr<float>("scale_0", &scale0_).IsOK());
ORT_ENFORCE(info.GetAttr<float>("scale_1", &scale1_).IsOK());
float scale2_tmp;
if (info.GetAttr<float>("scale_2", &scale2_tmp).IsOK()) {
scale2_ = scale2_tmp;
}
}
Status ComputeInternal(OpKernelContext* context) const override;
private:
float scale0_;
float scale1_;
std::optional<float> scale2_;
};
} // namespace cuda
} // namespace onnxruntime

153
orttraining/orttraining/training_ops/cuda/math/batch_scale_impl.cu Normal file
View file

@ -0,0 +1,153 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <vector>
#include "core/providers/cuda/cu_inc/common.cuh"
#include "core/providers/cuda/cuda_common.h"
#include "orttraining/training_ops/cuda/math/batch_scale_impl.h"
namespace onnxruntime {
namespace cuda {
constexpr int kBlockSize = 256;
constexpr int kNumUnroll = 4;
template <typename T, int NumUnroll, int OutputCount, bool IsVectorized>
struct BatchScaleFunctor {
BatchScaleFunctor(const T* input,
const std::vector<float>& scales,
int64_t N,
const std::vector<T*>& outputs)
: N_(static_cast<CUDA_LONG>(N)),
input_data_(input) {
for (int i = 0; i < OutputCount; i++) {
outputs_[i] = outputs[i];
scales_[i] = scales[i];
}
}
__device__ __inline__ void operator()(CUDA_LONG idx) const {
CUDA_LONG id = idx * NumUnroll;
if (id >= N_) {
return;
}
using LoadT = aligned_vector<T, NumUnroll>;
T input0_value[NumUnroll];
if (IsVectorized) {
LoadT* input0_value_ptr = reinterpret_cast<LoadT*>(&input0_value[0]);
*input0_value_ptr = *reinterpret_cast<const LoadT*>(&input_data_[id]);
} else {
#pragma unroll
for (int i = 0; i < NumUnroll; i++) {
CUDA_LONG li = id + i;
if (li < N_) {
input0_value[i] = input_data_[li];
}
}
}
if (IsVectorized) {
T output_values[OutputCount][NumUnroll];
#pragma unroll
for (int i = 0; i < NumUnroll; i++) {
CUDA_LONG li = id + i;
if (li < N_) {
output_values[0][i] = static_cast<T>(static_cast<float>(input0_value[i]) * scales_[0]);
output_values[1][i] = static_cast<T>(static_cast<float>(input0_value[i]) * scales_[1]);
if (OutputCount == 3)
output_values[2][i] = static_cast<T>(static_cast<float>(input0_value[i]) * scales_[2]);
}
}
*reinterpret_cast<LoadT*>(&outputs_[0][id]) = *reinterpret_cast<LoadT*>(&output_values[0][0]);
*reinterpret_cast<LoadT*>(&outputs_[1][id]) = *reinterpret_cast<LoadT*>(&output_values[1][0]);
if (OutputCount == 3)
*reinterpret_cast<LoadT*>(&outputs_[2][id]) = *reinterpret_cast<LoadT*>(&output_values[2][0]);
} else {
#pragma unroll
for (int i = 0; i < NumUnroll; i++) {
CUDA_LONG li = id + i;
if (li < N_) {
outputs_[0][li] = static_cast<T>(static_cast<float>(input0_value[i]) * scales_[0]);
outputs_[1][li] = static_cast<T>(static_cast<float>(input0_value[i]) * scales_[1]);
if (OutputCount == 3)
outputs_[2][li] = static_cast<T>(static_cast<float>(input0_value[i]) * scales_[2]);
}
}
}
}
private:
T* outputs_[OutputCount];
float scales_[OutputCount];
const CUDA_LONG N_;
const T* input_data_;
};
template <typename FuncT>
__global__ void BatchScaleKernel(const FuncT functor) {
CUDA_LONG idx = blockDim.x * blockIdx.x + threadIdx.x;
functor(idx);
}
template <typename T>
void BatchScaleImpl(cudaStream_t stream,
int64_t input_element_count,
const T* input_data,
const std::vector<float>& scales,
const std::vector<T*>& outputs) {
const int blocksPerGrid = static_cast<int>(CeilDiv(input_element_count, kBlockSize * kNumUnroll));
constexpr int vec_alignment = std::alignment_of<aligned_vector<T, kNumUnroll>>::value;
const bool use_vectorized = (input_element_count % kNumUnroll == 0) &&
(reinterpret_cast<uintptr_t>(input_data) % vec_alignment == 0) &&
(reinterpret_cast<uintptr_t>(outputs[0]) % vec_alignment == 0) &&
(reinterpret_cast<uintptr_t>(outputs[1]) % vec_alignment == 0) &&
(outputs.size() < 3 || (reinterpret_cast<uintptr_t>(outputs[2]) % vec_alignment == 0));
const int output_count = static_cast<int>(outputs.size());
using TwoOutputVectorizedFunctorType = BatchScaleFunctor<T, kNumUnroll, 2, true>;
using TwoOutputNonVectorizedFunctorType = BatchScaleFunctor<T, kNumUnroll, 2, false>;
using ThreeOutputVectorizedFunctorType = BatchScaleFunctor<T, kNumUnroll, 3, true>;
using ThreeOutputNonVectorizedFunctorType = BatchScaleFunctor<T, kNumUnroll, 3, false>;
if (output_count == 2) {
if (use_vectorized)
BatchScaleKernel<TwoOutputVectorizedFunctorType><<<blocksPerGrid, kBlockSize, 0, stream>>>(
TwoOutputVectorizedFunctorType(input_data, scales, input_element_count, outputs));
else
BatchScaleKernel<TwoOutputNonVectorizedFunctorType><<<blocksPerGrid, kBlockSize, 0, stream>>>(
TwoOutputNonVectorizedFunctorType(input_data, scales, input_element_count, outputs));
} else if (output_count == 3) {
if (use_vectorized) {
BatchScaleKernel<ThreeOutputVectorizedFunctorType><<<blocksPerGrid, kBlockSize, 0, stream>>>(
ThreeOutputVectorizedFunctorType(input_data, scales, input_element_count, outputs));
} else {
BatchScaleKernel<ThreeOutputNonVectorizedFunctorType><<<blocksPerGrid, kBlockSize, 0, stream>>>(
ThreeOutputNonVectorizedFunctorType(input_data, scales, input_element_count, outputs));
}
} else {
ORT_THROW("Unsupported output count: ", output_count);
}
}
#define SPECIALIZE_BATCH_SCALE_IMPL(T) \
template void BatchScaleImpl<T>(cudaStream_t stream, \
int64_t input_element_count, \
const T* input_data, \
const std::vector<float>& scales, \
const std::vector<T*>& outputs);
SPECIALIZE_BATCH_SCALE_IMPL(half);
SPECIALIZE_BATCH_SCALE_IMPL(float);
SPECIALIZE_BATCH_SCALE_IMPL(double);
SPECIALIZE_BATCH_SCALE_IMPL(BFloat16);
#undef SPECIALIZE_BATCH_SCALE_IMPL
} // namespace cuda
} // namespace onnxruntime

19
orttraining/orttraining/training_ops/cuda/math/batch_scale_impl.h Normal file
View file

@ -0,0 +1,19 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <vector>
namespace onnxruntime {
namespace cuda {
template <typename T>
void BatchScaleImpl(cudaStream_t stream,
int64_t input_element_count,
const T* input_data,
const std::vector<float>& scales,
const std::vector<T*>& outputs);
} // namespace cuda
} // namespace onnxruntime

79
orttraining/orttraining/training_ops/cuda/math/scaled_sum.cc Normal file
View file

@ -0,0 +1,79 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <vector>
#include "orttraining/training_ops/cuda/math/scaled_sum.h"
#include "orttraining/training_ops/cuda/math/scaled_sum_impl.h"
namespace onnxruntime {
namespace cuda {
ONNX_OPERATOR_KERNEL_EX(
ScaledSum,
kMSDomain,
1,
kCudaExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", BuildKernelDefConstraints<MLFloat16, float, double, BFloat16>()),
ScaledSum);
// Put implementation in the anonymous namespace to avoid name collision in the global namespace.
namespace {
template <typename T>
struct ScaledSumFunctor {
void operator()(cudaStream_t stream,
int64_t input_element_count,
const std::vector<const Tensor*>& input_tensors,
const std::vector<float>& scales,
Tensor* output_tensor) const {
typedef typename ToCudaType<T>::MappedType CudaT;
std::vector<const CudaT*> input_data_ptrs;
input_data_ptrs.reserve(input_tensors.size());
for (const Tensor* input_tensor : input_tensors) {
input_data_ptrs.push_back(reinterpret_cast<const CudaT*>(input_tensor->Data<T>()));
}
ScaledSumImpl<CudaT>(stream, input_element_count, input_data_ptrs, scales,
reinterpret_cast<CudaT*>(output_tensor->MutableData<T>()));
}
};
} // namespace
Status ScaledSum::ComputeInternal(OpKernelContext* context) const {
std::vector<const Tensor*> input_tensors;
input_tensors.reserve(3);
for (size_t i = 0; i < 3; ++i) {
const Tensor* input_tensor = context->Input<Tensor>(static_cast<int>(i));
if (!input_tensor)
continue;
input_tensors.push_back(input_tensor);
}
ORT_ENFORCE(input_tensors.size() > 1, "Number of input tensors must be greater than 1.");
const auto& first_input_tensor_shape = input_tensors[0]->Shape();
for (size_t i = 1; i < input_tensors.size(); ++i) {
ORT_ENFORCE(input_tensors[i]->Shape() == first_input_tensor_shape,
"Shape of input tensors must be the same.");
}
std::vector<float> scales{scale0_, scale1_};
if (input_tensors.size() == 3) {
ORT_ENFORCE(scale2_.has_value(), "Scale 2 must be specified.");
scales.push_back(scale2_.value());
}
Tensor* output_tensor = context->Output(0, first_input_tensor_shape);
utils::MLTypeCallDispatcher<float, MLFloat16, double, BFloat16> t_disp(input_tensors[0]->GetElementType());
t_disp.Invoke<ScaledSumFunctor>(Stream(context), first_input_tensor_shape.Size(),
input_tensors, scales, output_tensor);
return Status::OK();
}
} // namespace cuda
} // namespace onnxruntime

33
orttraining/orttraining/training_ops/cuda/math/scaled_sum.h Normal file
View file

@ -0,0 +1,33 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <optional>
#include "core/providers/cuda/cuda_kernel.h"
namespace onnxruntime {
namespace cuda {
class ScaledSum final : public CudaKernel {
public:
ScaledSum(const OpKernelInfo& info) : CudaKernel(info) {
ORT_ENFORCE(info.GetAttr<float>("scale_0", &scale0_).IsOK());
ORT_ENFORCE(info.GetAttr<float>("scale_1", &scale1_).IsOK());
float scale2_tmp;
if (info.GetAttr<float>("scale_2", &scale2_tmp).IsOK()) {
scale2_ = scale2_tmp;
}
}
Status ComputeInternal(OpKernelContext* context) const override;
private:
float scale0_;
float scale1_;
std::optional<float> scale2_;
};
} // namespace cuda
} // namespace onnxruntime

168
orttraining/orttraining/training_ops/cuda/math/scaled_sum_impl.cu Normal file
View file

@ -0,0 +1,168 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/providers/cuda/cu_inc/common.cuh"
#include "core/providers/cuda/cuda_common.h"
#include "orttraining/training_ops/cuda/math/scaled_sum_impl.h"
namespace onnxruntime {
namespace cuda {
constexpr int kBlockSize = 256;
constexpr int kNumUnroll = 4;
template <typename T, int NumUnroll, int InputCount, bool IsVectorized>
struct ScaledSumFunctor {
ScaledSumFunctor(const std::vector<const T*>& inputs,
const std::vector<float>& scales,
int64_t N,
T* output) {
output_data_ = output;
N_ = static_cast<CUDA_LONG>(N);
for (int i = 0; i < InputCount; i++) {
inputs_[i] = inputs[i];
scales_[i] = scales[i];
}
}
__device__ __inline__ void operator()(CUDA_LONG idx) const {
CUDA_LONG id = idx * NumUnroll;
if (id >= N_) {
return;
}
using LoadT = aligned_vector<T, NumUnroll>;
T input_values[InputCount][NumUnroll];
if (IsVectorized) {
LoadT* input0_value_ptr = reinterpret_cast<LoadT*>(&input_values[0][0]);
*input0_value_ptr = *reinterpret_cast<const LoadT*>(&inputs_[0][id]);
LoadT* input1_value_ptr = reinterpret_cast<LoadT*>(&input_values[1][0]);
*input1_value_ptr = *reinterpret_cast<const LoadT*>(&inputs_[1][id]);
if (InputCount == 3) {
LoadT* input2_value_ptr = reinterpret_cast<LoadT*>(&input_values[2][0]);
*input2_value_ptr = *reinterpret_cast<const LoadT*>(&inputs_[2][id]);
}
} else {
#pragma unroll
for (int i = 0; i < NumUnroll; i++) {
CUDA_LONG li = id + i;
if (li < N_) {
input_values[0][i] = inputs_[0][li];
input_values[1][i] = inputs_[1][li];
if (InputCount == 3)
input_values[2][i] = inputs_[2][li];
}
}
}
if (IsVectorized) {
T output_value[NumUnroll];
#pragma unroll
for (int i = 0; i < NumUnroll; i++) {
CUDA_LONG li = id + i;
if (li < N_) {
if (InputCount == 3)
output_value[i] = input_values[0][i] * static_cast<T>(scales_[0]) +
input_values[1][i] * static_cast<T>(scales_[1]) +
input_values[2][i] * static_cast<T>(scales_[2]);
else
output_value[i] = input_values[0][i] * static_cast<T>(scales_[0]) +
input_values[1][i] * static_cast<T>(scales_[1]);
}
}
*reinterpret_cast<LoadT*>(&output_data_[id]) = *reinterpret_cast<LoadT*>(&output_value[0]);
} else {
T* output_value = output_data_ + id;
#pragma unroll
for (int i = 0; i < NumUnroll; i++) {
CUDA_LONG li = id + i;
if (li < N_) {
if (InputCount == 3)
output_value[i] = input_values[0][i] * static_cast<T>(scales_[0]) +
input_values[1][i] * static_cast<T>(scales_[1]) +
input_values[2][i] * static_cast<T>(scales_[2]);
else
output_value[i] = input_values[0][i] * static_cast<T>(scales_[0]) +
input_values[1][i] * static_cast<T>(scales_[1]);
}
}
}
}
private:
const T* inputs_[InputCount];
float scales_[InputCount];
CUDA_LONG N_;
T* output_data_;
};
template <typename FuncT>
__global__ void ScaledSumKernel(const FuncT functor) {
CUDA_LONG idx = blockDim.x * blockIdx.x + threadIdx.x;
functor(idx);
}
template <typename T>
void ScaledSumImpl(cudaStream_t stream,
int64_t input_element_count,
const std::vector<const T*>& inputs,
const std::vector<float>& scales,
T* output_data) {
const int blocksPerGrid = static_cast<int>(CeilDiv(input_element_count, kBlockSize * kNumUnroll));
constexpr int vec_alignment = std::alignment_of<aligned_vector<T, kNumUnroll>>::value;
const bool use_vectorized = (input_element_count % kNumUnroll == 0) &&
(reinterpret_cast<uintptr_t>(output_data) % vec_alignment == 0) &&
(reinterpret_cast<uintptr_t>(inputs[0]) % vec_alignment == 0) &&
(reinterpret_cast<uintptr_t>(inputs[1]) % vec_alignment == 0) &&
(inputs.size() < 3 || (reinterpret_cast<uintptr_t>(inputs[2]) % vec_alignment == 0));
const int input_count = static_cast<int>(inputs.size());
using TwoInputTVectorizedFunctorType = ScaledSumFunctor<T, kNumUnroll, 2, true>;
using TwoInputTNonVectorizedFunctorType = ScaledSumFunctor<T, kNumUnroll, 2, false>;
using ThreeInputTVectorizedFunctorType = ScaledSumFunctor<T, kNumUnroll, 3, true>;
using ThreeInputTNonVectorizedFunctorType = ScaledSumFunctor<T, kNumUnroll, 3, false>;
if (input_count == 2) {
if (use_vectorized) {
ScaledSumKernel<TwoInputTVectorizedFunctorType><<<blocksPerGrid, kBlockSize, 0, stream>>>(
TwoInputTVectorizedFunctorType(inputs, scales, input_element_count, output_data));
} else {
ScaledSumKernel<TwoInputTNonVectorizedFunctorType><<<blocksPerGrid, kBlockSize, 0, stream>>>(
TwoInputTNonVectorizedFunctorType(inputs, scales, input_element_count, output_data));
}
} else if (input_count == 3) {
if (use_vectorized) {
ScaledSumKernel<ThreeInputTVectorizedFunctorType><<<blocksPerGrid, kBlockSize, 0, stream>>>(
ThreeInputTVectorizedFunctorType(inputs, scales, input_element_count, output_data));
} else {
ScaledSumKernel<ThreeInputTNonVectorizedFunctorType><<<blocksPerGrid, kBlockSize, 0, stream>>>(
ThreeInputTNonVectorizedFunctorType(inputs, scales, input_element_count, output_data));
}
} else {
ORT_THROW("Unsupported input count: ", input_count);
}
}
#define SPECIALIZE_SCALED_SUM_IMPL(T) \
template void ScaledSumImpl<T>(cudaStream_t stream, \
int64_t input_element_count, \
const std::vector<const T*>& inputs, \
const std::vector<float>& scales, \
T* output_data);
SPECIALIZE_SCALED_SUM_IMPL(half);
SPECIALIZE_SCALED_SUM_IMPL(float);
SPECIALIZE_SCALED_SUM_IMPL(double);
SPECIALIZE_SCALED_SUM_IMPL(BFloat16);
#undef SPECIALIZE_SCALED_SUM_IMPL
} // namespace cuda
} // namespace onnxruntime

19
orttraining/orttraining/training_ops/cuda/math/scaled_sum_impl.h Normal file
View file

@ -0,0 +1,19 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <vector>
namespace onnxruntime {
namespace cuda {
template <typename T>
void ScaledSumImpl(cudaStream_t stream,
int64_t input_element_count,
const std::vector<const T*>& inputs,
const std::vector<float>& scales,
T* output_data);
} // namespace cuda
} // namespace onnxruntime