saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add an option to save the training graph after optimization (#6410)

Browse source 
* expose optimized_model_filepath in SessionOptions as `debug.graph_save_paths.model_with_training_graph_after_optimization_path` in `ORTTrainerOptions`
This commit is contained in:
Ryota Tomioka 2021-01-27 07:39:46 +00:00 committed by GitHub
parent 0d20104a72
commit 9835b46a1d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 38 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
orttraining/orttraining/python/ort_trainer.py
View file

@ -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

2
orttraining/orttraining/python/training/orttrainer.py
View file

@ -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

18
orttraining/orttraining/python/training/orttrainer_options.py
View file

@ -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': ''
},
}
},
}

16
orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py
View file

@ -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: