diff --git a/orttraining/orttraining/core/framework/module_gradient_graph_builder.cc b/orttraining/orttraining/core/framework/module_gradient_graph_builder.cc index c43c759f63..2e9b33112e 100644 --- a/orttraining/orttraining/core/framework/module_gradient_graph_builder.cc +++ b/orttraining/orttraining/core/framework/module_gradient_graph_builder.cc @@ -50,7 +50,9 @@ Status ModuleGradientGraphBuilder::Initialize(std::istream& model_istream, // Remove the training initializers from the graph and move them to graph inputs. for (const auto& initializer_name : training_graph_info_.initializer_names_to_train) { - input_args.emplace_back(graph.GetNodeArg(initializer_name)); + const NodeArg* node_arg = graph.GetNodeArg(initializer_name); + ORT_ENFORCE(node_arg != nullptr); + input_args.emplace_back(node_arg); graph.RemoveInitializedTensor(initializer_name); } @@ -188,42 +190,35 @@ void ModuleGradientGraphBuilder::AddYieldOp() { } } - // Yield inputs include all user outputs, those require output gradients come first, so Yield Op can use their shapes - // to infer Op output shapes. - std::vector user_output_names_require_grad; - std::vector user_output_names_no_grad; - training_graph_info_.backward_output_grad_names.clear(); - for (const auto& name : training_graph_info_.user_output_names) { + // 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()) { - user_output_names_require_grad.emplace_back(name); - training_graph_info_.backward_output_grad_names.emplace_back(grad_name); - } else { - user_output_names_no_grad.emplace_back(name); + training_graph_info_.backward_output_grad_names_map.insert(std::make_pair(grad_name, i)); + required_grad.add_ints(static_cast(i)); } } - // Reorder the user outputs. - training_graph_info_.user_output_names.clear(); - for (const auto& name : user_output_names_require_grad) { - training_graph_info_.user_output_names.emplace_back(name); - } - - for (const auto& name : user_output_names_no_grad) { - training_graph_info_.user_output_names.emplace_back(name); - } - std::vector yield_input_node_args; std::vector yield_output_node_args; for (const auto& name : training_graph_info_.user_output_names) { yield_input_node_args.emplace_back(gradient_graph.GetNodeArg(name)); } - for (const auto& name : training_graph_info_.backward_output_grad_names) { - yield_output_node_args.emplace_back(gradient_graph.GetNodeArg(name)); + for (const auto& element : training_graph_info_.backward_output_grad_names_map) { + yield_output_node_args.emplace_back(gradient_graph.GetNodeArg(element.first)); } - gradient_graph.AddNode("YieldOp", "YieldOp", "Yield Op", yield_input_node_args, yield_output_node_args, {}, kMSDomain); + NodeAttributes attributes({{attribute_name, required_grad}}); + + gradient_graph.AddNode("YieldOp", "YieldOp", "Yield Op", yield_input_node_args, yield_output_node_args, &attributes, kMSDomain); } void ModuleGradientGraphBuilder::ReorderOutputs() { diff --git a/orttraining/orttraining/core/framework/module_gradient_graph_builder.h b/orttraining/orttraining/core/framework/module_gradient_graph_builder.h index 36b5ca422d..20b01d5b3b 100644 --- a/orttraining/orttraining/core/framework/module_gradient_graph_builder.h +++ b/orttraining/orttraining/core/framework/module_gradient_graph_builder.h @@ -41,8 +41,9 @@ struct TrainingGraphInfo { std::vector initializer_grad_names_to_train{}; // The user outputs. std::vector user_output_names{}; - // The user output grad names that are actual required by the backward graph. - std::vector backward_output_grad_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 backward_output_grad_names_map{}; }; class ModuleGradientGraphBuilder { diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc index c5326cae80..7b14dbc83d 100644 --- a/orttraining/orttraining/core/graph/training_op_defs.cc +++ b/orttraining/orttraining/core/graph/training_op_defs.cc @@ -2220,19 +2220,27 @@ 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) .TypeConstraint("T", OpSchema::all_tensor_types(), "Allow inputs and outputs to be any kind of tensor.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - // Assume the outputs and gradients are one-to-one matching - // TODO: The contrain is relaxed for now - // ORT_ENFORCE(ctx.getNumInputs() == ctx.getNumOutputs(), "Yield op doesn't have the same number of inputs and output"); - - for (size_t i = 0; i < ctx.getNumOutputs(); ++i) { - propagateElemTypeFromInputToOutput(ctx, i, i); - auto typeProto = ctx.getInputType(i); + const std::string attribute_name = "required_grad"; + auto required_grads = ctx.getAttribute(attribute_name); + if (nullptr == required_grads) { // attribute not present + fail_type_inference("Value of attribute ", attribute_name, " not specified"); + } + ORT_ENFORCE(ctx.getNumOutputs() == static_cast (required_grads->ints_size())); + for (size_t i = 0, n = static_cast (required_grads->ints_size()); i < n; ++i) { + size_t j = static_cast (required_grads->ints(static_cast(i))); + ORT_ENFORCE(ctx.getNumInputs() > j); + propagateElemTypeFromInputToOutput(ctx, j, i); + auto typeProto = ctx.getInputType(j); if (!hasShape(*typeProto)) { continue; } - propagateShapeFromInputToOutput(ctx, i, i); + propagateShapeFromInputToOutput(ctx, j, i); } }); } diff --git a/orttraining/orttraining/python/orttraining_pybind_state.cc b/orttraining/orttraining/python/orttraining_pybind_state.cc index 93c2f67e72..02b3cae098 100644 --- a/orttraining/orttraining/python/orttraining_pybind_state.cc +++ b/orttraining/orttraining/python/orttraining_pybind_state.cc @@ -491,7 +491,7 @@ void addObjectMethodsForTraining(py::module& m) { .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", &TrainingGraphInfo::backward_output_grad_names); + .def_readwrite("backward_output_grad_names_map", &TrainingGraphInfo::backward_output_grad_names_map); py::class_ module_gradient_graph_builder(m, "ModuleGradientGraphBuilder"); module_gradient_graph_builder.def(py::init([]() { return onnxruntime::make_unique(); })) diff --git a/orttraining/orttraining/python/training/ortmodule.py b/orttraining/orttraining/python/training/ortmodule.py index 7737abd090..2813e28948 100644 --- a/orttraining/orttraining/python/training/ortmodule.py +++ b/orttraining/orttraining/python/training/ortmodule.py @@ -258,11 +258,14 @@ class ORTModule(torch.nn.Module): # Use IO binding # Push user output grads to ONNX backend. backward_grad_output_ortvalue = [] - for grad_output in grad_outputs[:len(self._onnx_graphs_info.backward_output_grad_names)]: - # Force torch tensors to be contiguous before converting into 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. + + for _, i in self._onnx_graphs_info.backward_output_grad_names_map.items(): + grad_output = grad_outputs[i] if not grad_output.is_contiguous(): grad_output = grad_output.contiguous() - 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())) diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index 940e269d1e..ab70582b2a 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -34,6 +34,38 @@ class NeuralNetSinglePositionalArgument(torch.nn.Module): out = self.fc2(out) return out +class NeuralNetMultiplePositionalArgumentsMultipleOutputs0(torch.nn.Module): + def __init__(self, input_size, hidden_size, num_classes): + super(NeuralNetMultiplePositionalArgumentsMultipleOutputs0, self).__init__() + + self.fc1 = torch.nn.Linear(input_size, hidden_size) + self.fc2 = torch.nn.Linear(input_size, hidden_size) + self.relu1 = torch.nn.ReLU() + self.relu2 = torch.nn.ReLU() + + def forward(self, input1, input2): + model_input = input1 + input2 + out1 = self.fc1(model_input) + out2 = self.fc2(model_input) + out1 = self.relu1(out1) + out2 = self.relu2(out2) + return out1, out2 + +class NeuralNetMultiplePositionalArgumentsMultipleOutputs1(torch.nn.Module): + def __init__(self, input_size, hidden_size, num_classes): + super(NeuralNetMultiplePositionalArgumentsMultipleOutputs1, self).__init__() + + self.fc1 = torch.nn.Linear(input_size, hidden_size) + self.relu = torch.nn.ReLU() + self.fc2 = torch.nn.Linear(hidden_size, num_classes) + + def forward(self, input1, input2): + model_input = input1 + input2 + out1 = self.fc1(model_input) + out1 = self.relu(out1) + out2 = self.fc2(out1) + return out1, out2 + class NeuralNetMultiplePositionalArguments(torch.nn.Module): def __init__(self, input_size, hidden_size, num_classes): super(NeuralNetMultiplePositionalArguments, self).__init__() @@ -550,6 +582,71 @@ def test_changes_input_requires_grad_reinitializes_module_gradient_graph_builder model(x) assert module_gradient_graph_builder != model._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) + + y1, _ = model(x1, x2) + s1 = y1.sum() + s1.backward() + assert x1.grad is not None and x2.grad is not None + + # 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 + + # Reset gradients + for param in named_params: + param[1].grad = None + + _, 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 + + + +@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): + + def run_step(model, x1, x2): + y1, y2 = model(x1, x2) + s = y2.sum() + s.backward() + + N, D_in, H, D_out = 32, 784, 500, 10 + device = 'cuda' + pt_model = NeuralNetMultiplePositionalArgumentsMultipleOutputs1(D_in, H, D_out).to(device) + ort_model = ORTModule(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) + + ort_x1 = pt_x1.clone().detach() + ort_x2 = pt_x2.clone().detach() + ort_x1.requires_grad = x1_requires_grad + ort_x2.requires_grad = x2_requires_grad + + run_step(pt_model, pt_x1, pt_x2) + run_step(ort_model, ort_x1, ort_x2) + + assert not x1_requires_grad or ort_x1.grad is not None + assert not x2_requires_grad or ort_x2.grad is not None + assert not x1_requires_grad or torch.allclose(ort_x1.grad, pt_x1.grad) + assert not x2_requires_grad or torch.allclose(ort_x2.grad, pt_x2.grad) + _test_helpers.assert_gradients_match_and_reset_gradient(ort_model, pt_model) + def test_gpu_reserved_memory_with_torch_no_grad(): device = 'cuda'