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

Enable graph save for orttrainer (#6333)

Browse source 
* 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 <thiago.crepaldi@microsoft.com>
This commit is contained in:
ashbhandare 2021-01-14 10:07:54 -08:00 committed by GitHub
parent c24f2950bf
commit fd21c84eb8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 138 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
orttraining/orttraining/core/session/training_session.cc
View file

@ -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{};

4
orttraining/orttraining/core/session/training_session.h
View file

@ -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<PathString> model_with_loss_function_path{};
// The path at which to save the model after applying the graph transformations.
optional<PathString> model_after_graph_transforms_path{};
// The path at which to save the model with gradient graph added.
optional<PathString> model_with_gradient_graph_path{};
// The path at which to save the intermediate model with the whole training graph.
optional<PathString> model_with_training_graph_path{};

18
orttraining/orttraining/python/orttraining_pybind_state.cc
View file

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

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

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

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

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

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

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