mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
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:
parent
87e1063e19
commit
c875fe0919
2 changed files with 15 additions and 3 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
};
|
||||
|
|
|
|||
Loading…
Reference in a new issue