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

Optimize SCE loss compute (#15401)

Browse source 
### Optimize SCE loss compute

Compute optimization based on label data sparsity:
- Insert ShrunkenGather before SCELoss node, to filter out invalid
labels for compute.
- Support ShrunkenGather upstream.
- Added test for the above.
- Added flag to enable label sparsity optimization with env var, by
default disabled now. Will enable after comprehensive benchmarking
later.
- Extract common logic into test_optimizer_utils.h/cc from
core/optimizer/compute_optimzier_test.cc, then the common functions can
be shared by both core/optimizer/compute_optimzier_test.cc and
orttraining/core/optimizer/compute_optimzier_test.cc
- Extract common logic into shared_utils.h/cc: `GetONNXOpSetVersion` and
`Create1DInitializerFromVector`


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
This commit is contained in:
pengwa 2023-04-13 13:02:12 +08:00 committed by GitHub
parent 07b64d5275
commit 516c8e95fa
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
22 changed files with 1215 additions and 245 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
cmake/onnxruntime_optimizer.cmake
View file

@ -78,6 +78,8 @@ if (onnxruntime_ENABLE_TRAINING_APIS)
list(APPEND onnxruntime_optimizer_src_patterns
"${ORTTRAINING_SOURCE_DIR}/core/optimizer/*.h"
"${ORTTRAINING_SOURCE_DIR}/core/optimizer/*.cc"
"${ORTTRAINING_SOURCE_DIR}/core/optimizer/compute_optimizer/*.h"
"${ORTTRAINING_SOURCE_DIR}/core/optimizer/compute_optimizer/*.cc"
)
endif()

3
docs/ORTModule_Training_Guidelines.md
View file

@ -135,6 +135,9 @@ debugging). Disable it with following command:
export ORTMODULE_ENABLE_COMPUTE_OPTIMIZER=0
```
There are fine-grained flags to control different optimizations, but only applied when `ORTMODULE_ENABLE_COMPUTE_OPTIMIZER` is enabled.
- `ORTMODULE_ENABLE_LABEL_SPARSITY_OPT` is used to enable the optimization leveraging label sparsity.
#### ORTMODULE_ENABLE_INPUT_DENSITY_INSPECTOR
- **Feature Area**: *ORTMODULE/Runtime inspector*

41
onnxruntime/core/optimizer/compute_optimizer/shared_utils.cc
View file

@ -47,7 +47,11 @@ Node* InsertIntermediateNodeOnDestInput(Graph& graph,
for (size_t j = 0; j < new_node.MutableOutputDefs().size(); ++j) {
graph.UpdateProducerNode(new_node.MutableOutputDefs()[j]->Name(), new_node.Index());
}
graph.AddConsumerNode(src_node_arg->Name(), &new_node);
for (size_t j = 0; j < new_node.MutableInputDefs().size(); ++j) {
graph.AddConsumerNode(new_node.MutableInputDefs()[j]->Name(), &new_node);
}
const Node* src_node = graph.GetProducerNode(src_node_arg->Name());
if (src_node) {
int src_out_index = optimizer_utils::IndexOfNodeOutput(*src_node, *src_node_arg);
@ -150,6 +154,41 @@ std::pair<bool, std::vector<DimCompare>> CompareInputShapeWithOutputShape(
return std::make_pair<bool, std::vector<DimCompare>>(true, std::move(rets));
}
int GetONNXOpSetVersion(const Graph& graph) {
int onnx_opset = -1;
auto onnx_domain_it = graph.DomainToVersionMap().find(kOnnxDomain);
if (onnx_domain_it != graph.DomainToVersionMap().end()) {
onnx_opset = onnx_domain_it->second;
} else {
auto onnx_domain_alias_it = graph.DomainToVersionMap().find(kOnnxDomainAlias);
if (onnx_domain_alias_it != graph.DomainToVersionMap().end())
onnx_opset = onnx_domain_alias_it->second;
else
ORT_THROW("ONNX domain not found in this model");
}
return onnx_opset;
}
NodeArg* CreateInitializerFromVector(Graph& graph,
const InlinedVector<int64_t>& dims,
const InlinedVector<int64_t>& values,
const std::string& name) {
ONNX_NAMESPACE::TensorProto const_tensor;
const_tensor.set_name(name);
const_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64);
int64_t total_count = 1;
for (const int64_t dim : dims) {
const_tensor.add_dims(dim);
total_count *= dim;
}
ORT_ENFORCE(total_count == static_cast<int64_t>(values.size()));
const_tensor.set_raw_data(values.data(), values.size() * sizeof(int64_t));
return &graph_utils::AddInitializer(graph, const_tensor);
}
} // namespace onnxruntime::optimizer::compute_optimizer
#endif

18
onnxruntime/core/optimizer/compute_optimizer/shared_utils.h
View file

@ -156,5 +156,23 @@ std::pair<bool, std::vector<DimCompare>> CompareInputShapeWithOutputShape(
const ONNX_NAMESPACE::TensorShapeProto* full_broadcasted_shape,
const ONNX_NAMESPACE::TensorShapeProto* target_shape);
/**
* @brief Get opset version from the graph.
*/
int GetONNXOpSetVersion(const Graph& graph);
/**
* @brief Create an initializer from given dims/value vector and name.
*
* @param dims A int vector as the shape of the created initializer. If we want to create a scalar initializer,
* we should pass an empty vector.
* @param values A int vector containing the value buffer.
*/
NodeArg* CreateInitializerFromVector(Graph& graph,
const InlinedVector<int64_t>& dims,
const InlinedVector<int64_t>& values,
const std::string& name);
} // namespace onnxruntime::optimizer::compute_optimizer
#endif

42
onnxruntime/core/optimizer/compute_optimizer/upstream_gather.cc
View file

@ -307,6 +307,44 @@ std::optional<SliceInfo> IsSupportedGather(Graph& graph, Node& node,
return SliceInfo(graph, &node, dim_size == 0, "axis", axis, true);
}
std::optional<SliceInfo> IsSupportedShrunkenGather(Graph& graph, Node& node,
const InlinedHashSet<std::string_view>&
compatible_execution_providers,
const logging::Logger& logger) {
if (!graph_utils::IsSupportedOptypeVersionAndDomain(node, "ShrunkenGather", {1}, kMSDomain) ||
!graph_utils::IsSupportedProvider(node, compatible_execution_providers)) {
return std::nullopt;
}
auto data_shape = node.MutableInputDefs()[0]->Shape();
auto indices_shape = node.MutableInputDefs()[1]->Shape();
auto gather_out_shape = node.MutableOutputDefs()[0]->Shape();
if (data_shape == nullptr || indices_shape == nullptr || gather_out_shape == nullptr) {
LOG_DEBUG_INFO(logger, "Skip ShrunkenGather node " + node.Name() + " due to undefined shape." +
std::to_string(data_shape == nullptr) + std::to_string(indices_shape == nullptr) +
std::to_string(gather_out_shape == nullptr));
return std::nullopt;
}
const int data_rank = data_shape->dim_size();
if (data_rank <= 1) {
LOG_DEBUG_INFO(logger, "Skip ShrunkenGather node " + node.Name() + " due to data rank <= 1.");
return std::nullopt;
}
int axis = static_cast<int>(node.GetAttributes().at("axis").i());
axis = axis < 0 ? axis + data_rank : axis;
int dim_size = indices_shape->dim_size();
if (dim_size == 0) {
LOG_DEBUG_INFO(logger, "Skip ShrunkenGather node " + node.Name() + " due to unsupported dim size: " +
std::to_string(dim_size));
return std::nullopt;
}
return SliceInfo(graph, &node, false /*is_slice_scalar*/, "axis", axis, true);
}
} // namespace
std::optional<SliceInfo> UpStreamGatherGraphTransformer::IsSupportedForUpstream(
@ -317,7 +355,9 @@ std::optional<SliceInfo> UpStreamGatherGraphTransformer::IsSupportedForUpstream(
if (!gather_info.has_value()) {
gather_info = IsSupportedGather(graph, node, GetCompatibleExecutionProviders(), logger);
}
if (!gather_info.has_value()) {
gather_info = IsSupportedShrunkenGather(graph, node, GetCompatibleExecutionProviders(), logger);
}
return gather_info;
}

80
onnxruntime/core/optimizer/compute_optimizer/upstream_gather_actors.cc
View file

@ -83,29 +83,6 @@ TensorShapeProto CreateTensorShapeInsertDimAtAxis(const TensorShapeProto* src_sh
return updated_shape;
}
NodeArg* CreateUnsqueezeAxesInitializer(Graph& graph, const std::vector<int64_t>& values) {
ONNX_NAMESPACE::TensorProto axes_const_tensor;
axes_const_tensor.set_name(graph.GenerateNodeArgName("axes"));
axes_const_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64);
axes_const_tensor.add_dims(values.size());
axes_const_tensor.set_raw_data(values.data(), values.size() * sizeof(int64_t));
return &graph_utils::AddInitializer(graph, axes_const_tensor);
}
int GetONNXOpSetVersion(const Graph& graph) {
int onnx_opset = -1;
auto onnx_domain_it = graph.DomainToVersionMap().find(kOnnxDomain);
if (onnx_domain_it != graph.DomainToVersionMap().end()) {
onnx_opset = onnx_domain_it->second;
} else {
auto onnx_domain_alias_it = graph.DomainToVersionMap().find(kOnnxDomainAlias);
if (onnx_domain_alias_it != graph.DomainToVersionMap().end())
onnx_opset = onnx_domain_alias_it->second;
else
ORT_THROW("ONNX domain not found in this model");
}
return onnx_opset;
}
} // namespace
bool UpdateSliceOutputShape(NodeArg& arg_to_update, int axis_to_update,
@ -161,7 +138,8 @@ void AdaptInputAndOutputForScalarSlice(Graph& graph, Node& current_node, int cur
"Unsqueeze",
"Unsqueeze node",
{current_node.MutableInputDefs()[input_index],
CreateUnsqueezeAxesInitializer(graph, {pair.second.non_negative_axis})},
CreateInitializerFromVector(graph, {1}, {pair.second.non_negative_axis},
graph.GenerateNodeArgName("axes"))},
{&graph.GetOrCreateNodeArg(
graph.GenerateNodeArgName("unsqueeze_adaptor"),
current_node.MutableInputDefs()[input_index]->TypeAsProto())},
@ -218,7 +196,7 @@ void AdaptInputAndOutputForScalarSlice(Graph& graph, Node& current_node, int cur
"Squeeze",
"Squeeze node",
{consumer.MutableInputDefs()[index],
CreateUnsqueezeAxesInitializer(graph, {slice_axis})},
CreateInitializerFromVector(graph, {1}, {slice_axis}, graph.GenerateNodeArgName("axes"))},
{&graph.GetOrCreateNodeArg(
graph.GenerateNodeArgName("squeeze_adaptor"),
consumer.MutableInputDefs()[index]->TypeAsProto())},
@ -450,7 +428,7 @@ bool LayerNormalizationGatherActor::PreCheck(const Graph& /* graph */,
auto axis = static_cast<int64_t>(current_node.GetAttributes().at("axis").i());
axis = axis < 0 ? axis + current_node.InputDefs()[0]->Shape()->dim_size() : axis;
// Make sure layernorm/softmax's reduction happens after the axis we want to slice.
// Make sure LayerNormalization's reduction happens after the axis we want to slice.
if (axis <= info.non_negative_axis) {
return false;
}
@ -458,17 +436,8 @@ bool LayerNormalizationGatherActor::PreCheck(const Graph& /* graph */,
const NodeArg* gather_data_input_arg = current_node.OutputDefs()[info.GetDataProducerOutputIndex()];
const auto& gather_data_input_shape = gather_data_input_arg->Shape();
// Only handle the case where the first input is 3D.
auto ln_input_arg = current_node.InputDefs()[0];
auto ln_input_shape = ln_input_arg->Shape();
auto ln_input_rank = ln_input_shape->dim_size();
if (ln_input_rank != 3) {
LOG_DEBUG_INFO(logger, "Fail LayerNormalizationGatherActor::PreCheck for node " + current_node.Name() +
": data_input_rank is " + std::to_string(ln_input_rank));
return false;
}
auto [success, ret] = CompareInputShapeWithOutputShape(gather_data_input_shape, ln_input_shape);
auto [success, ret] = CompareInputShapeWithOutputShape(gather_data_input_shape,
current_node.InputDefs()[0]->Shape());
if (!success) {
// This should not happen!!!
LOG_DEBUG_INFO(logger, "Fail LayerNormalizationGatherActor::PreCheck for node " + current_node.Name() +
@ -481,9 +450,9 @@ bool LayerNormalizationGatherActor::PreCheck(const Graph& /* graph */,
all_input_cmp_rets[0] = std::move(ret);
shape_update_func = [&info](Node& node) -> void {
// Be noted: LayerNorm's 2nd and 3rd output have shape [batch, sequence, 1]. The dim is still kept even
// for reduced axes.
// So the slicing axis is same with the 1st output.
// Be noted: If LayerNorm's data input is [dim1, dim2, dim3], reduce axis is 1,
// then its 2nd and 3rd outputs have shape [dim1, dim2, 1]. The dim is still kept even
// for reduced axes, so the slicing axis is same with the 1st output.
for (size_t output_idx = 0; output_idx < node.MutableOutputDefs().size(); ++output_idx) {
UpdateSliceOutputShape(*node.MutableOutputDefs()[output_idx], info.non_negative_axis,
info.output_dim_on_axis);
@ -501,7 +470,7 @@ bool SoftmaxGatherActor::PreCheck(const Graph& graph, const Node& current_node,
auto axis = static_cast<int64_t>(current_node.GetAttributes().at("axis").i());
axis = axis < 0 ? axis + current_node.InputDefs()[0]->Shape()->dim_size() : axis;
// Make sure layernorm/softmax's reduction happens after the axis we want to slice.
// Make sure Softmax's reduction happens after the axis we want to slice.
if (axis <= info.non_negative_axis) {
return false;
}
@ -613,26 +582,6 @@ bool ReshapeGatherActor::PostProcess(
optimizer_utils::AppendTensorFromInitializer(graph, *current_node.InputDefs()[1], new_shape_const_values, true);
const int slice_axis = info_without_node.non_negative_axis;
auto create_new_initializer_from_vector = [&graph](NodeArg* arg_to_be_replaced,
const InlinedVector<int64_t>& new_values) -> NodeArg* {
// Create new TensorProto.
ONNX_NAMESPACE::TensorProto constant_tensor_proto;
constant_tensor_proto.set_name(graph.GenerateNodeArgName(arg_to_be_replaced->Name()));
constant_tensor_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64);
auto length = new_values.size();
constant_tensor_proto.add_dims(length);
constant_tensor_proto.set_raw_data(new_values.data(), length * sizeof(int64_t));
// Add initializer into Graph.
NodeArg* new_shape_arg = &graph_utils::AddInitializer(graph, constant_tensor_proto);
// Update the output arg shape.
ONNX_NAMESPACE::TensorShapeProto new_shape;
new_shape.add_dim()->set_dim_value(length);
new_shape_arg->SetShape(new_shape);
return new_shape_arg;
};
// If the shape constant on slice_axis is 0, then it keeps the original dim of input.
// If it is a scalar slice, then we just remove that dim. Otherwise, we don't need to update the dim value.
if (new_shape_const_values[slice_axis] == 0) {
@ -645,10 +594,11 @@ bool ReshapeGatherActor::PostProcess(
new_values.push_back(new_shape_const_values[i]);
}
}
auto new_shape_arg = create_new_initializer_from_vector(arg_to_be_replaced, new_values);
auto new_shape_arg = CreateInitializerFromVector(graph, {static_cast<int64_t>(new_values.size())}, new_values,
graph.GenerateNodeArgName(arg_to_be_replaced->Name()));
graph_utils::ReplaceNodeInput(current_node, 1, *new_shape_arg);
} else {
LOG_DEBUG_INFO(logger, "Reshape's shape has 0 specified for aixs: " + std::to_string(slice_axis) +
LOG_DEBUG_INFO(logger, "Reshape's shape has 0 specified for axis: " + std::to_string(slice_axis) +
", not need an update.");
}
return true;
@ -657,7 +607,9 @@ bool ReshapeGatherActor::PostProcess(
// If it selected shape is a dim value, we can update the shape tensor directory.
if (info_without_node.output_dim_on_axis.has_dim_value()) {
new_shape_const_values[slice_axis] = info_without_node.output_dim_on_axis.dim_value();
auto new_shape_arg = create_new_initializer_from_vector(current_node.MutableInputDefs()[1], new_shape_const_values);
auto new_shape_arg =
CreateInitializerFromVector(graph, {static_cast<int64_t>(new_shape_const_values.size())}, new_shape_const_values,
graph.GenerateNodeArgName(current_node.MutableInputDefs()[1]->Name()));
graph_utils::ReplaceNodeInput(current_node, 1, *new_shape_arg);
return true;
}

10
onnxruntime/core/optimizer/compute_optimizer/upstream_reshape.cc
View file

@ -140,19 +140,15 @@ ReshapeInfo UpStreamReshapeGraphTransformer::PropagateReshapeForInput(
input_args.push_back(current_node.MutableInputDefs()[current_node_input_index]);
// Prepare the target shape initializer. (Currently, only constant target shape is supported.)
std::vector<int64_t> new_shape;
InlinedVector<int64_t> new_shape;
new_shape.push_back(-1);
auto input_shape = current_node.MutableInputDefs()[current_node_input_index]->Shape();
for (int k = 2; k < input_shape->dim_size(); ++k) {
ORT_ENFORCE(input_shape->dim(k).has_dim_value());
new_shape.push_back(input_shape->dim(k).dim_value());
}
ONNX_NAMESPACE::TensorProto new_shape_const_tensor;
new_shape_const_tensor.set_name(graph.GenerateNodeArgName("new_shape"));
new_shape_const_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64);
new_shape_const_tensor.add_dims(new_shape.size());
new_shape_const_tensor.set_raw_data(new_shape.data(), new_shape.size() * sizeof(int64_t));
NodeArg* new_shape_arg = &graph_utils::AddInitializer(graph, new_shape_const_tensor);
NodeArg* new_shape_arg = CreateInitializerFromVector(graph, {static_cast<int64_t>(new_shape.size())}, new_shape,
graph.GenerateNodeArgName("new_shape"));
input_args.push_back(new_shape_arg);

8
onnxruntime/core/optimizer/compute_optimizer/upstream_transformer_base.cc
View file

@ -52,7 +52,7 @@ Status UpStreamGraphTransformerBase<T1, T2>::ApplyImpl(Graph& graph, bool& modif
size_t reordered_node_count = 0; // For summary
size_t passthrough_count = 0;
for (auto index : order) {
for (const auto index : order) {
auto* node_ptr = graph.GetNode(index);
if (!node_ptr)
// node was removed.
@ -91,7 +91,11 @@ Status UpStreamGraphTransformerBase<T1, T2>::ApplyImpl(Graph& graph, bool& modif
break;
}
if (graph.GetConsumerNodes(input_tensor_producer_node->MutableOutputDefs()[0]->Name()).size() > 1) {
// This condition implies few things:
// 1. The node's outputs are only used once (that's the node to upstream).
// 2. If the node has more than one outputs, only one of the outputs has output edge
// and that output is used by the node to upstream.
if (input_tensor_producer_node->GetOutputEdgesCount() > 1) {
LOG_DEBUG_INFO(logger, log_prefix + " stops at node " + input_tensor_producer_node->Name() +
" since multiple consumers found");
continue;

246
onnxruntime/test/optimizer/compute_optimizer_test.cc
View file

@ -21,26 +21,22 @@
#include "core/graph/graph_utils.h"
#include "core/graph/graph_viewer.h"
#include "core/graph/model.h"
#include "core/optimizer/common_subexpression_elimination.h"
#include "core/optimizer/compute_optimizer/upstream_gather.h"
#include "core/optimizer/compute_optimizer/upstream_reshape.h"
#include "core/optimizer/utils.h"
#include "core/platform/env.h"
#include "core/session/inference_session.h"
#include "core/util/math.h"
#include "test/common/tensor_op_test_utils.h"
#include "test/compare_ortvalue.h"
#include "test/framework/test_utils.h"
#include "test/optimizer/graph_transform_test_builder.h"
#include "test/optimizer/graph_transform_test_fixture.h"
#include "test/optimizer/test_optimizer_utils.h"
#include "test/providers/provider_test_utils.h"
#include "test/test_environment.h"
#include "test/util/include/temp_dir.h"
#include "test/util/include/asserts.h"
#include "test/util/include/default_providers.h"
#include "test/common/tensor_op_test_utils.h"
namespace onnxruntime {
namespace test {
@ -131,138 +127,6 @@ TEST(ComputeOptimizerTests, GatherND_MatMul) {
GatherNDComputationReductionTest("MatMul", *logger, SingleOpDefaultValidationFunc);
}
/**
* @brief Class represent a input data (dimensions, data type and value).
*/
struct TestInputData {
template <typename T>
TestInputData(const std::string& name, const TensorShapeVector& dims, const std::vector<T>& values)
: name_(name), dims_(dims), values_(values) {}
OrtValue ToOrtValue() {
OrtValue ortvalue;
std::vector<int64_t> dims;
dims.reserve(dims_.size());
dims.insert(dims.end(), dims_.begin(), dims_.end());
std::visit([&ortvalue, &dims](auto&& arg) {
using T = std::decay_t<decltype(arg)>;
if constexpr (std::is_same_v<T, std::vector<int64_t>> ||
std::is_same_v<T, std::vector<float>> ||
std::is_same_v<T, std::vector<MLFloat16>>)
CreateMLValue<typename T::value_type>(
TestCPUExecutionProvider()->GetAllocator(OrtMemTypeDefault), dims, arg, &ortvalue);
else
static_assert("Unspported types!");
},
values_);
return ortvalue;
}
std::string GetName() const {
return name_;
}
private:
std::string name_;
TensorShapeVector dims_;
std::variant<std::vector<float>, std::vector<MLFloat16>, std::vector<int64_t>> values_;
};
void RandomFillFloatVector(const TensorShapeVector& shape, std::vector<float>& data) {
static RandomValueGenerator random{1234};
data = random.Gaussian<float>(shape, 0.0f, 0.25f);
}
void RandomFillHalfVector(const TensorShapeVector& shape, std::vector<MLFloat16>& data) {
std::vector<float> data_float(TensorShape(shape).Size());
RandomFillFloatVector(shape, data_float);
std::transform(data_float.begin(), data_float.end(), data.begin(),
[](float value) { return MLFloat16(math::floatToHalf(value)); });
}
void RandomMasks(int64_t batch, int64_t sequence_length, std::vector<int64_t>& data) {
static RandomValueGenerator random{5678};
const std::vector<int64_t> num_count_to_random{batch};
std::vector<int64_t> random_seq_lens = random.Uniform<int64_t>(num_count_to_random, 0, sequence_length);
data.resize(batch * sequence_length); // fill with zeros first.
for (int64_t i = 0; i < batch; ++i) {
for (int64_t j = 0; j < sequence_length; ++j) {
if (j > random_seq_lens[i]) {
break;
}
data[i * sequence_length + j] = 1;
}
}
}
struct InputContainer {
InputContainer() = default;
template <typename T>
TestInputData& AddInput(const std::string& name, const TensorShapeVector dims, const std::vector<T>& values) {
inputs_.emplace_back(TestInputData(name, dims, values));
return inputs_.back();
}
template <typename T>
TestInputData& AddInput(const std::string& name, TensorShapeVector dims,
std::function<
void(const TensorShapeVector& shape, std::vector<T>& data)>
func = nullptr) {
std::vector<T> values(TensorShape(dims).Size());
if (func) {
func(dims, values);
}
inputs_.emplace_back(TestInputData(name, dims, values));
return inputs_.back();
}
void ToInputMap(NameMLValMap& feeds) const {
for (auto input : inputs_) {
feeds.insert({input.GetName(), input.ToOrtValue()});
}
}
private:
std::vector<TestInputData> inputs_;
};
static void RunModelWithData(const PathString& model_uri, const std::string session_log_id,
const std::string& provider_type, const InputContainer& input_container,
const std::vector<std::string>& output_names,
std::vector<OrtValue>& run_results) {
SessionOptions so;
// we don't want any transformation here.
so.graph_optimization_level = TransformerLevel::Default;
so.session_logid = session_log_id;
InferenceSession session_object{so, GetEnvironment()};
std::unique_ptr<IExecutionProvider> execution_provider;
if (provider_type == onnxruntime::kCpuExecutionProvider)
execution_provider = DefaultCpuExecutionProvider();
else if (provider_type == onnxruntime::kCudaExecutionProvider)
execution_provider = DefaultCudaExecutionProvider();
else if (provider_type == onnxruntime::kRocmExecutionProvider)
execution_provider = DefaultRocmExecutionProvider();
EXPECT_TRUE(session_object.RegisterExecutionProvider(std::move(execution_provider)).IsOK());
Status st;
ASSERT_TRUE((st = session_object.Load(model_uri)).IsOK()) << st.ErrorMessage();
ASSERT_TRUE((st = session_object.Initialize()).IsOK()) << st.ErrorMessage();
NameMLValMap feeds;
input_container.ToInputMap(feeds);
// Now run
RunOptions run_options;
st = session_object.Run(run_options, feeds, output_names, &run_results);
ASSERT_TRUE(st.IsOK()) << "RunModelWithData run graph failed with error: " << st.ErrorMessage();
}
TEST(ComputeOptimizerTests, GatherND_E2E) {
const logging::Logger* logger = &logging::LoggingManager::DefaultLogger();
auto model_uri = MODEL_FOLDER "computation_reduction/gathernd/e2e.onnx";
@ -1606,6 +1470,112 @@ TEST(ComputeOptimizerTests, GatherRobertaE2E) {
}
}
/*
Test graph include multiple equivalent subgraphs as below.
graph input [4, 32, 256] (float) graph input [4, 32, 256] (float)
| |
\_____________ ______________/
\ /
Add ______ [16]
| /
ShrunkenGather, axis = 1
|
Identity
|
graph output [4, 16, 256] (float)
Add an Identity node because currently we don't allow ShrunkenGather generates graph output.
*/
TEST(ComputeOptimizerTests, ShrunkenGatherElementwiseOps_PropagationOnTwoBranches) {
const logging::Logger* logger = &logging::LoggingManager::DefaultLogger();
InlinedVector<int64_t> gather_indices;
auto pre_graph_checker = [&gather_indices](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["Add"] == 1);
TEST_RETURN_IF_NOT(op_count_pre["com.microsoft.ShrunkenGather"] == 1);
TEST_RETURN_IF_NOT(op_count_pre["Identity"] == 1);
for (Node& node : graph.Nodes()) {
if (node.OpType() == "ShrunkenGather") {
TEST_RETURN_IF_NOT(gather_indices.empty());
constexpr bool require_constant = true;
NodeArg* initializer_node_arg = graph.GetNodeArg(node.InputDefs()[1]->Name());
TEST_RETURN_IF_NOT(optimizer_utils::AppendTensorFromInitializer(graph, *initializer_node_arg, gather_indices,
require_constant));
}
}
return Status::OK();
};
auto post_graph_checker = [&gather_indices](Graph& graph) {
auto op_count_post = CountOpsInGraph(graph);
TEST_RETURN_IF_NOT(op_count_post.size() == 3U);
TEST_RETURN_IF_NOT(op_count_post["Add"] == 1);
TEST_RETURN_IF_NOT(op_count_post["com.microsoft.ShrunkenGather"] == 2);
TEST_RETURN_IF_NOT(op_count_post["Identity"] == 1);
for (Node& node : graph.Nodes()) {
if (node.OpType() == "Add") {
const auto& input_defs = node.InputDefs();
{
auto producer_node = graph.GetProducerNode(input_defs[0]->Name());
TEST_RETURN_IF_NOT(producer_node != nullptr);
TEST_RETURN_IF_NOT(producer_node->OpType() == "ShrunkenGather");
InlinedVector<int64_t> values;
constexpr bool require_constant = true;
NodeArg* initializer_node_arg = graph.GetNodeArg(producer_node->InputDefs()[1]->Name());
TEST_RETURN_IF_NOT(optimizer_utils::AppendTensorFromInitializer(graph, *initializer_node_arg, values,
require_constant));
for (size_t i = 0; i < values.size(); i++) {
TEST_RETURN_IF_NOT(values[i] == gather_indices[i]);
}
}
{
auto producer_node = graph.GetProducerNode(input_defs[1]->Name());
TEST_RETURN_IF_NOT(producer_node != nullptr);
TEST_RETURN_IF_NOT(producer_node->OpType() == "ShrunkenGather");
InlinedVector<int64_t> values;
constexpr bool require_constant = true;
NodeArg* initializer_node_arg = graph.GetNodeArg(producer_node->InputDefs()[1]->Name());
TEST_RETURN_IF_NOT(optimizer_utils::AppendTensorFromInitializer(graph, *initializer_node_arg, values, require_constant));
for (size_t i = 0; i < values.size(); i++) {
TEST_RETURN_IF_NOT(values[i] == gather_indices[i]);
}
}
}
}
return Status::OK();
};
auto build_test_case = [](ModelTestBuilder& builder) {
auto* input1_arg = builder.MakeInput<int64_t>({{4, 32, 256}});
auto* input2_arg = builder.MakeInput<int64_t>({{4, 32, 256}});
auto* add_out = builder.MakeIntermediate();
builder.AddNode("Add", {input1_arg, input2_arg}, {add_out});
const std::vector<int64_t> slice_shape{16};
static RandomValueGenerator random{8888};
std::vector<int64_t> random_slices = random.Uniform<int64_t>(slice_shape, 0, 32);
auto* slice_initializer = builder.MakeInitializer<int64_t>(slice_shape, random_slices);
auto* gather_out = builder.MakeIntermediate();
builder.AddNode("ShrunkenGather", {add_out, slice_initializer}, {gather_out}, kMSDomain)
.AddAttribute("axis", static_cast<int64_t>(1));
auto* identity_out = builder.MakeOutput();
builder.AddNode("Identity", {gather_out}, {identity_out});
};
std::unique_ptr<GraphTransformer> transformer = std::make_unique<UpStreamGatherGraphTransformer>();
ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 14, *logger, std::move(transformer),
TransformerLevel::Level1,
1, pre_graph_checker, post_graph_checker));
}
/*
Test graph include multiple equivalent subgraphs as below.
graph input [4, 32, 256] (int64_t) graph input [4, 32, 256] (int64_t)

5
onnxruntime/test/optimizer/graph_transform_test_builder.h
View file

@ -186,6 +186,11 @@ class ModelTestBuilder {
return MakeInitializer({static_cast<int64_t>(data.size())}, data);
}
NodeArg* MakeEmptyInput() {
NodeArg* empty = &graph_.GetOrCreateNodeArg("", nullptr);
return empty;
}
Node& AddNode(const std::string& op_type,
const std::vector<NodeArg*>& input_args,
const std::vector<NodeArg*>& output_args,

87
onnxruntime/test/optimizer/test_optimizer_utils.cc Normal file
View file

@ -0,0 +1,87 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#ifdef _MSC_VER
#pragma warning(push)
#pragma warning(disable : 4244)
#endif
#include <random>
#include "core/framework/ort_value.h"
#include "core/graph/model.h"
#include "core/platform/env.h"
#include "core/session/inference_session.h"
#include "test/test_environment.h"
#include "test/util/include/asserts.h"
#include "test/util/include/default_providers.h"
#include "test/optimizer/test_optimizer_utils.h"
#include "test/common/tensor_op_test_utils.h"
namespace onnxruntime {
namespace test {
void RandomFillFloatVector(const TensorShapeVector& shape, std::vector<float>& data) {
static RandomValueGenerator random{1234};
data = random.Gaussian<float>(shape, 0.0f, 0.25f);
}
void RandomFillHalfVector(const TensorShapeVector& shape, std::vector<MLFloat16>& data) {
std::vector<float> data_float(TensorShape(shape).Size());
RandomFillFloatVector(shape, data_float);
std::transform(data_float.begin(), data_float.end(), data.begin(),
[](float value) { return MLFloat16(math::floatToHalf(value)); });
}
void RandomMasks(int64_t batch, int64_t sequence_length, std::vector<int64_t>& data) {
static RandomValueGenerator random{5678};
const std::vector<int64_t> num_count_to_random{batch};
std::vector<int64_t> random_seq_lens = random.Uniform<int64_t>(num_count_to_random, 0, sequence_length);
data.resize(batch * sequence_length); // fill with zeros first.
for (int64_t i = 0; i < batch; ++i) {
for (int64_t j = 0; j < sequence_length; ++j) {
if (j > random_seq_lens[i]) {
break;
}
data[i * sequence_length + j] = 1;
}
}
}
void RunModelWithData(const PathString& model_uri, const std::string session_log_id,
const std::string& provider_type, const InputContainer& input_container,
const std::vector<std::string>& output_names,
std::vector<OrtValue>& run_results) {
SessionOptions so;
// we don't want any transformation here.
so.graph_optimization_level = TransformerLevel::Default;
so.session_logid = session_log_id;
InferenceSession session_object{so, GetEnvironment()};
std::unique_ptr<IExecutionProvider> execution_provider;
if (provider_type == onnxruntime::kCpuExecutionProvider)
execution_provider = DefaultCpuExecutionProvider();
else if (provider_type == onnxruntime::kCudaExecutionProvider)
execution_provider = DefaultCudaExecutionProvider();
else if (provider_type == onnxruntime::kRocmExecutionProvider)
execution_provider = DefaultRocmExecutionProvider();
EXPECT_TRUE(session_object.RegisterExecutionProvider(std::move(execution_provider)).IsOK());
Status st;
ASSERT_TRUE((st = session_object.Load(model_uri)).IsOK()) << st.ErrorMessage();
ASSERT_TRUE((st = session_object.Initialize()).IsOK()) << st.ErrorMessage();
NameMLValMap feeds;
input_container.ToInputMap(feeds);
// Now run
RunOptions run_options;
st = session_object.Run(run_options, feeds, output_names, &run_results);
ASSERT_TRUE(st.IsOK()) << "RunModelWithData run graph failed with error: " << st.ErrorMessage();
}
} // namespace test
} // namespace onnxruntime

104
onnxruntime/test/optimizer/test_optimizer_utils.h Normal file
View file

@ -0,0 +1,104 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <variant>
#include "core/common/common.h"
#include "core/framework/tensor_shape.h"
#include "core/framework/float16.h"
#include "core/framework/framework_common.h"
#include "core/framework/ort_value.h"
#include "test/framework/test_utils.h"
namespace onnxruntime {
namespace test {
/**
* @brief Class represent a input data (dimensions, data type and value).
*/
struct TestInputData {
template <typename T>
TestInputData(const std::string& name, const TensorShapeVector& dims, const std::vector<T>& values)
: name_(name), dims_(dims), values_(values) {}
OrtValue ToOrtValue() {
OrtValue ortvalue;
std::vector<int64_t> dims;
dims.reserve(dims_.size());
dims.insert(dims.end(), dims_.begin(), dims_.end());
std::visit([&ortvalue, &dims](auto&& arg) {
using T = std::decay_t<decltype(arg)>;
if constexpr (std::is_same_v<T, std::vector<int64_t>> ||
std::is_same_v<T, std::vector<float>> ||
std::is_same_v<T, std::vector<MLFloat16>>)
CreateMLValue<typename T::value_type>(
TestCPUExecutionProvider()->GetAllocator(OrtMemTypeDefault), dims, arg, &ortvalue);
else
static_assert("Unspported types!");
},
values_);
return ortvalue;
}
std::string GetName() const {
return name_;
}
private:
std::string name_;
TensorShapeVector dims_;
std::variant<std::vector<float>, std::vector<MLFloat16>, std::vector<int64_t>> values_;
};
/**
* @brief A container for all input data.
*
*/
struct InputContainer {
InputContainer() = default;
template <typename T>
TestInputData& AddInput(const std::string& name, const TensorShapeVector dims, const std::vector<T>& values) {
inputs_.emplace_back(TestInputData(name, dims, values));
return inputs_.back();
}
template <typename T>
TestInputData& AddInput(const std::string& name, TensorShapeVector dims,
std::function<
void(const TensorShapeVector& shape, std::vector<T>& data)>
func = nullptr) {
std::vector<T> values(TensorShape(dims).Size());
if (func) {
func(dims, values);
}
inputs_.emplace_back(TestInputData(name, dims, values));
return inputs_.back();
}
void ToInputMap(NameMLValMap& feeds) const {
for (auto input : inputs_) {
feeds.insert({input.GetName(), input.ToOrtValue()});
}
}
private:
std::vector<TestInputData> inputs_;
};
void RandomFillFloatVector(const TensorShapeVector& shape, std::vector<float>& data);
void RandomFillHalfVector(const TensorShapeVector& shape, std::vector<MLFloat16>& data);
void RandomMasks(int64_t batch, int64_t sequence_length, std::vector<int64_t>& data);
void RunModelWithData(const PathString& model_uri, const std::string session_log_id,
const std::string& provider_type, const InputContainer& input_container,
const std::vector<std::string>& output_names,
std::vector<OrtValue>& run_results);
} // namespace test
} // namespace onnxruntime

2
orttraining/orttraining/core/framework/ortmodule_graph_builder.cc
View file

@ -152,7 +152,7 @@ Status OrtModuleGraphBuilder::OptimizeForwardGraph(std::unordered_set<std::strin
Graph& forward_graph = forward_model_->MainGraph();
ORT_RETURN_IF_ERROR(forward_graph.Resolve());
GraphTransformerManager graph_transformation_mgr{2};
GraphTransformerManager graph_transformation_mgr{3};
std::unique_ptr<CPUExecutionProvider> cpu_execution_provider =
std::make_unique<CPUExecutionProvider>(CPUExecutionProviderInfo());

225
orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.cc Normal file
View file

@ -0,0 +1,225 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <onnx/defs/attr_proto_util.h>
#include "orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h"
#include "core/graph/graph_utils.h"
#include "core/framework/random_seed.h"
#include "core/optimizer/initializer.h"
#include "core/optimizer/utils.h"
#include "core/optimizer/compute_optimizer/shared_utils.h"
#include "core/optimizer/compute_optimizer/upstream_gather_actors.h"
namespace onnxruntime {
// Put utilities in anonymous namespace.
namespace {
NodeArg* InsertNodesForValidLabelIndices(Graph& graph, Node& node, NodeArg* label_input, NodeArg* reduce_index_input) {
InlinedVector<NodeArg*> input_args{label_input, reduce_index_input};
InlinedVector<NodeArg*> output_args{&graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("label_sub_result"),
node.MutableInputDefs()[1]->TypeAsProto())};
Node& sub_node = graph.AddNode(graph.GenerateNodeName("labels_sub"), "Sub", "label sub padding idx", input_args,
output_args, nullptr, kOnnxDomain);
ORT_ENFORCE(graph.SetOpSchemaFromRegistryForNode(sub_node), "Failed to set op schema for " + sub_node.Name());
sub_node.SetExecutionProviderType(node.GetExecutionProviderType());
auto non_zero_out_arg = &graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("labels_filter_pad_result"),
node.MutableInputDefs()[1]->TypeAsProto());
Node& non_zero_node = graph.AddNode(graph.GenerateNodeName("labels_filter_pad"), "NonZero",
"labels filtering padding idx",
{sub_node.MutableOutputDefs()[0]},
{non_zero_out_arg}, nullptr, kOnnxDomain);
ORT_ENFORCE(graph.SetOpSchemaFromRegistryForNode(non_zero_node),
"Failed to set op schema for " + non_zero_node.Name());
const std::string dim_name = MakeString("valid_label_count_", utils::GetRandomSeed());
// 1D input NonZero generates output of shape (1,valid_token_count).
ONNX_NAMESPACE::TensorShapeProto non_zero_output_shape;
non_zero_output_shape.add_dim()->set_dim_value(1);
non_zero_output_shape.add_dim()->set_dim_param(dim_name);
non_zero_out_arg->SetShape(non_zero_output_shape);
non_zero_node.SetExecutionProviderType(node.GetExecutionProviderType());
InlinedVector<NodeArg*> squeeze_input_args;
squeeze_input_args.push_back(non_zero_out_arg);
bool opset_lower_than_13 = onnxruntime::optimizer::compute_optimizer::GetONNXOpSetVersion(graph) < 13;
onnxruntime::NodeAttributes attributes;
if (opset_lower_than_13) {
attributes["axes"] = ONNX_NAMESPACE::MakeAttribute("axes", std::vector<int64_t>{0});
} else {
squeeze_input_args.push_back(onnxruntime::optimizer::compute_optimizer::CreateInitializerFromVector(
graph, {1}, {0}, graph.GenerateNodeArgName("axes")));
}
auto squeeze_out_arg = &graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("squeeze_adaptor"),
non_zero_out_arg->TypeAsProto());
Node& squeeze_node = graph.AddNode(graph.GenerateNodeName("squeeze_adaptor"), "Squeeze", "nonzero_squeezer",
squeeze_input_args, {squeeze_out_arg}, &attributes, kOnnxDomain);
ORT_ENFORCE(graph.SetOpSchemaFromRegistryForNode(squeeze_node),
"Failed to set op schema for " + squeeze_node.Name());
// After Squeeze, the shape becomes (valid_token_count).
ONNX_NAMESPACE::TensorShapeProto squeeze_output_shape;
squeeze_output_shape.add_dim()->set_dim_param(dim_name);
squeeze_out_arg->SetShape(squeeze_output_shape);
squeeze_node.SetExecutionProviderType(node.GetExecutionProviderType());
return squeeze_out_arg;
}
} // namespace
Status InsertGatherBeforeSceLoss::ApplyImpl(Graph& graph, bool& modified, int /*graph_level*/,
const logging::Logger& logger) const {
LOG_DEBUG_INFO(logger, "Enter InsertGatherBeforeSceLoss");
GraphViewer graph_viewer(graph);
size_t handled_sce_node_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;
bool is_onnx_sce = graph_utils::IsSupportedOptypeVersionAndDomain(node, "SoftmaxCrossEntropyLoss", {12, 13});
bool is_internal_sce = graph_utils::IsSupportedOptypeVersionAndDomain(node, "SoftmaxCrossEntropyLossInternal", {1},
kMSDomain);
if ((!is_onnx_sce && !is_internal_sce) ||
!graph_utils::IsSupportedProvider(node, GetCompatibleExecutionProviders())) {
continue;
}
// Check whether this SCE node is handled or not.
const Node* labels_producer = graph.GetProducerNode(node.MutableInputDefs()[1]->Name());
// Skip if already inserted a ShrunkenGather node.
if (labels_producer && graph_utils::IsSupportedOptypeVersionAndDomain(
*labels_producer, "ShrunkenGather", {1}, kMSDomain)) {
LOG_DEBUG_INFO(logger, "Skip node " + node.Name() + "(" + node.OpType() +
") due to labels input is already consumed by a ShrunkenGather node.");
continue;
}
// Check shape requirements.
auto logits_shape = node.MutableInputDefs()[0]->Shape();
auto labels_shape = node.MutableInputDefs()[1]->Shape();
if (logits_shape == nullptr || labels_shape == nullptr) {
LOG_DEBUG_INFO(logger, "Skip node " + node.Name() + "(" + node.OpType() +
") due to undefined input shapes.");
continue;
}
if (logits_shape->dim_size() != 2 || labels_shape->dim_size() != 1) {
LOG_DEBUG_INFO(logger, "Skip node " + node.Name() + "(" + node.OpType() +
") due to unsupported input shape ranks.");
continue;
}
// Check attribute and input requirements.
std::string reduction = node.GetAttributes().at("reduction").s();
if (reduction.compare("mean") == 0 || reduction.compare("sum") == 0) {
// loss output is a scalar, don't need reset shape.
} else {
LOG_DEBUG_INFO(logger, "Skip node " + node.Name() + "(" + node.OpType() +
") due to loss [reduction=" + reduction + "].");
continue;
}
NodeArg* ignore_index_node_arg = nullptr;
if (is_internal_sce) {
if (node.InputDefs().size() < 4 || !graph_utils::IsConstantInitializer(
graph, node.InputDefs()[3]->Name(), /* check_outer_scope */ false)) {
LOG_DEBUG_INFO(logger, "Skip node " + node.Name() + "(" + node.OpType() +
") due to target padding idx is non-constant initializer. Input count: " + std::to_string(node.InputDefs().size()));
continue;
}
ignore_index_node_arg = node.MutableInputDefs()[3];
} else {
const auto ignore_index_attr = node.GetAttributes().find("ignore_index");
if (ignore_index_attr != node.GetAttributes().end()) {
int64_t ignore_index_value = (*ignore_index_attr).second.i();
ignore_index_node_arg = onnxruntime::optimizer::compute_optimizer::CreateInitializerFromVector(
graph, {}, {ignore_index_value}, graph.GenerateNodeArgName("ignore_index"));
} else {
LOG_DEBUG_INFO(logger, "Skip node " + node.Name() + "(" + node.OpType() +
") due to missing ignore_index attribute.");
continue;
}
}
std::vector<const Node*> sce_out1_consumers = graph.GetConsumerNodes(node.OutputDefs()[1]->Name());
if (sce_out1_consumers.size() != 0 || graph.IsOutput(node.OutputDefs()[1])) {
LOG_DEBUG_INFO(logger, "Skip node " + node.Name() + "(" + node.OpType() +
") due to log_prob output is graph output or consumed by other nodes.");
continue;
}
// SoftmaxCrossEntropyLossInternal op definition guarantees that the the first dimension of both inputs must match,
// we don't do the check explicitly here.
LOG_DEBUG_INFO(logger, "Inserting Sub+NonZero nodes for filtering valid tokens");
// It is possible a label input is used by multiple SoftmaxCrossEntropyLossInternal nodes, here we will create a
// subgraph retrieving valid tokens for each SoftmaxCrossEntropyLossInternal node.
// The duplication will be removed by CSE graph transformers.
NodeArg* valid_labels_input_arg =
InsertNodesForValidLabelIndices(graph, node, node.MutableInputDefs()[1], ignore_index_node_arg);
// Insert the ShrunkenGather node on the two inputs.
for (int i = 0; i < 2; ++i) {
InlinedVector<NodeArg*> input_args;
input_args.reserve(2);
input_args.push_back(node.MutableInputDefs()[i]);
input_args.push_back(valid_labels_input_arg);
InlinedVector<NodeArg*> output_args;
output_args.push_back(
&graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("label_filter_result"),
node.MutableInputDefs()[i]->TypeAsProto()));
/* new node input index to connect to node's input node*/
int new_gather_input_index_to_connect = 0;
/* new node output index to connect to node*/
int new_gather_output_index_to_connect = 0;
Node* new_gather_node = onnxruntime::optimizer::compute_optimizer::InsertIntermediateNodeOnDestInput(
graph, node,
i,
new_gather_input_index_to_connect,
new_gather_output_index_to_connect,
graph.GenerateNodeName("LabelsFilter"),
"ShrunkenGather",
"ShrunkenGather node to filter invalid tokens.",
input_args,
output_args,
{},
kMSDomain,
logger);
new_gather_node->SetExecutionProviderType(node.GetExecutionProviderType());
auto gather_out_arg = new_gather_node->MutableOutputDefs()[0];
onnxruntime::optimizer::compute_optimizer::UpdateSliceOutputShape(
*gather_out_arg, 0, valid_labels_input_arg->Shape()->dim(0));
}
modified = true;
handled_sce_node_count += 1;
}
LOG_DEBUG_INFO(logger, "Exit InsertGatherBeforeSceLoss, handled " + std::to_string(handled_sce_node_count) +
" SCE nodes");
return Status::OK();
}
} // namespace onnxruntime

77
orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h Normal file
View file

@ -0,0 +1,77 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/optimizer/graph_transformer.h"
namespace onnxruntime {
/**
* @brief Graph transformer that inserts ShrunkenGather before SCE nodes (e.g.
* SoftmaxCrossEntropyLossInternal/SoftmaxCrossEntropyLoss nodes).
*
* For label sparsity, we can remove them from the compute of loss function, by inserting ShrunkenGather
* operators for its two inputs.
*
* logits (float) [token_count, classes] labels (int64), [token_count]
* \ /
* SCE Node(ignore_index=-100)
* / \
* loss (shape is scalar or [token_count]) log_prob [token_count, classes]
*
* Be noted in Transformer-based models:
* > `token_count` usually equals with `batch size` x `sequence length`.
* > `classes` usually equals with `vocabulary`.
*
* Only insert ShrunkGather if all following conditions are met for SCE nodes`:
* 1. Its reduction attribute value is 'sum' or 'mean', to make sure loss is a scalar.
* Otherwise, the loss is in shape [token_count], changing on `token_count` will affect subsequent computations.
* 2. Its 2nd output (log_prob) MUST NOT be graph output and not consumed by other other nodes.
* 3. Its ignore_index exists and is a constant scalar value.
* 4. Its 2nd input label's input node is not `ShrunkGather` node (to avoid this transformer duplicated applied).
*
*
* After the transformation:
* labels [token_count]
* \_______
* \ \
* \ Sub(ignore_index)
* \ \
* | |
* | |
* | NonZero
* | |
* | Squeeze
* | |
* | indices of valid token [valid_token_count]
* \ |
* logits [token_count, classes] _________________\ _ _____/
* \ / \ /
* \ / \ /
* \ / \ /
* ShrunkenGather ShrunkenGather
* [valid_token_count, classes] [valid_token_count]
* \ /
* SCE Node (ignore_index=-100)
* / \
* / \
* loss (shape is scalar or [valid_token_count]) log_prob [valid_token_count, classes]
*
* In this specific scenario, it is easy to infer that valid_token_count <= token_count.
* After insertion, loss computation flop is reduced. Additionally, upstream Gather graph optimization
* will try to reduce flop for other ops further.
*/
class InsertGatherBeforeSceLoss : public GraphTransformer {
public:
InsertGatherBeforeSceLoss(const InlinedHashSet<std::string_view>& compatible_execution_providers = {}) noexcept
: GraphTransformer("InsertGatherBeforeSceLoss", compatible_execution_providers) {
}
/**
* @brief
*/
Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override;
};
} // namespace onnxruntime

3
orttraining/orttraining/core/optimizer/graph_transformer_config.h
View file

@ -24,6 +24,9 @@ struct TrainingGraphTransformerConfiguration : public GraphTransformerConfigurat
// Enable compute optimizer.
bool enable_compute_optimizer{false};
// Enable label sparsity compute optimization.
bool enable_label_sparsity_optimization{false};
};
} // namespace training

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

@ -11,8 +11,10 @@
#include "core/optimizer/cast_elimination.h"
#include "core/optimizer/common_subexpression_elimination.h"
#include "core/optimizer/compute_optimizer/upstream_gather.h"
#include "core/optimizer/compute_optimizer/upstream_reshape.h"
#include "core/optimizer/concat_slice_elimination.h"
#include "core/optimizer/constant_folding.h"
#include "core/optimizer/constant_sharing.h"
#include "core/optimizer/conv_activation_fusion.h"
#include "core/optimizer/conv_add_fusion.h"
#include "core/optimizer/conv_bn_fusion.h"
@ -51,6 +53,7 @@
#include "orttraining/core/framework/distributed_run_context.h"
#include "orttraining/core/optimizer/batchnorm_replacement.h"
#include "orttraining/core/optimizer/bitmask_dropout_replacement.h"
#include "orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h"
#include "orttraining/core/optimizer/concat_replacement.h"
#include "orttraining/core/optimizer/graph_transformer_registry.h"
#include "orttraining/core/optimizer/insert_output_rewriter.h"
@ -96,6 +99,10 @@ std::vector<std::unique_ptr<GraphTransformer>> GeneratePreTrainingTransformers(
ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique<InsertSoftmaxCrossEntropyLossOutput>()));
ORT_THROW_IF_ERROR(rule_transformer->Register(std::make_unique<LSTMReplacement>()));
// Put ConstantSharing before CommonSubexpressionElimination by intention as it can create more opportunities for
// CSE. For example, if A and B nodes both do Add operation with a same value but different initializers, by
// default, CSE will not merge them, because the different initializers are represented by different NodeArg.
transformers.emplace_back(std::make_unique<ConstantSharing>(compatible_eps));
// Remove duplicate nodes. Must be applied before any recompute transformations.
if (config.gelu_recompute || config.attn_dropout_recompute || config.transformer_layer_recompute) {
transformers.emplace_back(std::make_unique<CommonSubexpressionEliminationApplyOnce>(compatible_eps));
@ -136,9 +143,6 @@ std::vector<std::unique_ptr<GraphTransformer>> GeneratePreTrainingTransformers(
transformers.emplace_back(std::make_unique<ReshapeFusion>(compatible_eps));
transformers.emplace_back(std::make_unique<ConcatSliceElimination>(compatible_eps));
if (config.enable_compute_optimizer) {
transformers.emplace_back(std::make_unique<UpStreamGatherGraphTransformer>(compatible_eps));
}
if (config.gelu_recompute) {
transformers.emplace_back(std::make_unique<GeluRecompute>());
}
@ -150,12 +154,22 @@ std::vector<std::unique_ptr<GraphTransformer>> GeneratePreTrainingTransformers(
config.number_recompute_layers, compatible_eps));
}
if (config.propagate_cast_ops_config.level >= 0) {
const InlinedHashSet<std::string_view> cuda_execution_provider = {onnxruntime::kCudaExecutionProvider, onnxruntime::kRocmExecutionProvider};
const InlinedHashSet<std::string_view> cuda_execution_provider = {onnxruntime::kCudaExecutionProvider,
onnxruntime::kRocmExecutionProvider};
transformers.emplace_back(std::make_unique<PropagateCastOps>(config.propagate_cast_ops_config.strategy,
static_cast<size_t>(config.propagate_cast_ops_config.level),
config.propagate_cast_ops_config.allow,
cuda_execution_provider));
}
if (config.enable_compute_optimizer) {
transformers.emplace_back(std::make_unique<UpStreamGatherGraphTransformer>(compatible_eps));
if (config.enable_label_sparsity_optimization) {
transformers.emplace_back(std::make_unique<UpStreamReshapeGraphTransformer>(compatible_eps));
transformers.emplace_back(std::make_unique<InsertGatherBeforeSceLoss>(compatible_eps));
}
}
} break;
case TransformerLevel::Level2: {

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

@ -737,6 +737,7 @@ void addObjectMethodsForTraining(py::module& m, ExecutionProviderRegistrationFn
.def_readwrite("transformer_layer_recompute", &TrainingGraphTransformerConfiguration::transformer_layer_recompute)
.def_readwrite("number_recompute_layers", &TrainingGraphTransformerConfiguration::number_recompute_layers)
.def_readwrite("enable_compute_optimizer", &TrainingGraphTransformerConfiguration::enable_compute_optimizer)
.def_readwrite("enable_label_sparsity_optimization", &TrainingGraphTransformerConfiguration::enable_label_sparsity_optimization)
.def_readwrite("propagate_cast_ops_config", &TrainingGraphTransformerConfiguration::GraphTransformerConfiguration::propagate_cast_ops_config);
py::class_<OrtModuleGraphBuilderConfiguration> module_graph_builder_config(

5
orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py
View file

@ -185,6 +185,10 @@ class GraphExecutionManager(GraphExecutionInterface):
self._enable_compute_optimizer = (
ortmodule._defined_from_envvar("ORTMODULE_ENABLE_COMPUTE_OPTIMIZER", 1, warn=True) == 1
)
self._enable_label_sparsity_optimization = (
self._enable_compute_optimizer
and ortmodule._defined_from_envvar("ORTMODULE_ENABLE_LABEL_SPARSITY_OPT", 0, warn=True) == 1
)
# Flag to re-export the model due to attribute change on the original module.
# Re-export will be avoided if _skip_check is enabled.
@ -459,6 +463,7 @@ class GraphExecutionManager(GraphExecutionInterface):
graph_transformer_config.propagate_cast_ops_config.allow = self._propagate_cast_ops_allow
graph_transformer_config.propagate_cast_ops_config.strategy = self._propagate_cast_ops_strategy
graph_transformer_config.enable_compute_optimizer = self._enable_compute_optimizer
graph_transformer_config.enable_label_sparsity_optimization = self._enable_label_sparsity_optimization
return graph_transformer_config
def _initialize_graph_builder(self):

49
orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/torch_interop_utils/torch_interop_utils.cc
View file

@ -8,15 +8,15 @@
// In Torch forward run (e.g. THPFunction_apply), ctx of type THPFunction* (which is also a PyObject*)
// is created (https://github.com/pytorch/pytorch/blob/15532595209d2daf34d35e10f8d3d3b64966aea2/torch/csrc/autograd/python_function.cpp#L673).
// The ctx is used to run user-defined forward function and backward function as the first
// parameter. The same time, a cdata of type std::shared_ptr<PyNode> is created
// parameter. The same time, a cdata of type std::shared_ptr<PyNode> is created
// (https://github.com/pytorch/pytorch/blob/15532595209d2daf34d35e10f8d3d3b64966aea2/torch/csrc/autograd/python_function.cpp#L677),
// cdata is owned by:
// a). forward run output tensors as grad_fn_ property. (The full hierarchy is: Tensor owns
// a). forward run output tensors as grad_fn_ property. (The full hierarchy is: Tensor owns
// shared_pointer<TensorImpl>; TensorImpl owns std::unique_ptr<AutogradMeta>; AutogradMeta
// manages grad_/grad_fn_/grad_accumulator_. Among them, grad_fn_ is std::shared_ptr<PyNode>,
// e.g, the so called gradient function.)
// https://github.com/pytorch/pytorch/blob/15532595209d2daf34d35e10f8d3d3b64966aea2/torch/csrc/autograd/variable.h#L194
// b). the consumer operator of forward run outputs, will let its own PyNode/Node (gradident function)
// b). the consumer operator of forward run outputs, will let its own PyNode/Node (gradient function)
// owns the grad_fn_ (of type std::shared_ptr<PyNode>) of all inputs that require grad.
// https://github.com/pytorch/pytorch/blob/15532595209d2daf34d35e10f8d3d3b64966aea2/torch/csrc/autograd/function.h#L263
// BUT, if we run torch computation within PythonOp, b) is lost. SO, for some cases, where forward outputs
@ -25,7 +25,7 @@
// Then when PythonOpGrad runs, segment fault.
//
// So we add b)'s reference in this Pool when forward run returns; dereference from this Pool when backward
// completes, then ~PyNode() is called, which subsquently calls ~THPFunction() destorying ctx.
// completes, then ~PyNode() is called, which subsequently calls ~THPFunction() destroying ctx.
class PyNodeSharedPointerPool {
public:
static PyNodeSharedPointerPool& GetInstance() {
@ -33,7 +33,7 @@ class PyNodeSharedPointerPool {
return pool;
};
void RegisterGradFunc(const size_t& ctx_address, torch::autograd::AutogradMeta* autograd_meta){
void RegisterGradFunc(const size_t& ctx_address, torch::autograd::AutogradMeta* autograd_meta) {
auto it = grad_fns_.find(ctx_address);
TORCH_CHECK(it == grad_fns_.end(), "should not register grad_fn twice for ctx ", ctx_address);
@ -41,21 +41,20 @@ class PyNodeSharedPointerPool {
grad_fns_.emplace(ctx_address, std::move(autograd_meta->grad_fn_));
};
void UnRegisterGradFunc(const size_t& ctx_address){
void UnRegisterGradFunc(const size_t& ctx_address) {
auto it = grad_fns_.find(ctx_address);
TORCH_CHECK(it != grad_fns_.end(), "fail to find grad_fn for ctx ", ctx_address);
grad_fns_.erase(ctx_address);
};
void ClearAll(){
void ClearAll() {
grad_fns_.clear();
}
private:
PyNodeSharedPointerPool(){};
~PyNodeSharedPointerPool(){
};
~PyNodeSharedPointerPool(){};
PyNodeSharedPointerPool(const PyNodeSharedPointerPool&) = delete;
PyNodeSharedPointerPool& operator=(const PyNodeSharedPointerPool&) = delete;
@ -65,20 +64,19 @@ class PyNodeSharedPointerPool {
std::unordered_map<size_t, std::shared_ptr<torch::autograd::Node>> grad_fns_;
};
void clear_grad_fns_for_next_edges(at::Tensor target, std::vector<at::Tensor> saved_tensors) {
// For leaf tensor, there will be a AccumulateGrad (gradident function) created, which owns a
// reference to the tensor.
// For leaf tensor, there will be a AccumulateGrad (gradient function) created, which owns a
// reference to the tensor.
// For any user saved tensors (with save_for_backward), if the tensor is leaf, we put the map
// {AccumulateGrad*, Tensor*} into grad_fn_to_tensor_map.
std::unordered_map<torch::autograd::Node*, at::Tensor*> grad_fn_to_tensor_map;
for (auto& t: saved_tensors) {
std::unordered_map<torch::autograd::Node*, at::Tensor*> grad_fn_to_tensor_map;
for (auto& t : saved_tensors) {
auto grad_fn = t.grad_fn();
if (!grad_fn) {
grad_fn = torch::autograd::impl::try_get_grad_accumulator(t);
if (grad_fn) {
TORCH_CHECK(grad_fn_to_tensor_map.find(grad_fn.get()) == grad_fn_to_tensor_map.end(),
"found AccumulateGrad* is used by more than one tensors.");
TORCH_CHECK(grad_fn_to_tensor_map.find(grad_fn.get()) == grad_fn_to_tensor_map.end(),
"found AccumulateGrad* is used by more than one tensors.");
grad_fn_to_tensor_map.insert({grad_fn.get(), &t});
}
}
@ -103,14 +101,12 @@ void clear_grad_fns_for_next_edges(at::Tensor target, std::vector<at::Tensor> sa
}
}
void register_grad_fn(size_t ctx_address, at::Tensor target)
{
void register_grad_fn(size_t ctx_address, at::Tensor target) {
torch::autograd::AutogradMeta* autograd_meta = torch::autograd::impl::get_autograd_meta(target);
PyNodeSharedPointerPool::GetInstance().RegisterGradFunc(ctx_address, autograd_meta);
}
void unregister_grad_fn(size_t ctx_address)
{
void unregister_grad_fn(size_t ctx_address) {
PyNodeSharedPointerPool::GetInstance().UnRegisterGradFunc(ctx_address);
}
@ -118,7 +114,7 @@ void unregister_grad_fn(size_t ctx_address)
// When training program exits, PyNodeSharedPointerPool destructor is called, if grad_fns_ is not empty,
// PyNode::release_variables() will be called.
// (https://github.com/pytorch/pytorch/blob/15532595209d2daf34d35e10f8d3d3b64966aea2/torch/csrc/autograd/python_function.cpp#L168)
// The other hand, there is known issue when acquiring GIL in pybind11 destructors, there will be probabbly deadlock issue.
// The other hand, there is known issue when acquiring GIL in pybind11 destructors, there will be probably deadlock issue.
// (https://github.com/pybind/pybind11/issues/1446)
// The resolution here, we remove all maintained states before program exits.
@ -126,13 +122,13 @@ void unregister_grad_fn(size_t ctx_address)
// grad functions keeps accumulating without releasing, there might be memory (bound to those gradient function) leaks.
// Ideally this usually won't happen in real training case, so it should be fine.
// We CANNOT explictly clear grad functions before each forward pass to mitigate the known issue above.
// We CANNOT explicitly clear grad functions before each forward pass to mitigate the known issue above.
// For example:
// loss1 = forward_run(inputs1)
// loss2 = forward_run(inputs2)
// loss = loss1 + loss2
// loss.backward()
// If we clear grad functions in the beggining of the second `forward_run`, when `loss.backward()` runs,
// If we clear grad functions in the beginning of the second `forward_run`, when `loss.backward()` runs,
// the backward path of `loss1` will fail to run PythonOpGrad ops (if there is any).
void clear_all_grad_fns() {
PyNodeSharedPointerPool::GetInstance().ClearAll();
@ -140,7 +136,8 @@ void clear_all_grad_fns() {
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("register_grad_fn", &register_grad_fn, "increase grad_fn shared pointer reference.");
m.def("unregister_grad_fn", &unregister_grad_fn, "release grad_fn shared pointer referece.");
m.def("clear_all_grad_fns", &clear_all_grad_fns, "clear all grad_fn shared pointer refereces.");
m.def("clear_grad_fns_for_next_edges", &clear_grad_fns_for_next_edges, "remove reference on next edges' gradident funtions.");
m.def("unregister_grad_fn", &unregister_grad_fn, "release grad_fn shared pointer reference.");
m.def("clear_all_grad_fns", &clear_all_grad_fns, "clear all grad_fn shared pointer references.");
m.def("clear_grad_fns_for_next_edges", &clear_grad_fns_for_next_edges,
"remove reference on next edges' gradient functions.");
}

428
orttraining/orttraining/test/optimizer/compute_optimizer_test.cc Normal file
View file

@ -0,0 +1,428 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <random>
#include "core/graph/onnx_protobuf.h"
#include "gtest/gtest.h"
#include "gmock/gmock.h"
#include "asserts.h"
#include "core/common/span_utils.h"
#include "core/framework/data_types.h"
#include "core/framework/ort_value.h"
#include "core/graph/graph_utils.h"
#include "core/graph/graph_viewer.h"
#include "core/graph/model.h"
#include "core/optimizer/utils.h"
#include "core/util/math.h"
#include "orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h"
#include "test/common/tensor_op_test_utils.h"
#include "test/compare_ortvalue.h"
#include "test/framework/test_utils.h"
#include "test/optimizer/graph_transform_test_builder.h"
#include "test/optimizer/graph_transform_test_fixture.h"
#include "test/optimizer/test_optimizer_utils.h"
#include "test/providers/provider_test_utils.h"
#include "test/util/include/temp_dir.h"
#include "test/util/include/asserts.h"
#include "test/util/include/default_providers.h"
namespace onnxruntime {
namespace test {
#define MODEL_FOLDER ORT_TSTR("testdata/transform/")
/*
Test graph include multiple equivalent subgraphs as below.
graph input [32, 256] (float) graph input [32] (int64_t)
| |
\_____________ _______/ graph input -1, scalar (int64_t)
\ / _______/
\ / /
SCE Node, reduction = 'mean', output_type=1
|
|
graph output, loss, scalar (float)
*/
TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_Allowed) {
const logging::Logger* logger = &logging::LoggingManager::DefaultLogger();
for (const bool is_sce_internal : {true, false}) {
auto pre_graph_checker = [is_sce_internal](Graph& graph) -> Status {
auto op_count_pre = CountOpsInGraph(graph);
TEST_RETURN_IF_NOT(op_count_pre.size() == 1U);
if (is_sce_internal)
TEST_RETURN_IF_NOT(op_count_pre["com.microsoft.SoftmaxCrossEntropyLossInternal"] == 1);
else
TEST_RETURN_IF_NOT(op_count_pre["SoftmaxCrossEntropyLoss"] == 1);
return Status::OK();
};
auto post_graph_checker = [is_sce_internal](Graph& graph) {
auto op_count_post = CountOpsInGraph(graph);
TEST_RETURN_IF_NOT(op_count_post.size() == 5U);
TEST_RETURN_IF_NOT(op_count_post["Sub"] == 1);
TEST_RETURN_IF_NOT(op_count_post["NonZero"] == 1);
TEST_RETURN_IF_NOT(op_count_post["Squeeze"] == 1);
TEST_RETURN_IF_NOT(op_count_post["com.microsoft.ShrunkenGather"] == 2);
if (is_sce_internal)
TEST_RETURN_IF_NOT(op_count_post["com.microsoft.SoftmaxCrossEntropyLossInternal"] == 1);
else
TEST_RETURN_IF_NOT(op_count_post["SoftmaxCrossEntropyLoss"] == 1);
const NodeArg* squeeze_output_arg = nullptr;
for (Node& node : graph.Nodes()) {
if (node.OpType() == "Squeeze") {
squeeze_output_arg = node.OutputDefs()[0];
break;
}
}
TEST_RETURN_IF_NOT(squeeze_output_arg != nullptr);
for (Node& node : graph.Nodes()) {
if (node.OpType() == "SoftmaxCrossEntropyLossInternal" || node.OpType() == "SoftmaxCrossEntropyLoss") {
const auto& input_defs = node.InputDefs();
{
auto producer_node = graph.GetProducerNode(input_defs[0]->Name());
TEST_RETURN_IF_NOT(producer_node != nullptr);
TEST_RETURN_IF_NOT(producer_node->OpType() == "ShrunkenGather");
TEST_RETURN_IF_NOT(producer_node->InputDefs()[1] == squeeze_output_arg);
}
{
auto producer_node = graph.GetProducerNode(input_defs[1]->Name());
TEST_RETURN_IF_NOT(producer_node != nullptr);
TEST_RETURN_IF_NOT(producer_node->OpType() == "ShrunkenGather");
TEST_RETURN_IF_NOT(producer_node->InputDefs()[1] == squeeze_output_arg);
}
}
}
return Status::OK();
};
auto build_test_case = [is_sce_internal](ModelTestBuilder& builder) {
auto* input1_arg = builder.MakeInput<float>({{32, 256}});
auto* input2_arg = builder.MakeInput<int64_t>({{32}});
auto* sce_out1 = builder.MakeOutput();
NodeArg* empty = builder.MakeEmptyInput();
auto* sce_out2 = builder.MakeIntermediate();
if (is_sce_internal) {
auto* ignore_index_arg = builder.MakeScalarInitializer<int64_t>(-100);
Node& sce = builder.AddNode("SoftmaxCrossEntropyLossInternal",
{input1_arg, input2_arg, empty, ignore_index_arg},
{sce_out1, sce_out2}, kMSDomain);
sce.AddAttribute("reduction", "mean");
sce.AddAttribute("output_type", static_cast<int64_t>(1));
} else {
Node& sce = builder.AddNode("SoftmaxCrossEntropyLoss",
{input1_arg, input2_arg, empty},
{sce_out1, sce_out2});
sce.AddAttribute("reduction", "mean");
sce.AddAttribute("ignore_index", static_cast<int64_t>(-100));
}
};
std::vector<int> opsets{12, 13, 14, 15};
for (auto opset : opsets) {
std::unique_ptr<GraphTransformer> transformer = std::make_unique<InsertGatherBeforeSceLoss>();
ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, opset, *logger, std::move(transformer),
TransformerLevel::Level1,
1, pre_graph_checker, post_graph_checker));
}
}
}
/*
Test graph include multiple equivalent subgraphs as below.
graph input [32, 256] (float) graph input [32] (int64_t)
| |
\_____________ _______/ graph input -1, scalar (int64_t)
\ / _______/
\ / /
SCE Node, reduction = 'none', output_type=1
|
|
graph output, loss, [32] (float)
*/
TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_NotAllowed_ReduceNone) {
const logging::Logger* logger = &logging::LoggingManager::DefaultLogger();
for (const bool is_sce_internal : {true, false}) {
auto pre_graph_checker = [is_sce_internal](Graph& graph) -> Status {
auto op_count_pre = CountOpsInGraph(graph);
TEST_RETURN_IF_NOT(op_count_pre.size() == 1U);
if (is_sce_internal)
TEST_RETURN_IF_NOT(op_count_pre["com.microsoft.SoftmaxCrossEntropyLossInternal"] == 1);
else
TEST_RETURN_IF_NOT(op_count_pre["SoftmaxCrossEntropyLoss"] == 1);
return Status::OK();
};
auto post_graph_checker = [is_sce_internal](Graph& graph) {
auto op_count_post = CountOpsInGraph(graph);
TEST_RETURN_IF_NOT(op_count_post.size() == 1U);
if (is_sce_internal)
TEST_RETURN_IF_NOT(op_count_post["com.microsoft.SoftmaxCrossEntropyLossInternal"] == 1);
else
TEST_RETURN_IF_NOT(op_count_post["SoftmaxCrossEntropyLoss"] == 1);
return Status::OK();
};
auto build_test_case = [is_sce_internal](ModelTestBuilder& builder) {
auto* input1_arg = builder.MakeInput<float>({{32, 256}});
auto* input2_arg = builder.MakeInput<int64_t>({{32}});
auto* sce_out1 = builder.MakeOutput();
NodeArg* empty = builder.MakeEmptyInput();
auto* sce_out2 = builder.MakeIntermediate();
if (is_sce_internal) {
auto* ignore_index_arg = builder.MakeScalarInitializer<int64_t>(-100);
Node& sce = builder.AddNode("SoftmaxCrossEntropyLossInternal",
{input1_arg, input2_arg, empty, ignore_index_arg},
{sce_out1, sce_out2}, kMSDomain);
sce.AddAttribute("reduction", "none");
sce.AddAttribute("output_type", static_cast<int64_t>(1));
} else {
Node& sce = builder.AddNode("SoftmaxCrossEntropyLoss",
{input1_arg, input2_arg, empty},
{sce_out1, sce_out2});
sce.AddAttribute("reduction", "none");
sce.AddAttribute("ignore_index", static_cast<int64_t>(-100));
}
};
std::vector<int> opsets{12, 13, 14, 15};
for (auto opset : opsets) {
std::unique_ptr<GraphTransformer> transformer = std::make_unique<InsertGatherBeforeSceLoss>();
ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, opset, *logger, std::move(transformer),
TransformerLevel::Level1,
1, pre_graph_checker, post_graph_checker));
}
}
}
/*
Test graph include multiple equivalent subgraphs as below.
graph input [32, 256] (float) graph input [32] (int64_t)
| |
\_____________ _______/ graph input -1, scalar (int64_t)
\ / _______/
\ / /
SCE Node, reduction = 'none', output_type=1
|
|
graph output, loss, scalar (float)
*/
TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_NotAllowed_NoIgnoreIndex) {
const logging::Logger* logger = &logging::LoggingManager::DefaultLogger();
for (const bool is_sce_internal : {true, false}) {
auto pre_graph_checker = [is_sce_internal](Graph& graph) -> Status {
auto op_count_pre = CountOpsInGraph(graph);
TEST_RETURN_IF_NOT(op_count_pre.size() == 1U);
if (is_sce_internal)
TEST_RETURN_IF_NOT(op_count_pre["com.microsoft.SoftmaxCrossEntropyLossInternal"] == 1);
else
TEST_RETURN_IF_NOT(op_count_pre["SoftmaxCrossEntropyLoss"] == 1);
return Status::OK();
};
auto post_graph_checker = [is_sce_internal](Graph& graph) {
auto op_count_post = CountOpsInGraph(graph);
TEST_RETURN_IF_NOT(op_count_post.size() == 1U);
if (is_sce_internal)
TEST_RETURN_IF_NOT(op_count_post["com.microsoft.SoftmaxCrossEntropyLossInternal"] == 1);
else
TEST_RETURN_IF_NOT(op_count_post["SoftmaxCrossEntropyLoss"] == 1);
return Status::OK();
};
auto build_test_case = [is_sce_internal](ModelTestBuilder& builder) {
auto* input1_arg = builder.MakeInput<float>({{32, 256}});
auto* input2_arg = builder.MakeInput<int64_t>({{32}});
auto* sce_out1 = builder.MakeOutput();
auto* sce_out2 = builder.MakeIntermediate();
if (is_sce_internal) {
Node& sce = builder.AddNode("SoftmaxCrossEntropyLossInternal",
{input1_arg, input2_arg},
{sce_out1, sce_out2}, kMSDomain);
sce.AddAttribute("reduction", "sum");
sce.AddAttribute("output_type", static_cast<int64_t>(1));
} else {
Node& sce = builder.AddNode("SoftmaxCrossEntropyLoss",
{input1_arg, input2_arg},
{sce_out1, sce_out2});
sce.AddAttribute("reduction", "sum");
}
};
std::vector<int> opsets{12, 13, 14, 15};
for (auto opset : opsets) {
std::unique_ptr<GraphTransformer> transformer = std::make_unique<InsertGatherBeforeSceLoss>();
ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, opset, *logger, std::move(transformer),
TransformerLevel::Level1,
1, pre_graph_checker, post_graph_checker));
}
}
}
TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_MlmBertE2E) {
const logging::Logger* logger = &logging::LoggingManager::DefaultLogger();
// Be noted all dropout have a ratio be 0.0, to make it easier to compare when running with the session.
// This did not affect the transformer tests, because we did not remove the Dropout of ratio 0. in the middle.
auto model_uri = MODEL_FOLDER "computation_reduction/reshape/mlm_bert_e2e.onnx";
std::shared_ptr<Model> model;
ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger));
Graph& graph = model->MainGraph();
std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
onnxruntime::GraphTransformerManager graph_transformation_mgr{3};
ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::make_unique<InsertGatherBeforeSceLoss>(),
TransformerLevel::Level1));
ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger));
{
const NodeArg* squeeze_output_arg = nullptr;
for (Node& node : graph.Nodes()) {
if (node.OpType() == "NonZero") {
const std::vector<const Node*>& consumers = graph.GetConsumerNodes(node.OutputDefs()[0]->Name());
ASSERT_TRUE(consumers.size() == 1);
ASSERT_TRUE(consumers[0]->OpType() == "Squeeze");
squeeze_output_arg = consumers[0]->OutputDefs()[0];
break;
}
}
ASSERT_TRUE(squeeze_output_arg != nullptr);
for (Node& node : graph.Nodes()) {
if (node.OpType() == "SoftmaxCrossEntropyLossInternal") {
const auto& input_defs = node.InputDefs();
{
auto producer_node = graph.GetProducerNode(input_defs[0]->Name());
ASSERT_TRUE(producer_node != nullptr);
ASSERT_TRUE(producer_node->OpType() == "ShrunkenGather");
ASSERT_TRUE(producer_node->InputDefs()[1] == squeeze_output_arg);
}
{
auto producer_node = graph.GetProducerNode(input_defs[1]->Name());
ASSERT_TRUE(producer_node != nullptr);
ASSERT_TRUE(producer_node->OpType() == "ShrunkenGather");
ASSERT_TRUE(producer_node->InputDefs()[1] == squeeze_output_arg);
}
}
}
}
onnxruntime::test::TemporaryDirectory tmp_dir{ORT_TSTR("compute_optimizer_test_tmp_dir")};
PathString new_model_uri{ConcatPathComponent<PathChar>(
tmp_dir.Path(),
ORT_TSTR("insert_gather_before_sceloss_bert_e2e_optimized.onnx"))};
ASSERT_STATUS_OK(Model::Save(*model, new_model_uri));
int64_t batch_size = 8;
int64_t sequence_length = 16;
int64_t hidden_size = 1024;
InputContainer input_container;
input_container.AddInput<float>("input", {batch_size, sequence_length, hidden_size}, RandomFillFloatVector);
const TensorShapeVector dims_mask = {batch_size, sequence_length};
std::vector<int64_t> attention_mask(TensorShape(dims_mask).Size(), 0);
RandomMasks(batch_size, sequence_length, attention_mask);
input_container.AddInput<int64_t>("attention_mask", dims_mask, attention_mask);
input_container.AddInput<MLFloat16>("matmul1.weight", {hidden_size, 1024}, RandomFillHalfVector);
input_container.AddInput<MLFloat16>("add1.bias", {1024}, RandomFillHalfVector);
input_container.AddInput<MLFloat16>("matmul2.weight", {hidden_size, 1024}, RandomFillHalfVector);
input_container.AddInput<MLFloat16>("add2.bias", {1024}, RandomFillHalfVector);
input_container.AddInput<MLFloat16>("matmul3.weight", {hidden_size, 1024}, RandomFillHalfVector);
input_container.AddInput<MLFloat16>("add3.bias", {1024}, RandomFillHalfVector);
input_container.AddInput<MLFloat16>("matmul4.weight", {hidden_size, 1024}, RandomFillHalfVector);
input_container.AddInput<MLFloat16>("add4.bias", {1024}, RandomFillHalfVector);
input_container.AddInput<float>("layer_norm1.weight", {hidden_size}, RandomFillFloatVector);
input_container.AddInput<float>("layer_norm1.bias", {hidden_size}, RandomFillFloatVector);
input_container.AddInput<MLFloat16>("matmul7.weight", {hidden_size, hidden_size * 4}, RandomFillHalfVector);
input_container.AddInput<MLFloat16>("add7.bias", {hidden_size * 4}, RandomFillHalfVector);
input_container.AddInput<MLFloat16>("matmul8.weight", {hidden_size * 4, hidden_size}, RandomFillHalfVector);
input_container.AddInput<MLFloat16>("add8.bias", {hidden_size}, RandomFillHalfVector);
input_container.AddInput<float>("layer_norm2.weight", {hidden_size}, RandomFillFloatVector);
input_container.AddInput<float>("layer_norm2.bias", {hidden_size}, RandomFillFloatVector);
input_container.AddInput<MLFloat16>("matmul9.weight", {hidden_size, hidden_size}, RandomFillHalfVector);
input_container.AddInput<MLFloat16>("add9.bias", {hidden_size}, RandomFillHalfVector);
input_container.AddInput<float>("layer_norm3.weight", {hidden_size}, RandomFillFloatVector);
input_container.AddInput<float>("layer_norm3.bias", {hidden_size}, RandomFillFloatVector);
input_container.AddInput<MLFloat16>("matmul10.weight", {hidden_size, 30522}, RandomFillHalfVector);
input_container.AddInput<MLFloat16>("add10.bias", {30522}, RandomFillHalfVector);
const TensorShapeVector dims_labels = {batch_size * sequence_length};
static RandomValueGenerator random{8910};
std::vector<int64_t> labels = random.Uniform<int64_t>(dims_labels, 0, 30522);
const std::vector<int64_t> num_count_to_random{batch_size};
std::vector<int64_t> random_seq_lens = random.Uniform<int64_t>(num_count_to_random, 0, sequence_length);
for (int64_t i = 0; i < batch_size; ++i) {
for (int64_t j = 0; j < sequence_length; ++j) {
if (j > random_seq_lens[i]) {
labels[i * sequence_length + j] = -100;
}
}
}
input_container.AddInput<int64_t>("labels", dims_labels, labels);
const std::string all_provider_types[] = {
onnxruntime::kCpuExecutionProvider,
#ifdef USE_CUDA
onnxruntime::kCudaExecutionProvider,
#elif USE_ROCM
onnxruntime::kRocmExecutionProvider,
#endif
};
const std::vector<std::string> output_names = {"output-1"};
for (auto& provider_type : all_provider_types) {
std::vector<OrtValue> expected_ort_values;
RunModelWithData(model_uri, std::string("RawGraphRun"), provider_type,
input_container, output_names, expected_ort_values);
std::vector<OrtValue> actual_ort_values;
RunModelWithData(ToPathString(new_model_uri), std::string("OptimizedGraphRun"),
provider_type, input_container, output_names, actual_ort_values);
ASSERT_TRUE(expected_ort_values.size() == actual_ort_values.size());
constexpr double per_sample_tolerance = 1e-4;
constexpr double relative_per_sample_tolerance = 1e-4;
for (size_t i = 0; i < expected_ort_values.size(); i++) {
auto ret = CompareOrtValue(actual_ort_values[i], expected_ort_values[i],
per_sample_tolerance, relative_per_sample_tolerance, false);
EXPECT_EQ(ret.first, COMPARE_RESULT::SUCCESS) << ret.second;
}
}
}
} // namespace test
} // namespace onnxruntime

2
orttraining/orttraining/test/python/orttraining_test_ortmodule_autograd.py
View file

@ -126,7 +126,7 @@ def test_gelu_custom_func_rets_not_as_module_output():
def forward(self, model_input):
out = self.relu(model_input, self.bias)
# add * 9 by intention to make custom function's output
# NOT as module outputs (which are consumed by subsquent computations).
# NOT as module outputs (which are consumed by subsequent computations).
# This aims to trigger a GC for "out", saying, out is released,
# the underlying std::shared<PyNode> still have other references.
# Otherwise, a segment fault will be triggered.