From fd21c84eb814371ea0dd86b4902dd137846e07c9 Mon Sep 17 00:00:00 2001 From: ashbhandare Date: Thu, 14 Jan 2021 10:07:54 -0800 Subject: [PATCH] Enable graph save for orttrainer (#6333) * Enable graph save for orttrainer * Fix CI * Update orttraining/orttraining/python/training/orttrainer_options.py * Update orttraining/orttraining/python/training/orttrainer_options.py * Update orttraining/orttraining/python/training/orttrainer_options.py * Update orttraining/orttraining/python/training/orttrainer_options.py * Update orttraining/orttraining/python/training/orttrainer_options.py Co-authored-by: Thiago Crepaldi --- .../core/session/training_session.cc | 17 +++++-- .../core/session/training_session.h | 4 ++ .../python/orttraining_pybind_state.cc | 18 +++++++ .../orttraining/python/training/orttrainer.py | 4 ++ .../python/training/orttrainer_options.py | 48 +++++++++++++++++ .../orttraining_test_orttrainer_frontend.py | 51 ++++++++++++++++++- 6 files changed, 138 insertions(+), 4 deletions(-) diff --git a/orttraining/orttraining/core/session/training_session.cc b/orttraining/orttraining/core/session/training_session.cc index 7934b601f6..55b5375792 100644 --- a/orttraining/orttraining/core/session/training_session.cc +++ b/orttraining/orttraining/core/session/training_session.cc @@ -257,6 +257,11 @@ Status TrainingSession::ConfigureForTraining( config_result.mixed_precision_config_result = mp_result; } + if (IsRootNode(config) && config.model_with_loss_function_path.has_value()) { + ORT_IGNORE_RETURN_VALUE(Save( + config.model_with_loss_function_path.value(), SaveOption::NO_RELOAD)); + } + // We need to get trainable weights to prevent constant folding from them. This works well if trainable weights are passed from config. // For case we use GetTrainableModelInitializers to get trainable weights such as C++ frontend, it may get more initializers // than trainable weights here as it's before transformers. So the constant folding may miss some nodes we actually can fold. @@ -277,10 +282,11 @@ Status TrainingSession::ConfigureForTraining( ORT_RETURN_IF_ERROR(ApplyModelParallelTransformationsToMainGraph(trainable_initializers, config_result)); weight_partition_info_ = config_result.weight_partition_info; - - if (IsRootNode(config) && config.model_with_loss_function_path.has_value()) { + + // Save the model after graph transformations + if (IsRootNode(config) && config.model_after_graph_transforms_path.has_value()) { ORT_IGNORE_RETURN_VALUE(Save( - config.model_with_loss_function_path.value(), SaveOption::NO_RELOAD)); + config.model_after_graph_transforms_path.value(), SaveOption::NO_RELOAD)); } // derive actual set of weights to train @@ -318,6 +324,11 @@ Status TrainingSession::ConfigureForTraining( ORT_RETURN_IF_ERROR(BuildGradientGraph( weight_names_to_train, loss_name, config.gradient_graph_config, *session_logger_)); + + if (IsRootNode(config) && config.model_with_gradient_graph_path.has_value()) { + ORT_IGNORE_RETURN_VALUE(Save( + config.model_with_gradient_graph_path.value(), SaveOption::NO_RELOAD)); + } if (config.pipeline_config.has_value()) { TrainingConfigurationResult::PipelineConfigurationResult pipeline_result{}; diff --git a/orttraining/orttraining/core/session/training_session.h b/orttraining/orttraining/core/session/training_session.h index 315d2929b6..4a974baa9a 100644 --- a/orttraining/orttraining/core/session/training_session.h +++ b/orttraining/orttraining/core/session/training_session.h @@ -46,6 +46,10 @@ class TrainingSession : public InferenceSession { struct TrainingConfiguration { // The path at which to save the intermediate model with the added loss function. optional model_with_loss_function_path{}; + // The path at which to save the model after applying the graph transformations. + optional model_after_graph_transforms_path{}; + // The path at which to save the model with gradient graph added. + optional model_with_gradient_graph_path{}; // The path at which to save the intermediate model with the whole training graph. optional model_with_training_graph_path{}; diff --git a/orttraining/orttraining/python/orttraining_pybind_state.cc b/orttraining/orttraining/python/orttraining_pybind_state.cc index c6f2f2e159..666dae31c4 100644 --- a/orttraining/orttraining/python/orttraining_pybind_state.cc +++ b/orttraining/orttraining/python/orttraining_pybind_state.cc @@ -56,6 +56,11 @@ struct TrainingParameters { bool transformer_layer_recompute = false; int number_recompute_layers = 0; bool enable_adasum = false; + + // graph dumping + std::string model_after_graph_transforms_path; + std::string model_with_gradient_graph_path; + std::string model_with_training_graph_path; }; struct TrainingConfigurationResult { @@ -162,6 +167,16 @@ TrainingConfigurationResult ConfigureSessionForTraining( config.graph_transformer_config.transformer_layer_recompute = parameters.transformer_layer_recompute; config.graph_transformer_config.number_recompute_layers = parameters.number_recompute_layers; + if (!parameters.model_after_graph_transforms_path.empty()) { + config.model_after_graph_transforms_path = parameters.model_after_graph_transforms_path; + } + if (!parameters.model_with_gradient_graph_path.empty()) { + config.model_with_gradient_graph_path = parameters.model_with_gradient_graph_path; + } + if (!parameters.model_with_training_graph_path.empty()) { + config.model_with_training_graph_path = parameters.model_with_training_graph_path; + } + training::TrainingSession::TrainingConfigurationResult config_result{}; OrtPybindThrowIfError(sess->ConfigureForTraining(config, config_result)); @@ -260,6 +275,9 @@ void addObjectMethodsForTraining(py::module& m) { } parameters.optimizer_initial_state = optim_state; }) + .def_readwrite("model_after_graph_transforms_path", &TrainingParameters::model_after_graph_transforms_path) + .def_readwrite("model_with_gradient_graph_path", &TrainingParameters::model_with_gradient_graph_path) + .def_readwrite("model_with_training_graph_path", &TrainingParameters::model_with_training_graph_path) .def_readwrite("enable_adasum", &TrainingParameters::enable_adasum); #if defined(USE_MPI) diff --git a/orttraining/orttraining/python/training/orttrainer.py b/orttraining/orttraining/python/training/orttrainer.py index 4de0a9c210..6b0dde32d7 100644 --- a/orttraining/orttraining/python/training/orttrainer.py +++ b/orttraining/orttraining/python/training/orttrainer.py @@ -637,6 +637,10 @@ class ORTTrainer(object): ort_parameters.transformer_layer_recompute = self.options.graph_transformer.transformer_layer_recompute ort_parameters.number_recompute_layers = self.options.graph_transformer.number_recompute_layers + ort_parameters.model_after_graph_transforms_path = self.options.debug.graph_save_paths.model_after_graph_transforms_path + ort_parameters.model_with_gradient_graph_path = self.options.debug.graph_save_paths.model_with_gradient_graph_path + ort_parameters.model_with_training_graph_path = self.options.debug.graph_save_paths.model_with_training_graph_path + # SessionOptions session_options = ort.SessionOptions() session_options.use_deterministic_compute = self.options.debug.deterministic_compute diff --git a/orttraining/orttraining/python/training/orttrainer_options.py b/orttraining/orttraining/python/training/orttrainer_options.py index 02d76f1c62..98cc0f0579 100644 --- a/orttraining/orttraining/python/training/orttrainer_options.py +++ b/orttraining/orttraining/python/training/orttrainer_options.py @@ -176,6 +176,25 @@ class ORTTrainerOptions(object): 'type' : 'boolean', 'default' : False }, + 'graph_save_paths' : { + 'type' : 'dict', + 'default': {}, + 'required': False, + 'schema': { + 'model_after_graph_transforms_path': { + 'type': 'string', + 'default': '' + }, + 'model_with_gradient_graph_path':{ + 'type': 'string', + 'default': '' + }, + 'model_with_training_graph_path': { + 'type': 'string', + 'default': '' + } + } + } } }, '_internal_use' : { @@ -273,6 +292,16 @@ class ORTTrainerOptions(object): debug.check_model_export (bool, default is False) compares PyTorch model outputs with ONNX model outputs in inference before the first train step to ensure successful model export + debug.graph_save_paths (dict): + paths used for dumping ONNX graphs for debugging purposes + debug.graph_save_paths.model_after_graph_transforms_path (str, default is "") + path to export the ONNX graph after training-related graph transforms have been applied. + No output when it is empty. + debug.graph_save_paths.model_with_gradient_graph_path (str, default is "") + path to export the ONNX graph with the gradient graph added. No output when it is empty. + debug.graph_save_paths.model_with_training_graph_path (str, default is "") + path to export the training ONNX graph with forward, gradient and optimizer nodes. + No output when it is empty. _internal_use (dict): internal options, possibly undocumented, that might be removed without notice _internal_use.enable_internal_postprocess (bool, default is True): @@ -530,6 +559,25 @@ _ORTTRAINER_OPTIONS_SCHEMA = { 'type': 'boolean', 'default': False }, + 'graph_save_paths' : { + 'type' : 'dict', + 'default_setter': lambda _: {}, + 'required': False, + 'schema': { + 'model_after_graph_transforms_path': { + 'type': 'string', + 'default': '' + }, + 'model_with_gradient_graph_path':{ + 'type': 'string', + 'default': '' + }, + 'model_with_training_graph_path': { + 'type': 'string', + 'default': '' + } + } + }, } }, '_internal_use': { diff --git a/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py b/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py index 78b51197fd..7958720b89 100644 --- a/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py +++ b/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py @@ -5,6 +5,7 @@ from numpy.testing import assert_allclose import onnx import os import pytest +import tempfile import torch import torch.nn.functional as F @@ -73,7 +74,12 @@ def testORTTrainerOptionsDefaultValues(test_input): }, 'debug': { 'deterministic_compute': False, - 'check_model_export': False + 'check_model_export': False, + 'graph_save_paths' : { + 'model_after_graph_transforms_path': '', + 'model_with_gradient_graph_path': '', + 'model_with_training_graph_path': '' + } }, '_internal_use': { 'enable_internal_postprocess': True, @@ -1409,3 +1415,46 @@ def testORTTrainerUnusedInput(): trainer.train_step(torch.FloatTensor([1.0]), torch.FloatTensor([1.0])) except RuntimeError: pytest.fail("RuntimeError doing train_step with an unused input.") + +@pytest.mark.parametrize("debug_files", [ + ({'model_after_graph_transforms_path': 'transformed.onnx', + 'model_with_gradient_graph_path': 'transformed_grad.onnx', + 'model_with_training_graph_path': 'training.onnx' + }), + ({'model_after_graph_transforms_path': 'transformed.onnx', + 'model_with_training_graph_path': '' + }), + ]) +def testTrainingGraphExport(debug_files): + device = 'cuda' + model, model_desc, my_loss, batcher_fn, train_data, _, _ = _test_commons._load_pytorch_transformer_model(device) + + with tempfile.TemporaryDirectory() as tempdir: + debug_paths = {} + for k,v in debug_files.items(): + debug_paths[k] = os.path.join(tempdir, v) + opts = orttrainer.ORTTrainerOptions( + { + "device": {"id": device}, + "debug": {"graph_save_paths": debug_paths} + } + ) + optim_config = optim.AdamConfig() + trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, loss_fn=my_loss, options=opts) + data, targets = batcher_fn(train_data, 0) + trainer.train_step(data, targets) + for k,v in debug_files.items(): + path = debug_paths[k] + if len(v) > 0: + assert os.path.isfile(path) + saved_graph = onnx.load(path).graph + if k == 'model_with_training_graph_path': + assert any("AdamOptimizer" in n.op_type for n in saved_graph.node) + elif k == 'model_with_gradient_graph_path': + assert any("Grad" in n.name for n in saved_graph.node) + elif k == 'model_after_graph_transforms_path': + assert any("LayerNormalization" in n.op_type for n in saved_graph.node) + # remove saved file + os.remove(path) + else: + assert not os.path.isfile(path)