mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-22 19:23:30 +00:00
Disable Materializing Grads (#6822)
* disable materialize grads * gradient builder bugfix * fix ut * fix ut * resolve comments and bugfix * add more assert * disable forward compare for now
This commit is contained in:
Vincent Wang
GitHub
parent
dfc7c18e31
commit
56c5620fd2
13 changed files with 203 additions and 138 deletions
|
|
@ -262,15 +262,24 @@ static void MoveAllNodeOutputs(Graph& graph, Node& src_node, Node& target_node)
|
|||
//--- end of local helpers ---
|
||||
//----------------------------
|
||||
|
||||
int GetNodeInputIndexFromInputName(const Node& node, const std::string& input_name) {
|
||||
auto itr = std::find_if(node.InputDefs().begin(), node.InputDefs().end(),
|
||||
[&input_name](const NodeArg* input) { return input->Name() == input_name; });
|
||||
ORT_ENFORCE(itr != node.InputDefs().end(),
|
||||
"Attempting to get index for an input which does not exist.");
|
||||
auto index = std::distance(node.InputDefs().begin(), itr);
|
||||
int GetIndexFromName(const Node& node, const std::string& name, bool is_input) {
|
||||
const auto& node_args = is_input ? node.InputDefs() : node.OutputDefs();
|
||||
auto itr = std::find_if(node_args.begin(), node_args.end(),
|
||||
[&name](const NodeArg* node_arg) { return node_arg->Name() == name; });
|
||||
ORT_ENFORCE(itr != node_args.end(),
|
||||
"Attempting to get index by a name which does not exist.");
|
||||
auto index = std::distance(node_args.begin(), itr);
|
||||
return static_cast<int>(index);
|
||||
}
|
||||
|
||||
int GetNodeInputIndexFromInputName(const Node& node, const std::string& input_name) {
|
||||
return GetIndexFromName(node, input_name, true);
|
||||
}
|
||||
|
||||
int GetNodeOutputIndexFromOutputName(const Node& node, const std::string& output_name) {
|
||||
return GetIndexFromName(node, output_name, false);
|
||||
}
|
||||
|
||||
const std::string& GetNodeInputName(const Node& node, int index) {
|
||||
const auto& inputs = node.InputDefs();
|
||||
ORT_ENFORCE(index >= 0 && static_cast<size_t>(index) < inputs.size(),
|
||||
|
|
|
|||
|
|
@ -76,6 +76,9 @@ const std::string& GetNodeInputName(const Node& node, int index);
|
|||
/** Gets the index of an input arg with the specified input arg name. */
|
||||
int GetNodeInputIndexFromInputName(const Node& node, const std::string& input_name);
|
||||
|
||||
/** Gets the index of an output arg with the specified output arg name. */
|
||||
int GetNodeOutputIndexFromOutputName(const Node& node, const std::string& output_name);
|
||||
|
||||
/** Gets the name of the outgoing NodeArg with the specified index for the given node. */
|
||||
const std::string& GetNodeOutputName(const Node& node, int index);
|
||||
|
||||
|
|
|
|||
|
|
@ -266,16 +266,15 @@ TEST_F(ExecutionFrameTest, MemPatternWithExternalOutputsTest) {
|
|||
onnxruntime::Graph& graph = model.MainGraph();
|
||||
TypeProto tensor_float;
|
||||
tensor_float.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT);
|
||||
onnxruntime::NodeArg input_def("X", &tensor_float),
|
||||
yield_out_def("T", &tensor_float),
|
||||
onnxruntime::NodeArg input_def("X", &tensor_float), yield_out_def("T", &tensor_float),
|
||||
gemm_out_def("Y", &tensor_float);
|
||||
|
||||
ONNX_NAMESPACE::AttributeProto required_grad;
|
||||
const std::string attribute_name = "required_grad";
|
||||
required_grad.set_name(attribute_name);
|
||||
required_grad.set_type(ONNX_NAMESPACE::AttributeProto::INTS);
|
||||
required_grad.add_ints(static_cast<int64_t>(0));
|
||||
NodeAttributes attributes({{attribute_name, required_grad}});
|
||||
ONNX_NAMESPACE::AttributeProto full_shape_outputs;
|
||||
const std::string attribute_name = "full_shape_outputs";
|
||||
full_shape_outputs.set_name(attribute_name);
|
||||
full_shape_outputs.set_type(ONNX_NAMESPACE::AttributeProto::INTS);
|
||||
full_shape_outputs.add_ints(static_cast<int64_t>(0));
|
||||
NodeAttributes attributes({{attribute_name, full_shape_outputs}});
|
||||
graph.AddNode("node1", "YieldOp", "yield", ArgMap{&input_def}, ArgMap{&yield_out_def}, &attributes, kMSDomain)
|
||||
.SetExecutionProviderType(xp_type);
|
||||
// Add another node after YieldOp as YieldOp should not be graph output.
|
||||
|
|
@ -292,8 +291,8 @@ TEST_F(ExecutionFrameTest, MemPatternWithExternalOutputsTest) {
|
|||
|
||||
DataTransferManager dtm;
|
||||
profiling::Profiler profiler;
|
||||
SessionState state(graph, execution_providers, true, &tp_, nullptr, dtm,
|
||||
DefaultLoggingManager().DefaultLogger(), profiler);
|
||||
SessionState state(graph, execution_providers, true, &tp_, nullptr, dtm, DefaultLoggingManager().DefaultLogger(),
|
||||
profiler);
|
||||
|
||||
ASSERT_STATUS_OK(state.FinalizeSessionState(ORT_TSTR(""), kernel_registry_manager));
|
||||
|
||||
|
|
@ -307,12 +306,8 @@ TEST_F(ExecutionFrameTest, MemPatternWithExternalOutputsTest) {
|
|||
auto cpu_allocator = execution_providers.Get(xp_type)->GetAllocator(0, OrtMemTypeDefault);
|
||||
|
||||
OrtValue x_value, t_value;
|
||||
CreateMLValue<float>(cpu_allocator,
|
||||
std::vector<int64_t>{2, 2},
|
||||
std::vector<float>(4, 2.0f), &x_value);
|
||||
CreateMLValue<float>(cpu_allocator,
|
||||
std::vector<int64_t>{2, 2},
|
||||
std::vector<float>(4, 1.0f), &t_value);
|
||||
CreateMLValue<float>(cpu_allocator, std::vector<int64_t>{2, 2}, std::vector<float>(4, 2.0f), &x_value);
|
||||
CreateMLValue<float>(cpu_allocator, std::vector<int64_t>{2, 2}, std::vector<float>(4, 1.0f), &t_value);
|
||||
|
||||
vector<OrtValue> outputs;
|
||||
ExecutionFrame frame({x_idx}, {x_value}, {y_idx}, outputs, {}, state);
|
||||
|
|
@ -322,10 +317,8 @@ TEST_F(ExecutionFrameTest, MemPatternWithExternalOutputsTest) {
|
|||
ASSERT_TRUE(frame.GetMutableNodeInputOrOutputMLValue(t_idx)->IsTensor());
|
||||
|
||||
OrtValue& y_value = *frame.GetMutableNodeInputOrOutputMLValue(y_idx);
|
||||
ASSERT_STATUS_OK(frame.AllocateMLValueTensorSelfOwnBuffer(y_value, y_idx,
|
||||
DataTypeImpl::GetType<float>(),
|
||||
cpu_allocator->Info(),
|
||||
TensorShape(std::vector<int64_t>{2, 2})));
|
||||
ASSERT_STATUS_OK(frame.AllocateMLValueTensorSelfOwnBuffer(
|
||||
y_value, y_idx, DataTypeImpl::GetType<float>(), cpu_allocator->Info(), TensorShape(std::vector<int64_t>{2, 2})));
|
||||
|
||||
MemoryPatternGroup pattern;
|
||||
ASSERT_STATUS_OK(frame.GeneratePatterns(&pattern));
|
||||
|
|
|
|||
|
|
@ -41,7 +41,7 @@ Status ModuleGradientGraphBuilder::Initialize(std::istream& model_istream,
|
|||
}
|
||||
|
||||
training_graph_info_.initializer_names_to_train.assign(config.initializer_names_to_train.begin(),
|
||||
config.initializer_names_to_train.end());
|
||||
config.initializer_names_to_train.end());
|
||||
|
||||
std::vector<const NodeArg*> input_args;
|
||||
for (const auto& input_name : training_graph_info_.user_input_names) {
|
||||
|
|
@ -78,8 +78,8 @@ Status ModuleGradientGraphBuilder::Build(const std::vector<std::vector<int64_t>>
|
|||
// Build the gradient graph.
|
||||
ORT_RETURN_IF_ERROR(BuildGradientGraph());
|
||||
|
||||
// Add Yield Op.
|
||||
AddYieldOp();
|
||||
// Handle user outputs and output grads.
|
||||
HandleOutputsAndGrads();
|
||||
|
||||
// Reorder outputs.
|
||||
ReorderOutputs();
|
||||
|
|
@ -170,59 +170,70 @@ Status ModuleGradientGraphBuilder::BuildGradientGraph() {
|
|||
return Status::OK();
|
||||
}
|
||||
|
||||
void ModuleGradientGraphBuilder::AddYieldOp() {
|
||||
void ModuleGradientGraphBuilder::HandleOutputsAndGrads() {
|
||||
Graph& gradient_graph = gradient_model_->MainGraph();
|
||||
GraphViewer gradient_graph_viewer(gradient_graph);
|
||||
const auto& gradient_node_topology_list = gradient_graph_viewer.GetNodesInTopologicalOrder();
|
||||
std::unordered_set<std::string> user_output_grad_names_set;
|
||||
for (const auto& name : training_graph_info_.user_output_names) {
|
||||
user_output_grad_names_set.insert(name + "_grad");
|
||||
user_output_grad_names_set.insert(GradientBuilderBase::GradientName(name));
|
||||
}
|
||||
|
||||
// If an NodeArg is output of one of nodes, it's not the user output gradient needed by backward graph.
|
||||
std::unordered_set<std::string> non_backward_user_output_grad_names;
|
||||
// If an output gradient is output of one of nodes, need to add this output to PT's output gradient.
|
||||
std::unordered_set<std::string> internal_output_grad_names;
|
||||
for (auto node_index : gradient_node_topology_list) {
|
||||
auto& node = *gradient_graph.GetNode(node_index);
|
||||
for (const auto& node_arg : node.OutputDefs()) {
|
||||
if (user_output_grad_names_set.find(node_arg->Name()) != user_output_grad_names_set.end()) {
|
||||
non_backward_user_output_grad_names.insert(node_arg->Name());
|
||||
internal_output_grad_names.insert(node_arg->Name());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// YieldOps required_grad attribute specifies the indices of the required gradients.
|
||||
ONNX_NAMESPACE::AttributeProto required_grad;
|
||||
const std::string attribute_name = "required_grad";
|
||||
required_grad.set_name(attribute_name);
|
||||
required_grad.set_type(ONNX_NAMESPACE::AttributeProto::INTS);
|
||||
|
||||
training_graph_info_.backward_output_grad_names_map.clear();
|
||||
for (std::size_t i = 0; i < training_graph_info_.user_output_names.size(); ++i) {
|
||||
const auto& name = training_graph_info_.user_output_names[i];
|
||||
std::string grad_name = name + "_grad";
|
||||
if (non_backward_user_output_grad_names.find(grad_name) == non_backward_user_output_grad_names.end()) {
|
||||
training_graph_info_.backward_output_grad_names_map.insert(std::make_pair(grad_name, i));
|
||||
required_grad.add_ints(static_cast<int64_t>(i));
|
||||
}
|
||||
for (const auto& output_grad_name : internal_output_grad_names) {
|
||||
Node* producer_node = gradient_graph.GetMutableProducerNode(output_grad_name);
|
||||
int producer_node_arg_index = graph_utils::GetNodeOutputIndexFromOutputName(*producer_node, output_grad_name);
|
||||
const TypeProto* type_info = producer_node->MutableOutputDefs()[producer_node_arg_index]->TypeAsProto();
|
||||
auto& external_node_arg = gradient_graph.GetOrCreateNodeArg(
|
||||
gradient_graph.GenerateNodeArgName(GradientBuilderBase::ExternalOutputName(output_grad_name)), type_info);
|
||||
auto& output_node_arg = gradient_graph.GetOrCreateNodeArg(
|
||||
gradient_graph.GenerateNodeArgName(output_grad_name + "_add_output"), type_info);
|
||||
Node& add_node = gradient_graph.AddNode(
|
||||
output_grad_name + "_add", "Add", "",
|
||||
{&external_node_arg, producer_node->MutableOutputDefs()[producer_node_arg_index]}, {&output_node_arg});
|
||||
graph_utils::ReplaceDownstreamNodeInput(gradient_graph, *producer_node, producer_node_arg_index, add_node, 0);
|
||||
}
|
||||
|
||||
// YieldOps full_shape_outputs attribute specifies the indices of outputs that must be full shape.
|
||||
// We need this info to set make TypeAndShapeInferenceFunction work properly.
|
||||
ONNX_NAMESPACE::AttributeProto full_shape_outputs;
|
||||
const std::string attribute_name = "full_shape_outputs";
|
||||
full_shape_outputs.set_name(attribute_name);
|
||||
full_shape_outputs.set_type(ONNX_NAMESPACE::AttributeProto::INTS);
|
||||
|
||||
std::vector<NodeArg*> yield_input_node_args;
|
||||
std::vector<NodeArg*> yield_output_node_args;
|
||||
for (const auto& name : training_graph_info_.user_output_names) {
|
||||
training_graph_info_.output_grad_indices_require_full_shape.clear();
|
||||
for (size_t i = 0; i < training_graph_info_.user_output_names.size(); i++) {
|
||||
std::string name = training_graph_info_.user_output_names[i];
|
||||
yield_input_node_args.emplace_back(gradient_graph.GetNodeArg(name));
|
||||
std::string grad_name = GradientBuilderBase::GradientName(name);
|
||||
if (internal_output_grad_names.find(grad_name) != internal_output_grad_names.end()) {
|
||||
grad_name = GradientBuilderBase::ExternalOutputName(grad_name);
|
||||
} else {
|
||||
// If output grad is the direct input of backward graph, we need to materialize it
|
||||
// to a all-0 tensor with same shape of output, otherwise, since it will be an input of
|
||||
// Add node, it's OK to use scalar-0 tensor to save memory.
|
||||
training_graph_info_.output_grad_indices_require_full_shape.emplace_back(i);
|
||||
full_shape_outputs.add_ints(static_cast<int64_t>(i));
|
||||
}
|
||||
|
||||
yield_output_node_args.emplace_back(gradient_graph.GetNodeArg(grad_name));
|
||||
}
|
||||
|
||||
for (const auto& name : training_graph_info_.user_output_names) {
|
||||
std::string grad_name = name + "_grad";
|
||||
auto element = training_graph_info_.backward_output_grad_names_map.find(grad_name);
|
||||
if (element != training_graph_info_.backward_output_grad_names_map.end()) {
|
||||
yield_output_node_args.emplace_back(gradient_graph.GetNodeArg(element->first));
|
||||
}
|
||||
}
|
||||
|
||||
NodeAttributes attributes({{attribute_name, required_grad}});
|
||||
|
||||
gradient_graph.AddNode("YieldOp", "YieldOp", "Yield Op", yield_input_node_args, yield_output_node_args, &attributes, kMSDomain);
|
||||
NodeAttributes attributes({{attribute_name, full_shape_outputs}});
|
||||
gradient_graph.AddNode("YieldOp", "YieldOp", "Yield Op", yield_input_node_args, yield_output_node_args, &attributes,
|
||||
kMSDomain);
|
||||
}
|
||||
|
||||
void ModuleGradientGraphBuilder::ReorderOutputs() {
|
||||
|
|
@ -243,7 +254,7 @@ void ModuleGradientGraphBuilder::ReorderOutputs() {
|
|||
training_graph_info_.user_input_grad_names.clear();
|
||||
for (const auto& input_name : training_graph_info_.user_input_names) {
|
||||
if (user_input_require_grad_set.find(input_name) != user_input_require_grad_set.end()) {
|
||||
std::string input_gradient_name = input_name + "_grad";
|
||||
std::string input_gradient_name = GradientBuilderBase::GradientName(input_name);
|
||||
ORT_ENFORCE(gradient_output_arg_map.find(input_gradient_name) != gradient_output_arg_map.end(),
|
||||
"Required user input grad is not found on gradient graph.");
|
||||
training_graph_info_.user_input_grad_names[input_name] = input_gradient_name;
|
||||
|
|
@ -254,7 +265,7 @@ void ModuleGradientGraphBuilder::ReorderOutputs() {
|
|||
// Add initializer gradients to graph outputs.
|
||||
training_graph_info_.initializer_grad_names_to_train.clear();
|
||||
for (const auto& initializer_name : training_graph_info_.initializer_names_to_train) {
|
||||
std::string initializer_gradient_name = initializer_name + "_grad";
|
||||
std::string initializer_gradient_name = GradientBuilderBase::GradientName(initializer_name);
|
||||
ORT_ENFORCE(gradient_output_arg_map.find(initializer_gradient_name) != gradient_output_arg_map.end(),
|
||||
"Trainable initializer grad is not found on gradient graph.");
|
||||
training_graph_info_.initializer_grad_names_to_train.emplace_back(initializer_gradient_name);
|
||||
|
|
|
|||
|
|
@ -41,9 +41,9 @@ struct TrainingGraphInfo {
|
|||
std::vector<std::string> initializer_grad_names_to_train{};
|
||||
// The user outputs.
|
||||
std::vector<std::string> user_output_names{};
|
||||
// The user output grad names that are actually required by the backward graph
|
||||
// mapped to the index of the correspoinding output of inference graph.
|
||||
std::unordered_map<std::string, size_t> backward_output_grad_names_map{};
|
||||
// Indices of output grads that need to be materialized to full size all-0 tensor.
|
||||
// Otherwise, we can use scalar-0 tensor.
|
||||
std::vector<size_t> output_grad_indices_require_full_shape{};
|
||||
};
|
||||
|
||||
class ModuleGradientGraphBuilder {
|
||||
|
|
@ -83,8 +83,8 @@ class ModuleGradientGraphBuilder {
|
|||
// Build gradient graph.
|
||||
Status BuildGradientGraph();
|
||||
|
||||
// Add Yield Op.
|
||||
void AddYieldOp();
|
||||
// Handle user outputs and output grads.
|
||||
void HandleOutputsAndGrads();
|
||||
|
||||
// Reorder gradient graph outputs.
|
||||
void ReorderOutputs();
|
||||
|
|
|
|||
|
|
@ -71,6 +71,10 @@ class GradientBuilderBase {
|
|||
return name + "_grad";
|
||||
}
|
||||
|
||||
static std::string ExternalOutputName(const std::string& name) {
|
||||
return name + "_external";
|
||||
}
|
||||
|
||||
protected:
|
||||
virtual GradientDef GetGradientDefsImpl() const = 0;
|
||||
|
||||
|
|
|
|||
|
|
@ -2220,27 +2220,26 @@ Return true if all elements are true and false otherwise.
|
|||
.Output(0, "outputs_grad", "Gradient of outputs returned from pytorch.", "T", OpSchema::Variadic,
|
||||
/*is_homogeneous*/ false,
|
||||
/*min_arity*/ 1)
|
||||
.Attr(
|
||||
"required_grad",
|
||||
"The indices of the outputs that require gradient outputs.",
|
||||
AttributeProto::INTS)
|
||||
.Attr("full_shape_outputs", "The indices of the outputs that must have full shape.", AttributeProto::INTS)
|
||||
.TypeConstraint("T", OpSchema::all_tensor_types(), "Allow inputs and outputs to be any kind of tensor.")
|
||||
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
|
||||
const std::string attribute_name = "required_grad";
|
||||
auto required_grads = ctx.getAttribute(attribute_name);
|
||||
if (nullptr == required_grads) { // attribute not present
|
||||
ORT_ENFORCE(ctx.getNumInputs() == ctx.getNumOutputs());
|
||||
for (size_t i = 0; i < ctx.getNumInputs(); ++i) {
|
||||
propagateElemTypeFromInputToOutput(ctx, i, i);
|
||||
}
|
||||
|
||||
const std::string attribute_name = "full_shape_outputs";
|
||||
auto full_shape_outputs = ctx.getAttribute(attribute_name);
|
||||
if (nullptr == full_shape_outputs) { // attribute not present
|
||||
fail_type_inference("Value of attribute ", attribute_name, " not specified");
|
||||
}
|
||||
ORT_ENFORCE(ctx.getNumOutputs() == static_cast<size_t> (required_grads->ints_size()));
|
||||
for (size_t i = 0, n = static_cast<size_t> (required_grads->ints_size()); i < n; ++i) {
|
||||
size_t j = static_cast<size_t> (required_grads->ints(static_cast<int>(i)));
|
||||
ORT_ENFORCE(ctx.getNumInputs() > j);
|
||||
propagateElemTypeFromInputToOutput(ctx, j, i);
|
||||
|
||||
for (size_t i = 0, n = static_cast<size_t>(full_shape_outputs->ints_size()); i < n; ++i) {
|
||||
size_t j = static_cast<size_t>(full_shape_outputs->ints(static_cast<int>(i)));
|
||||
auto typeProto = ctx.getInputType(j);
|
||||
if (!hasShape(*typeProto)) {
|
||||
continue;
|
||||
if (hasShape(*typeProto)) {
|
||||
propagateShapeFromInputToOutput(ctx, j, j);
|
||||
}
|
||||
propagateShapeFromInputToOutput(ctx, j, i);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
|
|
|||
|
|
@ -514,7 +514,7 @@ py::class_<TrainingAgent>(m, "TrainingAgent", R"pbdoc(This is the main class use
|
|||
.def_readwrite("initializer_names_to_train", &TrainingGraphInfo::initializer_names_to_train)
|
||||
.def_readwrite("initializer_grad_names_to_train", &TrainingGraphInfo::initializer_grad_names_to_train)
|
||||
.def_readwrite("user_output_names", &TrainingGraphInfo::user_output_names)
|
||||
.def_readwrite("backward_output_grad_names_map", &TrainingGraphInfo::backward_output_grad_names_map);
|
||||
.def_readwrite("output_grad_indices_require_full_shape", &TrainingGraphInfo::output_grad_indices_require_full_shape);
|
||||
|
||||
py::class_<ModuleGradientGraphBuilder> module_gradient_graph_builder(m, "ModuleGradientGraphBuilder");
|
||||
module_gradient_graph_builder.def(py::init([]() { return onnxruntime::make_unique<ModuleGradientGraphBuilder>(); }))
|
||||
|
|
|
|||
|
|
@ -142,6 +142,11 @@ class ORTModule(torch.nn.Module):
|
|||
user_outputs = tuple(_ort_output_to_torch_tensor(forward_output) for forward_output in forward_outputs)
|
||||
ctx.run_id = run_id
|
||||
|
||||
# Disable materializing grads then None object will not be converted to a tensor filled with zeros prior to calling backward.
|
||||
# Also save shape, device and type info to ctx for materializing tensor in backward if output grad is None.
|
||||
ctx.set_materialize_grads(False)
|
||||
ctx.output_info = [(output.shape, output.device, output.dtype) for output in user_outputs]
|
||||
|
||||
return user_outputs
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -152,20 +157,17 @@ class ORTModule(torch.nn.Module):
|
|||
# Use IO binding
|
||||
# Push user output grads to ONNX backend.
|
||||
backward_grad_output_ortvalue = []
|
||||
|
||||
# backward_output_grad_names_map only contains the subset of module outputs that need a gradient,
|
||||
# we filter out the invalid entries in grad_outputs, accessing using the mapped index.
|
||||
contiguous_grad_outputs = []
|
||||
for i in range(len(grad_outputs)):
|
||||
if i in self._onnx_graphs_info.backward_output_grad_names_map.values():
|
||||
grad_output = grad_outputs[i]
|
||||
if not grad_output.is_contiguous():
|
||||
grad_output = grad_output.contiguous()
|
||||
contiguous_grad_outputs.append(grad_output)
|
||||
# in the original logic, the first grad_output above would be out of scope in next loop, thus memory for
|
||||
# grad_output.data_ptr() in the first call would be corrupted in YieldOp since Torch may reclaim the memory
|
||||
# the solution is to store grad_output in another object, thus memory allocated for the grad_output in the second loop
|
||||
# would be new and will not impact the memory of the first grad_output
|
||||
for idx, grad_output in enumerate(grad_outputs):
|
||||
if grad_output is None:
|
||||
shape, device, dtype = ctx.output_info[idx]
|
||||
if idx in self._onnx_graphs_info.output_grad_indices_require_full_shape:
|
||||
grad_output = torch.zeros(shape, device=device, dtype=dtype)
|
||||
else:
|
||||
grad_output = torch.tensor(0., device=device, dtype=dtype)
|
||||
elif not grad_output.is_contiguous():
|
||||
grad_output = grad_output.contiguous()
|
||||
contiguous_grad_outputs.append(grad_output)
|
||||
for grad_output in contiguous_grad_outputs:
|
||||
backward_grad_output_ortvalue.append(onnxruntime.OrtValue.ortvalue_from_data_ptr(list(grad_output.size()), _utils.dtype_torch_to_numpy(
|
||||
grad_output.dtype), grad_output.device.type, _utils.get_device_index(grad_output.device), grad_output.data_ptr()))
|
||||
|
|
|
|||
|
|
@ -146,7 +146,9 @@ def _get_name(name):
|
|||
return res
|
||||
raise FileNotFoundError("Unable to find '{0}' or '{1}' or '{2}'".format(name, rel, res))
|
||||
|
||||
def assert_gradients_match_and_reset_gradient(ort_model, pt_model, reset_gradient=True, rtol=1e-05, atol=1e-06):
|
||||
# Depending on calling backward() from which outputs, it's possible that grad of some weights are not calculated.
|
||||
# none_pt_params is to tell what these weights are, so we will not compare the tensors.
|
||||
def assert_gradients_match_and_reset_gradient(ort_model, pt_model, none_pt_params=[], reset_gradient=True, rtol=1e-05, atol=1e-06):
|
||||
ort_named_params = list(ort_model.named_parameters())
|
||||
pt_named_params = list(pt_model.named_parameters())
|
||||
assert len(ort_named_params) == len(pt_named_params)
|
||||
|
|
@ -156,7 +158,11 @@ def assert_gradients_match_and_reset_gradient(ort_model, pt_model, reset_gradien
|
|||
pt_name, pt_param = pt_named_param
|
||||
|
||||
assert pt_name in ort_name
|
||||
assert torch.allclose(ort_param.grad, pt_param.grad, rtol=rtol, atol=atol)
|
||||
if pt_name in none_pt_params:
|
||||
assert pt_param.grad is None
|
||||
assert not torch.is_nonzero(torch.count_nonzero(ort_param.grad))
|
||||
else:
|
||||
assert torch.allclose(ort_param.grad, pt_param.grad, rtol=rtol, atol=atol)
|
||||
|
||||
if reset_gradient:
|
||||
ort_param.grad = None
|
||||
|
|
|
|||
|
|
@ -36,9 +36,9 @@ class NeuralNetSinglePositionalArgument(torch.nn.Module):
|
|||
out = self.fc2(out)
|
||||
return out
|
||||
|
||||
class NeuralNetMultiplePositionalArgumentsMultipleOutputs0(torch.nn.Module):
|
||||
class NeuralNetMultiplePositionalArgumentsMultiOutputsWithoutDependency(torch.nn.Module):
|
||||
def __init__(self, input_size, hidden_size, num_classes):
|
||||
super(NeuralNetMultiplePositionalArgumentsMultipleOutputs0, self).__init__()
|
||||
super(NeuralNetMultiplePositionalArgumentsMultiOutputsWithoutDependency, self).__init__()
|
||||
|
||||
self.fc1 = torch.nn.Linear(input_size, hidden_size)
|
||||
self.fc2 = torch.nn.Linear(input_size, hidden_size)
|
||||
|
|
@ -53,9 +53,9 @@ class NeuralNetMultiplePositionalArgumentsMultipleOutputs0(torch.nn.Module):
|
|||
out2 = self.relu2(out2)
|
||||
return out1, out2
|
||||
|
||||
class NeuralNetMultiplePositionalArgumentsMultipleOutputs1(torch.nn.Module):
|
||||
class NeuralNetMultiplePositionalArgumentsMultiOutputsWithDependency(torch.nn.Module):
|
||||
def __init__(self, input_size, hidden_size, num_classes):
|
||||
super(NeuralNetMultiplePositionalArgumentsMultipleOutputs1, self).__init__()
|
||||
super(NeuralNetMultiplePositionalArgumentsMultiOutputsWithDependency, self).__init__()
|
||||
|
||||
self.fc1 = torch.nn.Linear(input_size, hidden_size)
|
||||
self.relu = torch.nn.ReLU()
|
||||
|
|
@ -654,36 +654,70 @@ def test_changes_input_requires_grad_reinitializes_module_gradient_graph_builder
|
|||
@pytest.mark.parametrize("device", ['cuda'])
|
||||
def test_input_requires_grad_backward_creates_input_grad_as_required0(device):
|
||||
N, D_in, H, D_out = 32, 784, 500, 10
|
||||
model = NeuralNetMultiplePositionalArgumentsMultipleOutputs0(D_in, H, D_out).to(device)
|
||||
model = ORTModule(model)
|
||||
x1 = torch.randn(N, D_in, device=device, requires_grad=True)
|
||||
x2 = torch.randn(N, D_in, device=device, requires_grad=True)
|
||||
pt_model = NeuralNetMultiplePositionalArgumentsMultiOutputsWithoutDependency(D_in, H, D_out).to(device)
|
||||
ort_model = ORTModule(copy.deepcopy(pt_model))
|
||||
pt_x1 = torch.randn(N, D_in, device=device, requires_grad=True)
|
||||
pt_x2 = torch.randn(N, D_in, device=device, requires_grad=True)
|
||||
ort_x1 = pt_x1.clone().detach()
|
||||
ort_x2 = pt_x2.clone().detach()
|
||||
ort_x1.requires_grad = True
|
||||
ort_x2.requires_grad = True
|
||||
|
||||
y1, _ = model(x1, x2)
|
||||
s1 = y1.sum()
|
||||
s1.backward()
|
||||
assert x1.grad is not None and x2.grad is not None
|
||||
def run_step0(model, x1, x2):
|
||||
y1, _ = model(x1, x2)
|
||||
s1 = y1.sum()
|
||||
s1.backward() # y2's gradient will be materialized to full shape.
|
||||
return y1
|
||||
|
||||
# named_params[0] and named_params[1] correspond to weight and bias for fc1, similarly
|
||||
# named_params[2] and named_params[3] correspond to weight and bias for fc2.
|
||||
named_params = list(model.named_parameters())
|
||||
assert torch.count_nonzero(named_params[0][1].grad) > 0
|
||||
assert torch.count_nonzero(named_params[1][1].grad) > 0
|
||||
assert named_params[2][1].grad is None or torch.count_nonzero(named_params[2][1].grad) == 0
|
||||
assert named_params[3][1].grad is None or torch.count_nonzero(named_params[3][1].grad) == 0
|
||||
pt_y1 = run_step0(pt_model, pt_x1, pt_x2)
|
||||
ort_y1 = run_step0(ort_model, ort_x1, ort_x2)
|
||||
|
||||
# Reset gradients
|
||||
for param in named_params:
|
||||
param[1].grad = None
|
||||
#assert torch.allclose(pt_y1, ort_y1) # TODO: this assert is failing, need to investigate!!
|
||||
assert torch.allclose(ort_x1.grad, pt_x1.grad)
|
||||
assert torch.allclose(ort_x2.grad, pt_x2.grad)
|
||||
# backward() is from y1, so grad of fc2.weight and fc2.bias will not be calculated.
|
||||
_test_helpers.assert_gradients_match_and_reset_gradient(ort_model, pt_model, none_pt_params=['fc2.weight', 'fc2.bias'], reset_gradient=True)
|
||||
|
||||
_, y2 = model(x1,x2)
|
||||
s2 = y2.sum()
|
||||
s2.backward()
|
||||
named_params = list(model.named_parameters())
|
||||
assert named_params[0][1].grad is None or torch.count_nonzero(named_params[0][1].grad) == 0
|
||||
assert named_params[1][1].grad is None or torch.count_nonzero(named_params[1][1].grad) == 0
|
||||
assert torch.count_nonzero(named_params[2][1].grad) > 0
|
||||
assert torch.count_nonzero(named_params[3][1].grad) > 0
|
||||
def run_step1(model, x1, x2):
|
||||
_, y2 = model(x1, x2)
|
||||
s2 = y2.sum()
|
||||
s2.backward() # y1's gradient will be materialized to full shape.
|
||||
return y2
|
||||
|
||||
pt_y2 = run_step1(pt_model, pt_x1, pt_x2)
|
||||
ort_y2 = run_step1(ort_model, ort_x1, ort_x2)
|
||||
|
||||
#assert torch.allclose(pt_y2, ort_y2) # TODO: this assert is failing, need to investigate!!
|
||||
assert torch.allclose(ort_x1.grad, pt_x1.grad)
|
||||
assert torch.allclose(ort_x2.grad, pt_x2.grad)
|
||||
# backward() is from y2, so grad of fc1.weight and fc1.bias will not be calculated.
|
||||
_test_helpers.assert_gradients_match_and_reset_gradient(ort_model, pt_model, none_pt_params=['fc1.weight', 'fc1.bias'])
|
||||
|
||||
@pytest.mark.parametrize("device", ['cuda'])
|
||||
def test_loss_combines_two_outputs_with_dependency(device):
|
||||
|
||||
def run_step(model, x1, x2):
|
||||
y1, y2 = model(x1, x2)
|
||||
loss = y1.sum() + y2.sum()
|
||||
loss.backward()
|
||||
return y1, y2
|
||||
|
||||
N, D_in, H, D_out = 32, 784, 500, 10
|
||||
pt_model = NeuralNetMultiplePositionalArgumentsMultiOutputsWithDependency(D_in, H, D_out).to(device)
|
||||
ort_model = ORTModule(copy.deepcopy(pt_model))
|
||||
|
||||
pt_x1 = torch.randn(N, D_in, device=device, requires_grad=False)
|
||||
pt_x2 = torch.randn(N, D_in, device=device, requires_grad=False)
|
||||
ort_x1 = pt_x1.clone()
|
||||
ort_x2 = pt_x2.clone()
|
||||
|
||||
# Both y1 and y2's gradients are not None.
|
||||
pt_y1, pt_y2 = run_step(pt_model, pt_x1, pt_x2)
|
||||
ort_y1, ort_y2 = run_step(ort_model, ort_x1, ort_x2)
|
||||
|
||||
#assert torch.allclose(pt_y1, ort_y1) # TODO: this assert is failing, need to investigate!!
|
||||
#assert torch.allclose(pt_y2, ort_y2) # TODO: this assert is failing, need to investigate!!
|
||||
_test_helpers.assert_gradients_match_and_reset_gradient(ort_model, pt_model)
|
||||
|
||||
@pytest.mark.parametrize("x1_requires_grad, x2_requires_grad", [(True, True), (True, False), (False, False), (False, True)])
|
||||
def test_input_requires_grad_backward_creates_input_grad_as_required1(x1_requires_grad, x2_requires_grad):
|
||||
|
|
@ -696,7 +730,7 @@ def test_input_requires_grad_backward_creates_input_grad_as_required1(x1_require
|
|||
|
||||
N, D_in, H, D_out = 32, 784, 500, 10
|
||||
device = 'cuda'
|
||||
pt_model = NeuralNetMultiplePositionalArgumentsMultipleOutputs1(D_in, H, D_out).to(device)
|
||||
pt_model = NeuralNetMultiplePositionalArgumentsMultiOutputsWithDependency(D_in, H, D_out).to(device)
|
||||
ort_model = ORTModule(copy.deepcopy(pt_model))
|
||||
pt_x1 = torch.randn(N, D_in, device=device, requires_grad=x1_requires_grad)
|
||||
pt_x2 = torch.randn(N, D_in, device=device, requires_grad=x2_requires_grad)
|
||||
|
|
|
|||
|
|
@ -9,14 +9,8 @@ namespace onnxruntime {
|
|||
namespace contrib {
|
||||
|
||||
ONNX_OPERATOR_KERNEL_EX(
|
||||
YieldOp,
|
||||
kMSDomain,
|
||||
1,
|
||||
kCpuExecutionProvider,
|
||||
KernelDefBuilder()
|
||||
.TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
|
||||
.ExternalOutputs(),
|
||||
YieldOp);
|
||||
YieldOp, kMSDomain, 1, kCpuExecutionProvider,
|
||||
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()).ExternalOutputs(), YieldOp);
|
||||
|
||||
Status YieldOp::Compute(OpKernelContext* ctx) const {
|
||||
auto* ctx_internal = static_cast<OpKernelContextInternal*>(ctx);
|
||||
|
|
@ -39,6 +33,10 @@ Status YieldOp::Compute(OpKernelContext* ctx) const {
|
|||
} else {
|
||||
ORT_ENFORCE(backward_inputs.second.size() == static_cast<size_t>(ctx->OutputCount()));
|
||||
for (int i = 0; i < ctx->OutputCount(); ++i) {
|
||||
if (std::find(full_shape_outputs_.begin(), full_shape_outputs_.end(), static_cast<int64_t>(i)) !=
|
||||
full_shape_outputs_.end()) {
|
||||
ORT_ENFORCE(ctx->Input<Tensor>(i)->Shape() == backward_inputs.second[i].Get<Tensor>().Shape());
|
||||
}
|
||||
ctx_internal->SetOutputMLValue(i, backward_inputs.second[i]);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -11,8 +11,14 @@ namespace contrib {
|
|||
|
||||
class YieldOp final : public OpKernel {
|
||||
public:
|
||||
YieldOp(const OpKernelInfo& info) : OpKernel(info) {}
|
||||
YieldOp(const OpKernelInfo& info) : OpKernel(info) {
|
||||
ORT_ENFORCE(info.GetAttrs<int64_t>("full_shape_outputs", full_shape_outputs_).IsOK());
|
||||
}
|
||||
|
||||
Status Compute(OpKernelContext* context) const override;
|
||||
|
||||
private:
|
||||
std::vector<int64_t> full_shape_outputs_;
|
||||
};
|
||||
|
||||
} // namespace contrib
|
||||
|
|
|
|||
Loading…
Reference in a new issue