Add option to dump activations on all ranks (#5455)

* Add option to dump activations on all ranks

* address review comments

* review comments

* Fix review comment

Co-authored-by: Aishwarya <aibhanda@OrtTrainingDev4.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
This commit is contained in:
ashbhandare 2020-11-02 18:03:05 -08:00 committed by GitHub
parent 87e1063e19
commit c875fe0919
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 15 additions and 3 deletions

View file

@ -90,14 +90,14 @@ void DumpTensorToStdOut(const Tensor& tensor) {
std::cout << std::endl;
}
PathString MakeTensorFileName(const std::string& tensor_name) {
PathString MakeTensorFileName(const std::string& tensor_name, const NodeDumpOptions& dump_options) {
auto make_valid_name = [](std::string name) {
std::replace_if(
name.begin(), name.end(), [](char c) { return !std::isalnum(c); }, '_');
return name;
};
return path_utils::MakePathString(make_valid_name(tensor_name), ".tensorproto");
return path_utils::MakePathString(make_valid_name(tensor_name), dump_options.file_suffix, ".tensorproto");
}
void DumpTensorToFile(const Tensor& tensor, const std::string& tensor_name, const Path& file_path) {
@ -125,7 +125,7 @@ void DumpCpuTensor(
break;
}
case NodeDumpOptions::DataDestination::TensorProtoFiles: {
const Path tensor_file = dump_options.output_dir / Path::Parse(MakeTensorFileName(tensor_name));
const Path tensor_file = dump_options.output_dir / Path::Parse(MakeTensorFileName(tensor_name, dump_options));
DumpTensorToFile(tensor, tensor_name, tensor_file);
break;
}
@ -203,6 +203,15 @@ const NodeDumpOptions& NodeDumpOptionsFromEnvironmentVariables() {
opts.data_destination = NodeDumpOptions::DataDestination::TensorProtoFiles;
}
if (get_bool_env_var(env_vars::kAppendRankToFileName)) {
std::string rank = Env::Default().GetEnvironmentVar("OMPI_COMM_WORLD_RANK");
if (rank.empty()) {
opts.file_suffix = "_default_rank_0";
} else {
opts.file_suffix = "_rank_" + rank;
}
}
opts.output_dir = Path::Parse(ToPathString(Env::Default().GetEnvironmentVar(env_vars::kOutputDir)));
// check for confirmation for dumping data to files for all nodes

View file

@ -30,6 +30,8 @@ constexpr const char* kNameFilter = "ORT_DEBUG_NODE_IO_NAME_FILTER";
constexpr const char* kOpTypeFilter = "ORT_DEBUG_NODE_IO_OP_TYPE_FILTER";
// set to non-zero to dump data to files instead of stdout
constexpr const char* kDumpDataToFiles = "ORT_DEBUG_NODE_IO_DUMP_DATA_TO_FILES";
// set to non-zero to append OpenMPI world rank to filename
constexpr const char* kAppendRankToFileName = "ORT_DEBUG_NODE_IO_APPEND_RANK_TO_FILE_NAME";
// specify the output directory for any data files produced
constexpr const char* kOutputDir = "ORT_DEBUG_NODE_IO_OUTPUT_DIR";
// set to non-zero to confirm that dumping data files for all nodes is acceptable
@ -75,6 +77,7 @@ struct NodeDumpOptions {
TensorProtoFiles,
} data_destination{DataDestination::StdOut};
std::string file_suffix;
// the output directory for dumped data files
Path output_dir;
};