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

Discover trainable parameters using reverse DFS from loss node (#4116)

Browse source 
Discover trainable parameters using reverse DFS from loss node, omitting recursion along untrainable inputs.

Co-authored-by: suffian khan <sukha@OrtTrainingDev4.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
Co-authored-by: suffian khan <sukha@microsoft.com>
This commit is contained in:
suffiank 2020-06-08 14:16:10 -07:00 committed by GitHub
parent 842be1535d
commit 7f5339505e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 118 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
include/onnxruntime/core/graph/graph.h
View file

@ -737,6 +737,20 @@ class Graph {
const std::function<void(const Node*)>& leave,
const std::function<bool(const Node*, const Node*)>& 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<const Node*>& from,
const std::function<void(const Node*)>& enter,
const std::function<void(const Node*)>& leave,
const std::function<bool(const Node*, const Node*)>& comp,
const std::function<bool(const Node*, const Node*)>& stop) const;
/** Gets the map of operator domains to their opset versions. */
const std::unordered_map<std::string, int>& DomainToVersionMap() const noexcept {
return domain_to_version_;

13
onnxruntime/core/graph/graph.cc
View file

@ -1283,13 +1283,22 @@ void Graph::ReverseDFSFrom(const std::vector<NodeIndex>& 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<const Node*>& from,
const std::function<void(const Node*)>& enter,
const std::function<void(const Node*)>& leave,
const std::function<bool(const Node*, const Node*)>& comp) const {
ReverseDFSFrom(from, enter, leave, comp, {});
}
void Graph::ReverseDFSFrom(const std::vector<const Node*>& from,
const std::function<void(const Node*)>& enter,
const std::function<void(const Node*)>& leave,
const std::function<bool(const Node*, const Node*)>& comp,
const std::function<bool(const Node* from, const Node* to)>& stop) const {
using WorkEntry = std::pair<const Node*, bool>; // bool represents leave or not
std::vector<WorkEntry> stack(from.size());
for (size_t i = 0; i < from.size(); i++) {
@ -1323,6 +1332,7 @@ void Graph::ReverseDFSFrom(const std::vector<const Node*>& from,
if (comp) {
std::vector<const Node*> 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<const Node*>& 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);

38
onnxruntime/test/ir/graph_test.cc
View file

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

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

@ -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<std::string> 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<std::string> 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<std::string> 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<std::string> 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<std::string> 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<std::string> trainable_initializers;
std::unordered_set<std::string> 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;
}

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

@ -426,7 +426,12 @@ class TrainingSession : public InferenceSession {
bool use_fp16_initializer,
std::unordered_map<std::string, NodeArg*>& fp32_weight_name_to_fp16_node_arg);
std::unordered_set<std::string> 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<std::string> GetTrainableModelInitializers(const ImmutableWeights& immutable_weights,
const std::string& backprop_source_name) const;
std::unordered_set<std::string> GetStateTensorNames() const;

5
orttraining/orttraining/models/bert/main.cc
View file

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