From 6404aba5ae962604477e27833069279e1e0ccde7 Mon Sep 17 00:00:00 2001 From: Tixxx Date: Fri, 29 May 2020 12:28:12 -0700 Subject: [PATCH] Orttraining rc1 master merge (#4080) * fixed seg fault when using concrete shape disable gradient as output * fix evaluation hang issue for multiple gpu run * Remove dead code, ORTModel and improve docstrings (#3814) * Refine ORTTrainer docstring descriptions (#3907) --- onnxruntime/core/graph/graph_utils.cc | 2 +- orttraining/orttraining/models/bert/main.cc | 9 +- orttraining/orttraining/python/ort_trainer.py | 94 +++++++++++++------ 3 files changed, 71 insertions(+), 34 deletions(-) diff --git a/onnxruntime/core/graph/graph_utils.cc b/onnxruntime/core/graph/graph_utils.cc index 7ee51e8c10..c63d003e24 100644 --- a/onnxruntime/core/graph/graph_utils.cc +++ b/onnxruntime/core/graph/graph_utils.cc @@ -420,7 +420,7 @@ bool CanReplaceNodeWithInitializer(const Graph& graph, const Node& node, const s const logging::Logger& logger) { // we have no way to handle replacing multiple outputs so check only one is used const std::string* output_name = nullptr; - if (!IsOnlyOneOutputUsed(graph, node, output_name)) { + if (!IsOnlyOneOutputUsed(graph, node, output_name) || output_name == nullptr) { return false; } diff --git a/orttraining/orttraining/models/bert/main.cc b/orttraining/orttraining/models/bert/main.cc index cd719efee9..1a98f81e32 100644 --- a/orttraining/orttraining/models/bert/main.cc +++ b/orttraining/orttraining/models/bert/main.cc @@ -743,10 +743,11 @@ static Status RunTraining(const BertParameters& params, const Environment& env) ORT_RETURN_IF_ERROR(runner->ResetLossScaler()); } - auto test_data_loader = onnxruntime::make_unique(params_for_phase.input_name_map, - params_for_phase.test_data_dir, - max_num_files_preload); - ORT_RETURN_IF_ERROR(runner->EndTraining(test_data_loader.get())); + if (params_for_phase.mpi_context.world_rank == 0) { + // Pass in empty dataloader to disable evaluation in EndTraining + // to avoid a redundant synchronization caused by Tensorboard's SummaryMerge Op. + ORT_RETURN_IF_ERROR(runner->EndTraining(nullptr)); + } return Status::OK(); } diff --git a/orttraining/orttraining/python/ort_trainer.py b/orttraining/orttraining/python/ort_trainer.py index 877a95ef1a..4d2d614b47 100644 --- a/orttraining/orttraining/python/ort_trainer.py +++ b/orttraining/orttraining/python/ort_trainer.py @@ -8,8 +8,10 @@ from onnx import helper import torch import torch.nn import torch.onnx -import onnxruntime as ort from distutils.version import LooseVersion +import warnings + +import onnxruntime as ort from .checkpointing_utils import list_checkpoint_files, get_checkpoint_name, CombineZeroCheckpoint import onnxruntime.capi.pt_patch @@ -22,13 +24,11 @@ class IODescription(): self.dtype_ = dtype self.num_classes_ = num_classes - class ModelDescription(): def __init__(self, inputs, outputs): self.inputs_ = inputs self.outputs_ = outputs - def resolve_symbolic_dimensions(inputs, input_descs, output_descs): import copy output_descs_copy = copy.deepcopy(output_descs) @@ -511,34 +511,70 @@ class ORTTrainer(): enable_grad_norm_clip=True, frozen_weights=[], _opset_version=DEFAULT_OPSET_VERSION): super(ORTTrainer, self).__init__() """ - Initializes ORTTrainer. + Initialize ORTTrainer. - Params: - model: either: - - a PyTorch model: if 'loss_fn' is not provided, the 'model's first output must be the loss. - if 'loss_fn' is provided, 'model' and 'loss_fn' are combined as described in 'loss_fn:' - - an ONNX model: the 'model's first output must be the loss. - loss_fn: a PyTorch loss function. It takes two inputs [prediction, label] and ouput a loss tensor. - If provided, 'loss_fn' is combined with the PyTorch 'model' to form a Combined PyTorch model. - Inputs to the Combined PyTorch model are concatination of the 'model's input and 'loss_fn's label input. - Outputs of the Combined PyTorch model are concatination of 'loss_fn's loss output and 'model's outputs. - model_desc: model input and output description. it is used to specify input/output shapes, types, and names. - 'model_desc' must be consistent with the training model. - training_optimizer_name: one of: 'SGDOptimizer', 'AdamOptimizer', 'LambOptimizer'. - map_optimizer_attributes: for optimizers with weight dependent parameters, - 'map_optimizer_attributes' maps weight name to a set of optimization parameters. - learning_rate_description: in form of IODescription(Learning_Rate_Name, [1,], torch.float32). - Because learning_rate is an input to the training model, Learning_Rate_Name shall be set so that - there is no name conflict within the above model. - device: cuda device to store tensors - gradient_accumulation_steps: number of training steps to accumulate gradients before averaging and applying them (default is 1) - postprocess_model: a callable to postprocess the ONNX model that is converted from PyTorch (default is None) - world_rank: rank id used for distributed training (default is 0) - world_size: number of ranks participating in distributed training (default is 1) - use_mixed_precision: flag to enable mixed precision (aka fp16) (default is False) - allreduce_post_accumulation: controls whether overlaping gradient computation with allreduce (default is False) - partition_optimizer: controls whether to partition the optimizer state (default is False) + Args: + + model: one of + - a PyTorch model (class that inherits from torch.nn.Module) + - a combined PyTorch model and loss function. + Inputs to this combined PyTorch model are a concatenation of the + model's input and the loss function's label input. + Outputs are a concatenation of the loss function's output and the + model's output. + - a combined ONNX model and loss function. + loss_fn: one of + - a PyTorch loss function if 'model' is a PyTorch model. A loss + function takes two inputs (prediction, label) and outputs a loss + tensor. + - None if model is already combined with a loss function. + model_desc: Specify input/output shapes, types, and names. + Must be consistent with the training model. + training_optimizer_name: one of + - 'SGDOptimizer' + - 'AdamOptimizer' + - 'LambOptimizer' + map_optimizer_attributes: for optimizers with weight-dependent + parameters. A callable that maps weight name to a set of optimization + parameters. + Defaults to None. + learning_rate_description: the name, shape and type of the learning + rate in form of IODescription(Learning_Rate_Name, [1,], torch.float32). + Because learning_rate is an input to the training model, + Learning_Rate_Name must be specified so that there is no name conflict + within the model. + device: device to store tensors (e.g. 'cpu', 'cuda', 'cuda:'). + gradient_accumulation_steps: number of training steps to accumulate + gradients before averaging and applying them. + Defaults to 1. + postprocess_model: a callable to postprocess the ONNX model that is + converted from PyTorch. + Defaults to None. + world_rank: rank id used for distributed training. + Defaults to 0. + world_size: number of ranks participating in distributed training. + Defaults to 1. + use_mixed_precision: flag to enable mixed precision (aka fp16). + Defaults to False. + allreduce_post_accumulation: controls whether overlaping gradient + computation is applied with allreduce. + Defaults to False. + global_step: training step that is used as input to 'get_lr_this_step'. + Defaults to 0. + get_lr_this_step: functor used as learning rate scheduler. + It uses 'global_step' as input. + Defaults to None. + loss_scaler: updates loss scale automatically when 'use_mixed_precision' + is specified. + Defaults to None. + partition_optimizer: controls whether to partition the optimizer state. + Defaults to False. + enable_grad_norm_clip: enables gradient norm clipping. + Defaults to True. + frozen_weights: list of model parameters to be frozen (not trained). + Defaults to []. """ + 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 self.torch_model_ = None