From 9835b46a1df5c1dfeb634c61668c644390577d47 Mon Sep 17 00:00:00 2001 From: Ryota Tomioka Date: Wed, 27 Jan 2021 07:39:46 +0000 Subject: [PATCH] Add an option to save the training graph after optimization (#6410) * expose optimized_model_filepath in SessionOptions as `debug.graph_save_paths.model_with_training_graph_after_optimization_path` in `ORTTrainerOptions` --- orttraining/orttraining/python/ort_trainer.py | 15 ++++++++++++--- .../orttraining/python/training/orttrainer.py | 2 ++ .../python/training/orttrainer_options.py | 18 ++++++++++++++---- .../orttraining_test_orttrainer_frontend.py | 16 ++++++++++------ 4 files changed, 38 insertions(+), 13 deletions(-) diff --git a/orttraining/orttraining/python/ort_trainer.py b/orttraining/orttraining/python/ort_trainer.py index b0f56d5086..b0ff0e11e4 100644 --- a/orttraining/orttraining/python/ort_trainer.py +++ b/orttraining/orttraining/python/ort_trainer.py @@ -393,7 +393,8 @@ def create_ort_training_session_with_optimizer(model, device, training_optimizer frozen_weights=[], opset_version=DEFAULT_OPSET_VERSION, use_deterministic_compute=False, use_invertible_layernorm_grad=False, - enable_adasum=False): + enable_adasum=False, + optimized_model_filepath=""): output_name = model.graph.output[0].name ort_parameters = ort.TrainingParameters() ort_parameters.loss_output_name = output_name @@ -459,6 +460,8 @@ def create_ort_training_session_with_optimizer(model, device, training_optimizer sessionOptions = ort.SessionOptions() sessionOptions.use_deterministic_compute = use_deterministic_compute + if len(optimized_model_filepath) > 0: + sessionOptions.optimized_model_filepath = optimized_model_filepath session = ort.TrainingSession(model.SerializeToString(), ort_parameters, sessionOptions) train_io_binding = session.io_binding() eval_io_binding = session.io_binding() @@ -548,7 +551,8 @@ class ORTTrainer(): global_step=0, get_lr_this_step=None, loss_scaler=None, deepspeed_zero_stage=0, enable_grad_norm_clip=True, frozen_weights=[], _opset_version=DEFAULT_OPSET_VERSION, _enable_internal_postprocess=True, _extra_postprocess=None, _use_deterministic_compute=False, - use_invertible_layernorm_grad=False, run_symbolic_shape_infer=False, enable_adasum=False): + use_invertible_layernorm_grad=False, run_symbolic_shape_infer=False, enable_adasum=False, + optimized_model_filepath=""): super(ORTTrainer, self).__init__() """ Initialize ORTTrainer. @@ -618,6 +622,8 @@ class ORTTrainer(): Defaults to False run_symbolic_shape_infer: run symbolic shape inference Defaults to False + optimized_model_filepath: path to output the optimized training graph. + Defaults to "" (no output). """ warnings.warn('DISCLAIMER: This is an early version of an experimental training API and it is subject to change. DO NOT create production applications with it') self.is_train = True @@ -681,6 +687,7 @@ class ORTTrainer(): self.use_invertible_layernorm_grad = use_invertible_layernorm_grad self.run_symbolic_shape_infer = run_symbolic_shape_infer self.enable_adasum = enable_adasum + self.optimized_model_filepath = optimized_model_filepath # use this special string to workaround a corner case that external loss_scale is passed into train_step as kwargs. # see prepare_input_and_fetches for more details. @@ -711,7 +718,9 @@ class ORTTrainer(): enable_grad_norm_clip=self.enable_grad_norm_clip_, frozen_weights=self.frozen_weights_, opset_version=self.opset_version_, use_deterministic_compute=self._use_deterministic_compute, - use_invertible_layernorm_grad=self.use_invertible_layernorm_grad, enable_adasum=self.enable_adasum) + use_invertible_layernorm_grad=self.use_invertible_layernorm_grad, + enable_adasum=self.enable_adasum, + optimized_model_filepath=self.optimized_model_filepath) self.loss_scale_input_name = self.session.loss_scale_input_name diff --git a/orttraining/orttraining/python/training/orttrainer.py b/orttraining/orttraining/python/training/orttrainer.py index 6e321d965a..823977c7c4 100644 --- a/orttraining/orttraining/python/training/orttrainer.py +++ b/orttraining/orttraining/python/training/orttrainer.py @@ -663,6 +663,8 @@ class ORTTrainer(object): self.options.graph_transformer.gelu_recompute or self.options.graph_transformer.transformer_layer_recompute): session_options.execution_order = ort.ExecutionOrder.PRIORITY_BASED + if len(self.options.debug.graph_save_paths.model_with_training_graph_after_optimization_path) > 0: + session_options.optimized_model_filepath = self.options.debug.graph_save_paths.model_with_training_graph_after_optimization_path # old ort session may already exists and occupies GPU memory when creating new session, this may cause OOM error. # for example, load_state_dict will be called before returing the function, and it calls _init_session again diff --git a/orttraining/orttraining/python/training/orttrainer_options.py b/orttraining/orttraining/python/training/orttrainer_options.py index 04f7e62ca9..eb74879e0f 100644 --- a/orttraining/orttraining/python/training/orttrainer_options.py +++ b/orttraining/orttraining/python/training/orttrainer_options.py @@ -243,9 +243,13 @@ class ORTTrainerOptions(object): 'model_with_training_graph_path': { 'type': 'string', 'default': '' - } + }, + 'model_with_training_graph_after_optimization_path': { + 'type': 'string', + 'default': '' + }, } - } + }, } }, '_internal_use' : { @@ -365,6 +369,8 @@ class ORTTrainerOptions(object): 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. + debug.graph_save_paths.model_with_training_graph_after_optimization_path (str, default is "") + outputs the optimized training graph to the path if nonempty. _internal_use (dict): internal options, possibly undocumented, that might be removed without notice _internal_use.enable_internal_postprocess (bool, default is True): @@ -677,7 +683,7 @@ _ORTTRAINER_OPTIONS_SCHEMA = { }, 'graph_save_paths' : { 'type' : 'dict', - 'default_setter': lambda _: {}, + 'default_setter': lambda _: {}, 'required': False, 'schema': { 'model_after_graph_transforms_path': { @@ -691,7 +697,11 @@ _ORTTRAINER_OPTIONS_SCHEMA = { 'model_with_training_graph_path': { 'type': 'string', 'default': '' - } + }, + 'model_with_training_graph_after_optimization_path': { + 'type': 'string', + 'default': '' + }, } }, } diff --git a/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py b/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py index c0c8c8a89d..e13ccd6f6a 100644 --- a/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py +++ b/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py @@ -90,7 +90,8 @@ def testORTTrainerOptionsDefaultValues(test_input): 'graph_save_paths' : { 'model_after_graph_transforms_path': '', 'model_with_gradient_graph_path': '', - 'model_with_training_graph_path': '' + 'model_with_training_graph_path': '', + 'model_with_training_graph_after_optimization_path': '' } }, '_internal_use': { @@ -1431,13 +1432,14 @@ def testORTTrainerUnusedInput(): pytest.fail("RuntimeError doing train_step with an unused input.") @pytest.mark.parametrize("debug_files", [ - ({'model_after_graph_transforms_path': 'transformed.onnx', + {'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': 'training.onnx', + 'model_with_training_graph_after_optimization_path': 'training_optimized.onnx' + }, + {'model_after_graph_transforms_path': 'transformed.onnx', 'model_with_training_graph_path': '' - }), + }, ]) def testTrainingGraphExport(debug_files): device = 'cuda' @@ -1468,6 +1470,8 @@ def testTrainingGraphExport(debug_files): 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) + elif k == 'model_with_training_graph_after_optimization_path': + assert any("FusedMatMul" in n.op_type for n in saved_graph.node) # remove saved file os.remove(path) else: