From 7f5339505e146bfd303b629cd8f7d052ac72b461 Mon Sep 17 00:00:00 2001 From: suffiank Date: Mon, 8 Jun 2020 14:16:10 -0700 Subject: [PATCH] Discover trainable parameters using reverse DFS from loss node (#4116) Discover trainable parameters using reverse DFS from loss node, omitting recursion along untrainable inputs. Co-authored-by: suffian khan Co-authored-by: suffian khan --- include/onnxruntime/core/graph/graph.h | 14 ++++ onnxruntime/core/graph/graph.cc | 13 ++- onnxruntime/test/ir/graph_test.cc | 38 +++++++-- .../core/session/training_session.cc | 84 ++++++++++++------- .../core/session/training_session.h | 7 +- orttraining/orttraining/models/bert/main.cc | 5 -- 6 files changed, 118 insertions(+), 43 deletions(-) diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h index 717f5b21ce..048218744f 100644 --- a/include/onnxruntime/core/graph/graph.h +++ b/include/onnxruntime/core/graph/graph.h @@ -737,6 +737,20 @@ class Graph { const std::function& leave, const std::function& comp = {}) const; + /** Performs a reverse depth-first search (DFS) traversal from a set of nodes, via their inputs, + up to their source node/s. + @param from Set of Nodes to traverse from. + @param enter Visit function that will be invoked on a node when it is visited but its parents haven't been. + @param leave Visit function invoked on the node after its parents have all been visited. + @param stop Stop traversal from node n to input node p if stop(n, p) is true. + @param comp Comparison function to stabilize the traversal order by making Node ordering deterministic. + */ + void ReverseDFSFrom(const std::vector& from, + const std::function& enter, + const std::function& leave, + const std::function& comp, + const std::function& stop) const; + /** Gets the map of operator domains to their opset versions. */ const std::unordered_map& DomainToVersionMap() const noexcept { return domain_to_version_; diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index 89e59c0850..cf1a90c051 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -1283,13 +1283,22 @@ void Graph::ReverseDFSFrom(const std::vector& from, node_vec.push_back(GetNode(i)); } - ReverseDFSFrom(node_vec, enter, leave, comp); + ReverseDFSFrom(node_vec, enter, leave, comp, {}); } void Graph::ReverseDFSFrom(const std::vector& from, const std::function& enter, const std::function& leave, const std::function& comp) const { + + ReverseDFSFrom(from, enter, leave, comp, {}); +} + +void Graph::ReverseDFSFrom(const std::vector& from, + const std::function& enter, + const std::function& leave, + const std::function& comp, + const std::function& stop) const { using WorkEntry = std::pair; // bool represents leave or not std::vector stack(from.size()); for (size_t i = 0; i < from.size(); i++) { @@ -1323,6 +1332,7 @@ void Graph::ReverseDFSFrom(const std::vector& from, if (comp) { std::vector sorted_nodes; for (auto iter = n.InputNodesBegin(); iter != n.InputNodesEnd(); ++iter) { + if (stop && stop(&n, &(*iter))) continue; sorted_nodes.push_back(&(*iter)); } std::sort(sorted_nodes.begin(), sorted_nodes.end(), comp); @@ -1334,6 +1344,7 @@ void Graph::ReverseDFSFrom(const std::vector& from, } } else { for (auto iter = n.InputNodesBegin(); iter != n.InputNodesEnd(); ++iter) { + if (stop && stop(&n, &(*iter))) continue; const NodeIndex idx = (*iter).Index(); if (!visited[idx]) { stack.emplace_back(GetNode(idx), false); diff --git a/onnxruntime/test/ir/graph_test.cc b/onnxruntime/test/ir/graph_test.cc index efa6a894a1..f91a8e5e32 100644 --- a/onnxruntime/test/ir/graph_test.cc +++ b/onnxruntime/test/ir/graph_test.cc @@ -360,13 +360,14 @@ TEST_F(GraphTest, ReverseDFS) { auto& graph = model.MainGraph(); /* Case 1: A normal graph. + * * SouceNode * / \ - * node_1 (Variable) node_2 (Variable) - * \ / - * node_3 (Add) - * | - * node_4 (NoOp) + * node_1 (Variable) node_2 (Variable) node_5 (Variable) + * \ / | + * node_3 (Add) node_6 (NoOp) + * | | + * node_4 (Add) ------------------- <-- request stop * | * SinkNode */ @@ -399,12 +400,31 @@ TEST_F(GraphTest, ReverseDFS) { outputs.push_back(&output_arg3); auto& node_3 = graph.AddNode("node_3", "Add_DFS", "node 3", inputs, outputs); + // side path + inputs.clear(); + auto& input_arg5 = graph.GetOrCreateNodeArg("node_5_in_1", &tensor_int32); + inputs.push_back(&input_arg5); + auto& output_arg5 = graph.GetOrCreateNodeArg("node_5_out_1", &tensor_int32); + outputs.clear(); + outputs.push_back(&output_arg5); + graph.AddNode("node_5", "Variable_DFS", "node 5", inputs, outputs); + + inputs.clear(); + inputs.push_back(&output_arg5); + auto& output_arg6 = graph.GetOrCreateNodeArg("node_6_out_1", &tensor_int32); + outputs.clear(); + outputs.push_back(&output_arg6); + graph.AddNode("node_6", "NoOp_DFS", "node 6", inputs, outputs); + + // merged inputs.clear(); inputs.push_back(&output_arg3); + inputs.push_back(&output_arg6); auto& output_arg4 = graph.GetOrCreateNodeArg("node_4_out_1", &tensor_int32); outputs.clear(); outputs.push_back(&output_arg4); - graph.AddNode("node_4", "NoOp_DFS", "node 4", inputs, outputs); + graph.AddNode("node_4", "Add_DFS", "node 4", inputs, outputs); + auto status = graph.Resolve(); EXPECT_TRUE(status.IsOK()) << status.ErrorMessage(); @@ -440,7 +460,11 @@ TEST_F(GraphTest, ReverseDFS) { s += n->Name(); enter_leave_sequence.push_back(s); }, - NodeCompareName()); + NodeCompareName(), + // don't traverse side path + [](const Node* from, const Node* to) { + return from->Name() == "node_4" && to->Name() == "node_6"; + }); EXPECT_EQ(enter_leave_sequence.size(), 8u); EXPECT_EQ("enter:node_4", enter_leave_sequence.at(0)); diff --git a/orttraining/orttraining/core/session/training_session.cc b/orttraining/orttraining/core/session/training_session.cc index 67201f6a57..e200866088 100644 --- a/orttraining/orttraining/core/session/training_session.cc +++ b/orttraining/orttraining/core/session/training_session.cc @@ -130,17 +130,6 @@ Status TrainingSession::ConfigureForTraining( config.distributed_config.horizontal_parallel_size, config.distributed_config.pipeline_parallel_size}); - // 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. - std::unordered_set excluded_initializers = - !config.weight_names_to_train.empty() - ? config.weight_names_to_train - : GetTrainableModelInitializers(config.immutable_weights); - for (const auto& weight_name_to_not_train : config.weight_names_to_not_train) { - excluded_initializers.erase(weight_name_to_not_train); - } - if (config.pipeline_config.has_value() && config.pipeline_config.value().do_partition) { // Apply online pipeline partition to graph obj. This needs to be done first before any graph // transportation which may alter node_arg and invalidate cut_list info from the original graph. @@ -150,8 +139,6 @@ Status TrainingSession::ConfigureForTraining( config.distributed_config.world_size)); } - ORT_RETURN_IF_ERROR(ApplyTransformationsToMainGraph(excluded_initializers)); - is_mixed_precision_enabled_ = config.mixed_precision_config.has_value(); std::string loss_name{}; @@ -188,11 +175,30 @@ Status TrainingSession::ConfigureForTraining( 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. + std::unordered_set trainable_initializers = + !config.weight_names_to_train.empty() + ? config.weight_names_to_train + : GetTrainableModelInitializers(config.immutable_weights, loss_name); + if (config.weight_names_to_not_train.size() > 0) + { + LOGS(*session_logger_, INFO) << "Excluding following weights from trainable list as specified in configuration:\n"; + for (const auto& weight_name_to_not_train : config.weight_names_to_not_train) { + trainable_initializers.erase(weight_name_to_not_train); + LOGS(*session_logger_, INFO) << weight_name_to_not_train; + } + LOGS(*session_logger_, INFO) << std::endl; + } + + ORT_RETURN_IF_ERROR(ApplyTransformationsToMainGraph(trainable_initializers)); + // derive actual set of weights to train std::unordered_set weight_names_to_train = !config.weight_names_to_train.empty() ? config.weight_names_to_train - : GetTrainableModelInitializers(config.immutable_weights); + : GetTrainableModelInitializers(config.immutable_weights, loss_name); for (const auto& weight_name_to_not_train : config.weight_names_to_not_train) { weight_names_to_train.erase(weight_name_to_not_train); } @@ -937,26 +943,46 @@ bool TrainingSession::IsImmutableWeight(const ImmutableWeights& immutable_weight } std::unordered_set TrainingSession::GetTrainableModelInitializers( - const ImmutableWeights& immutable_weights) const { + const ImmutableWeights& immutable_weights, const std::string& loss_name) const { + const Graph& graph = model_->MainGraph(); const auto& initialized_tensors = graph.GetAllInitializedTensors(); - std::unordered_set model_initializers; - std::transform(initialized_tensors.begin(), - initialized_tensors.end(), - std::inserter(model_initializers, model_initializers.end()), - [](const auto& pair) { return pair.first; }); + std::unordered_set trainable_initializers; - std::unordered_set trainable_initializers(model_initializers); - for (const std::string& initializer_name : model_initializers) { - const auto& nodes = graph.GetConsumerNodes(initializer_name); - for (const Node* node : nodes) { - if (IsUntrainable(node, initializer_name, session_logger_) || - IsImmutableWeight(immutable_weights, node, initialized_tensors.at(initializer_name), session_logger_)) { - trainable_initializers.erase(initializer_name); + auto add_trainable_initializers = [&](const Node* node) { + for (auto input : node->InputDefs()) { + std::string initializer_name = input->Name(); + if (initialized_tensors.count(initializer_name) == 0) + continue; + + if (IsUntrainable(node, initializer_name, session_logger_) || + IsImmutableWeight(immutable_weights, node, initialized_tensors.at(initializer_name), session_logger_)) + continue; + + trainable_initializers.insert(initializer_name); } - } - } + }; + auto stop_at_untrainable = [&](const Node* from, const Node* to) { + + auto is_trainable_from_to_link = [&](Node::EdgeEnd e) { + if (&e.GetNode() != to) + return false; + + std::string input_name = from->InputDefs()[e.GetDstArgIndex()]->Name(); + return !IsUntrainable(from, input_name, session_logger_); + }; + + bool proceed = std::any_of(from->InputEdgesBegin(), from->InputEdgesEnd(), is_trainable_from_to_link); + if (!proceed && session_logger_) { + VLOGS(*session_logger_, 1) << "Stopping training parameters discovery traversal from " << from->Name() << " to " << to->Name() << std::endl; + } + + return !proceed; + }; + + // perform reverse dfs from output node to discover trainable parameters + graph.ReverseDFSFrom({graph.GetProducerNode(loss_name)}, add_trainable_initializers, {}, {}, stop_at_untrainable); return trainable_initializers; } diff --git a/orttraining/orttraining/core/session/training_session.h b/orttraining/orttraining/core/session/training_session.h index 4dfdd5f3d8..b5c03c51a1 100644 --- a/orttraining/orttraining/core/session/training_session.h +++ b/orttraining/orttraining/core/session/training_session.h @@ -426,7 +426,12 @@ class TrainingSession : public InferenceSession { bool use_fp16_initializer, std::unordered_map& fp32_weight_name_to_fp16_node_arg); - std::unordered_set GetTrainableModelInitializers(const ImmutableWeights& immutable_weights) const; + /** Discover all trainable initializers by reverse DFS starting from a given tensor (for example, the loss value) + @param immutable_weights do not include initializers matching an (op_type, input_index, value) entry from this table + @param backprop_source_name reverse DFS back propagation source name (i.e. loss name or pipeline send output name) + */ + std::unordered_set GetTrainableModelInitializers(const ImmutableWeights& immutable_weights, + const std::string& backprop_source_name) const; std::unordered_set GetStateTensorNames() const; diff --git a/orttraining/orttraining/models/bert/main.cc b/orttraining/orttraining/models/bert/main.cc index 1a98f81e32..edf2bcdf8b 100644 --- a/orttraining/orttraining/models/bert/main.cc +++ b/orttraining/orttraining/models/bert/main.cc @@ -551,11 +551,6 @@ void setup_training_params(BertParameters& params) { /*mlm_loss*/ "mlm_loss", /*nsp_loss*/ "nsp_loss"}); - params.weights_not_to_train = { - "position_01", // Slice's dat input - "op_min_ends_expand_10", //op_min_ends_expand_10 - "72", // [BERT-tiny only] input of expand - }; params.fetch_names = {"total_loss", "mlm_loss", "nsp_loss"}; if (params.EnableTensorboard()) {