diff --git a/orttraining/orttraining/core/framework/module_gradient_graph_builder.cc b/orttraining/orttraining/core/framework/module_gradient_graph_builder.cc index c2ef02ab0e..99ca1bf2c9 100644 --- a/orttraining/orttraining/core/framework/module_gradient_graph_builder.cc +++ b/orttraining/orttraining/core/framework/module_gradient_graph_builder.cc @@ -22,15 +22,15 @@ Status ModuleGradientGraphBuilder::Initialize(std::istream& model_istream, config_ = config; // Handle original model inputs, outputs and trainable initializers. - // We need to move the trainable initializers to graph inputs and keep the order in config, - // it's possible that the graph already has some trainable initializers in graph inputs, - // so we need to NOT assign these trainable initializers to the user inputs list. + // We need to move all the initializers to graph inputs and keep the order in config, + // it's possible that the graph already has some initializers in graph inputs, + // so we need to NOT assign these initializers to the user inputs list. Graph& graph = model_->MainGraph(); - std::unordered_set initializer_names_to_train_set(config.initializer_names_to_train.begin(), - config.initializer_names_to_train.end()); + std::unordered_set initializer_names(config.initializer_names.begin(), + config.initializer_names.end()); const std::vector& graph_inputs = graph.GetInputsIncludingInitializers(); for (auto& node_arg : graph_inputs) { - if (initializer_names_to_train_set.find(node_arg->Name()) == initializer_names_to_train_set.end()) { + if (initializer_names.find(node_arg->Name()) == initializer_names.end()) { training_graph_info_.user_input_names.emplace_back(node_arg->Name()); } } @@ -42,14 +42,16 @@ 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()); + training_graph_info_.initializer_names.assign(config.initializer_names.begin(), + config.initializer_names.end()); std::vector input_args; for (const auto& input_name : training_graph_info_.user_input_names) { input_args.emplace_back(graph.GetNodeArg(input_name)); } - // 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) { + // Remove all the initializers from the graph and move them to graph inputs. + for (const auto& initializer_name : training_graph_info_.initializer_names) { const NodeArg* node_arg = graph.GetNodeArg(initializer_name); ORT_ENFORCE(node_arg != nullptr); input_args.emplace_back(node_arg); diff --git a/orttraining/orttraining/core/framework/module_gradient_graph_builder.h b/orttraining/orttraining/core/framework/module_gradient_graph_builder.h index 11f849fe7f..3379ab65ef 100644 --- a/orttraining/orttraining/core/framework/module_gradient_graph_builder.h +++ b/orttraining/orttraining/core/framework/module_gradient_graph_builder.h @@ -16,6 +16,8 @@ namespace training { * The training configuration options. */ struct ModuleGradientGraphBuilderConfiguration { + // The names of the weights. + std::vector initializer_names{}; // The names of the weights to train. std::vector initializer_names_to_train{}; // The names of inputs that require gradient. @@ -35,6 +37,8 @@ struct TrainingGraphInfo { std::vector user_input_names{}; // Map from user input names to corresponding user input grad names for those user inputs that require grad. std::unordered_map user_input_grad_names{}; + // All initializers (trainable as well as non trainable). + std::vector initializer_names{}; // Trainable initializers. std::vector initializer_names_to_train{}; // Trainable initializer grad names, ordered according to initializer_names_to_train. @@ -60,9 +64,9 @@ class ModuleGradientGraphBuilder { Status Initialize(std::istream& model_istream, const ModuleGradientGraphBuilderConfiguration& config); /** - * Build the gradient graph and split it to forward and backward graphs. + * Build the gradient graph. * @param input_shapes_ptr The pointer to vector of concrete shapes of the user inputs. - * @return The status of the gradient graph building and forward/backward graphs splitting. + * @return The status of the gradient graph building. */ Status Build(const std::vector>* input_shapes_ptr = nullptr); @@ -73,8 +77,8 @@ class ModuleGradientGraphBuilder { std::string GetGradientModel() const; /** - * Get the split graphs information. - * @return The split graphs information. + * Get the training graphs information. + * @return The training graphs information. */ TrainingGraphInfo GetTrainingGraphInfo() const { return training_graph_info_; } diff --git a/orttraining/orttraining/python/orttraining_pybind_state.cc b/orttraining/orttraining/python/orttraining_pybind_state.cc index bdcbdece4b..61d7d993cb 100644 --- a/orttraining/orttraining/python/orttraining_pybind_state.cc +++ b/orttraining/orttraining/python/orttraining_pybind_state.cc @@ -501,6 +501,7 @@ py::class_(m, "TrainingAgent", R"pbdoc(This is the main class use m, "ModuleGradientGraphBuilderConfiguration", R"pbdoc(Configuration information for module gradient graph builder.)pbdoc"); module_gradient_graph_builder_config.def(py::init()) + .def_readwrite("initializer_names", &ModuleGradientGraphBuilderConfiguration::initializer_names) .def_readwrite("initializer_names_to_train", &ModuleGradientGraphBuilderConfiguration::initializer_names_to_train) .def_readwrite("input_names_require_grad", &ModuleGradientGraphBuilderConfiguration::input_names_require_grad) .def_readwrite("use_invertible_layernorm_grad", @@ -511,6 +512,7 @@ py::class_(m, "TrainingAgent", R"pbdoc(This is the main class use training_graph_info.def(py::init()) .def_readwrite("user_input_names", &TrainingGraphInfo::user_input_names) .def_readwrite("user_input_grad_names", &TrainingGraphInfo::user_input_grad_names) + .def_readwrite("initializer_names", &TrainingGraphInfo::initializer_names) .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) diff --git a/orttraining/orttraining/python/training/ortmodule.py b/orttraining/orttraining/python/training/ortmodule.py index ba303dab93..98b26f29a9 100644 --- a/orttraining/orttraining/python/training/ortmodule.py +++ b/orttraining/orttraining/python/training/ortmodule.py @@ -107,7 +107,9 @@ class ORTModule(torch.nn.Module): Finally, we instantiate the ONNX Runtime InferenceSession. ''' # TODO: using pytorch for evaluation for now. We will use ORT for evaluation later. - if not self._is_training: + # TODO: If the model is being executed with the gradient disabled (inside torch.no_grad() context for example), + # leverage pytorch model for now. + if not self._is_training(): return self._original_module(*inputs, **kwargs) # Exporting module to ONNX for the first time @@ -122,16 +124,25 @@ class ORTModule(torch.nn.Module): self._get_inference_graph_and_init_gradient_graph_builder( *inputs, **kwargs) + # Flag to indicate whether the gradient_graph needs to be built + build_gradient_graph = self._current_input_shape is None _, _, input_names_require_grad, new_input_shape = \ _ortmodule_io.parse_inputs_for_onnx_export( self._original_module_parameters, self._onnx_inference, *inputs, **kwargs) + initializer_names_to_train_set_user_model = {name for name, param in + self._flattened_output_module.named_parameters() if param.requires_grad} + initializer_names_to_train_set_onnx_graph = set(self._onnx_graphs_info.initializer_names_to_train) \ + if self._onnx_graphs_info else None # If inputs requiring gradient change from forward to the next, the module_gradient_graph_builder # needs to be reinitialized so it can compute the backward output for the new inputs that require_grad - if input_names_require_grad != self._input_names_require_grad: + if input_names_require_grad != self._input_names_require_grad or \ + initializer_names_to_train_set_user_model != initializer_names_to_train_set_onnx_graph: self._input_names_require_grad = input_names_require_grad self._initialize_module_gradient_graph_builder() + # Trigger the rebuilding of the gradient graph + build_gradient_graph = True - if self._current_input_shape is None: + if build_gradient_graph: self._current_input_shape = new_input_shape self._build_training_graph() self._create_training_session() @@ -238,9 +249,19 @@ class ORTModule(torch.nn.Module): # input_name is not found in the self._input_names_require_grad list # Append None to results for each input that did not require grad results.append(None) + # Append gradients of initializer to results - results += [_ortvalue_to_torch_tensor(backward_output) - for backward_output in backward_outputs[num_user_input_grads:]] + # Go over each initializer, check if it required grad and append to results accordingly + initializer_names_to_train_set = set(self._onnx_graphs_info.initializer_names_to_train) \ + if self._onnx_graphs_info else None + initializer_index = num_user_input_grads + for initializer_name in self._onnx_graphs_info.initializer_names: + if initializer_name in initializer_names_to_train_set: + results.append(_ortvalue_to_torch_tensor(backward_outputs[initializer_index])) + initializer_index += 1 + else: + results.append(None) + # The OrtValue has a shared_ptr to the data. # At this point there are two shared_ptrs to the data, one through the # OrtValue in the output iobinding, and the other through the copy in OrtDLManagedTensor. @@ -287,7 +308,6 @@ class ORTModule(torch.nn.Module): "The model's forward method has **kwargs parameter which is currently not supported.") self._onnx_inference = None - self._is_training = True # Related to training graph shape inference self._current_input_shape = None @@ -297,6 +317,7 @@ class ORTModule(torch.nn.Module): self._module_gradient_graph_builder = None self._input_names_require_grad = None self._original_module_output_schema = None + self._onnx_graphs_info = None # Training model self._onnx_training = None @@ -325,18 +346,26 @@ class ORTModule(torch.nn.Module): self._torch_alloc = self._torch_cuda_allocator.cuda_caching_allocator_raw_alloc_address() self._torch_free = self._torch_cuda_allocator.cuda_caching_allocator_raw_delete_address() + def _is_training(self): + return self._flattened_output_module.training and torch.is_grad_enabled() + def _initialize_module_gradient_graph_builder(self): # TODO: PyTorch exporter bug: changes the initializer order in ONNX model - initializer_names = [p[0] - for p in self._flattened_output_module.named_parameters()] + initializer_names = [name + for name, _ in self._flattened_output_module.named_parameters()] + initializer_names_to_train = [] + if self._is_training(): + initializer_names_to_train = [name + for name, param in self._flattened_output_module.named_parameters() if param.requires_grad] onnx_initializer_names = { p.name for p in self._onnx_inference.graph.initializer} - initializer_names = [ - p for p in initializer_names if p in onnx_initializer_names] + initializer_names_to_train = [ + p for p in initializer_names_to_train if p in onnx_initializer_names] # Build full training graph grad_builder_config = C.ModuleGradientGraphBuilderConfiguration() - grad_builder_config.initializer_names_to_train = initializer_names + grad_builder_config.initializer_names = initializer_names + grad_builder_config.initializer_names_to_train = initializer_names_to_train grad_builder_config.input_names_require_grad = self._input_names_require_grad self._module_gradient_graph_builder = C.ModuleGradientGraphBuilder() self._module_gradient_graph_builder.initialize( @@ -402,11 +431,9 @@ class ORTModule(torch.nn.Module): self._save_onnx_prefix + '_training.onnx') def eval(self: T) -> T: - self._is_training = False self._flattened_output_module.eval() def train(self: T, mode: bool = True) -> T: - self._is_training = mode self._flattened_output_module.train(mode) def _convert_training_graph_input_to_list(self, *inputs, **kwargs): diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index f2a749b61d..90b8bf419a 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -174,6 +174,19 @@ class NeuralNetNonDifferentiableOutput(torch.nn.Module): return out1, mask1, out2, mask2 # intentionally place the non-differentiable output in the middle +class NeuralNetPartialNoGradModel(torch.nn.Module): + def __init__(self, input_size, hidden_size, num_classes): + super(NeuralNetPartialNoGradModel, self).__init__() + + self.fc1 = torch.nn.Linear(input_size, hidden_size).requires_grad_(False) + self.relu = torch.nn.ReLU() + self.fc2 = torch.nn.Linear(hidden_size, num_classes) + + def forward(self, model_input): + out = self.relu(self.fc1(model_input)) + out = self.fc2(out) + return out + # TODO: This is a workaround for the problem that pytest is still cleaning up the previous test # while the next task already start. @pytest.fixture(autouse=True) @@ -1404,3 +1417,59 @@ def test_uint8_input_and_output(): assert y1 is not None assert y2 is not None and y2.dtype == torch.uint8 + +def test_model_partially_requires_grad(): + device = 'cuda' + N, D_in, H, D_out = 64, 784, 500, 10 + model = NeuralNetPartialNoGradModel(D_in, H, D_out).to(device) + model = ORTModule(model) + x = torch.randn(N, D_in, device=device) + + # Make sure no exception is raised + output = model(x) + + loss = torch.sum(output) + loss.backward() + +def test_model_wrapped_inside_torch_no_grad(): + device = 'cuda' + N, D_in, H, D_out = 64, 784, 500, 10 + model = NeuralNetSinglePositionalArgument(D_in, H, D_out).to(device) + model = ORTModule(model) + x = torch.randn(N, D_in, device=device) + + # Make sure no exception is raised + with torch.no_grad(): + output = model(x) + +def test_model_initializer_requires_grad_changes_from_one_forward_to_next(): + device = 'cuda' + N, D_in, H, D_out = 64, 784, 500, 10 + model = NeuralNetPartialNoGradModel(D_in, H, D_out).to(device) + model.fc1.requires_grad_(True) + model = ORTModule(model) + x = torch.randn(N, D_in, device=device) + assert model._original_module.fc1.weight.grad is None + assert model._original_module.fc1.bias.grad is None + + # Make sure no exception is raised + output = model(x) + loss = torch.sum(output) + loss.backward() + training_session1 = model._training_session + weight_grad_2 = model._original_module.fc1.weight.grad + bias_grad_2 = model._original_module.fc1.bias.grad + assert weight_grad_2 is not None + assert bias_grad_2 is not None + + model._original_module.fc1.requires_grad_(False) + output = model(x) + loss = torch.sum(output) + loss.backward() + training_session2 = model._training_session + weight_grad_3 = model._original_module.fc1.weight.grad + bias_grad_3 = model._original_module.fc1.bias.grad + + assert training_session1 != training_session2 + assert torch.equal(weight_grad_2, weight_grad_3) + assert torch.equal(bias_grad_2, bias_grad_3)