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

Make memory profiler work with multiple session runs. (#12317)

Browse source 
* make memory profiler work with multiple session runs.

(cherry picked from commit 5b636b4dd6fe91b75c063696dc73eda33ec36c8d)

* minor fix

* fix build

* fix window build

* 1. fix cpplint issues;
2. give unique filesname for each session profiler result.
This commit is contained in:
pengwa 2022-07-29 18:36:31 +08:00 committed by GitHub
parent e4bd41fb3b
commit 6514069749
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 283 additions and 212 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
onnxruntime/core/framework/execution_frame.cc
View file

@ -357,7 +357,7 @@ ExecutionFrame::ExecutionFrame(gsl::span<const int> feed_mlvalue_idxs, gsl::span
fetches);
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
MemoryInfo::IncreaseIteration();
session_state.GetMemoryProfiler()->GetMemoryInfo().IncreaseIteration();
#endif
// map the custom allocators to ort_value_idx entries
@ -385,7 +385,7 @@ ExecutionFrame::ExecutionFrame(gsl::span<const int> feed_mlvalue_idxs, gsl::span
}
}
//if there are some traditional ml value type in inputs disable the memory pattern optimization.
// if there are some traditional ml value type in inputs disable the memory pattern optimization.
if (all_tensors) {
mem_patterns_ = session_state.GetMemoryPatternGroup(feeds, feed_mlvalue_idxs, inferred_shapes_);
// if no existing patterns, generate one in this execution frame
@ -439,16 +439,19 @@ ExecutionFrame::ExecutionFrame(gsl::span<const int> feed_mlvalue_idxs, gsl::span
buffers_[location] = BufferUniquePtr(buffer, alloc);
}
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
//Record activation memory pattern
MemoryInfo::ClearMemoryInfoPerExecution();
// Record activation memory pattern
auto mem_profier_ptr = session_state.GetMemoryProfiler();
mem_profier_ptr->GetMemoryInfo().ClearMemoryInfoPerExecution();
if (mem_patterns_ && buffer != nullptr) {
MemoryInfo::RecordPatternInfo(*mem_patterns_, MemoryInfo::MapType::StaticActivation);
MemoryInfo::MemoryInfoProfile::CreateEvents("static activations_" + std::to_string(MemoryInfo::GetIteration()),
MemoryInfo::MemoryInfoProfile::GetAndIncreasePid(), MemoryInfo::MapType::StaticActivation, "", 0);
mem_profier_ptr->GetMemoryInfo().RecordPatternInfo(*mem_patterns_, MemoryInfo::MapType::StaticActivation);
mem_profier_ptr->CreateEvents(
"static activations_" + std::to_string(mem_profier_ptr->GetMemoryInfo().GetIteration()),
mem_profier_ptr->GetAndIncreasePid(), MemoryInfo::MapType::StaticActivation, "", 0);
}
#endif
// log size of activation. Keep it commented out for now to avoid log flooding.
// VLOGS(session_state_.Logger(), 1) << "**** Allocated memory for activations, size: " <<mem_patterns_->patterns[i].PeakSize();
// VLOGS(session_state_.Logger(), 1) << "**** Allocated memory for activations, size: "
// << mem_patterns_->patterns[i].PeakSize();
}
}
}
@ -542,7 +545,7 @@ Status ExecutionFrame::AllocateMLValueTensorSelfOwnBufferHelper(OrtValue& ort_va
}
}
//no memory pattern, or the pattern is not correct.
// no memory pattern, or the pattern is not correct.
if (!alloc) alloc = GetAllocator(location);
Tensor::InitOrtValue(element_type, shape, std::move(alloc), ort_value);
@ -560,7 +563,7 @@ Status ExecutionFrame::AllocateMLValueTensorSelfOwnBufferHelper(OrtValue& ort_va
// if parallel executor is used.
std::unique_lock<std::mutex> lock(mtx_);
dynamic_activation_memory_sizes_in_byte_[location.name] += size;
MemoryInfo::SetDynamicAllocation(ort_value_index);
session_state_.GetMemoryProfiler()->GetMemoryInfo().SetDynamicAllocation(ort_value_index);
#endif
}
@ -753,7 +756,7 @@ Status ExecutionFrame::AllocateAsPerAllocationPlan(OrtValue& ort_value, int ort_
}
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
MemoryInfo::RecordActivationAllocInfo(ort_value_index, ort_value);
session_state_.GetMemoryProfiler()->GetMemoryInfo().RecordActivationAllocInfo(ort_value_index, ort_value);
#endif
return Status::OK();

110
onnxruntime/core/framework/memory_info.cc
View file

@ -1,34 +1,30 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
#include "core/framework/memory_info.h"
#include "core/framework/mem_pattern.h"
#include "core/framework/ort_value.h"
#include <fstream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
namespace onnxruntime {
size_t MemoryInfo::iteration_ = 0;
size_t MemoryInfo::num_node_size_ = 0;
int MemoryInfo::local_rank_ = 0;
std::map<OrtMemoryInfo, std::map<MemoryInfo::MapType, MemoryInfoMap> > MemoryInfo::tensors_memory_info_map_;
std::map<OrtValueIndex, MemoryInfo::AllocInfoPerTensor> MemoryInfo::tensor_alloc_info_map_;
std::map<MemoryInfo::MapType, std::set<size_t> > MemoryInfo::time_step_trace_;
std::map<const std::string, std::map<const std::string, bool> > MemoryInfo::customized_recording_group_;
// Record allocation information for each tensor, based on the execution plan
void MemoryInfo::GenerateTensorMap(const SequentialExecutionPlan* execution_plan,
const OrtValueNameIdxMap& value_name_idx_map) {
if (!tensor_alloc_info_map_.empty()) {
return;
}
void MemoryInfo::Init(const SequentialExecutionPlan* execution_plan,
const OrtValueNameIdxMap& value_name_idx_map) {
num_node_size_ = execution_plan->execution_plan.size();
for (OrtValueIndex value_idx = 0; value_idx < OrtValueIndex(execution_plan->allocation_plan.size()); ++value_idx) {
// Only store tensor information
if (!(execution_plan->allocation_plan[value_idx].value_type) ||
!(execution_plan->allocation_plan[value_idx].value_type->IsTensorType()))
continue;
AllocInfoPerTensor mem_info;
mem_info.mlvalue_index = value_idx;
ORT_THROW_IF_ERROR(value_name_idx_map.GetName(mem_info.mlvalue_index, mem_info.mlvalue_name));
@ -36,6 +32,7 @@ void MemoryInfo::GenerateTensorMap(const SequentialExecutionPlan* execution_plan
mem_info.reused_buffer = (execution_plan->allocation_plan[value_idx].alloc_kind != AllocKind::kReuse)
? value_idx
: execution_plan->allocation_plan[value_idx].reused_buffer;
// If the tensor is using memory outside of the scope, do not store it
if (execution_plan->allocation_plan[mem_info.reused_buffer].alloc_kind == AllocKind::kPreExisting) continue;
if (execution_plan->allocation_plan[mem_info.reused_buffer].alloc_kind == AllocKind::kAllocateOutput) continue;
@ -167,7 +164,7 @@ void MemoryInfo::PrintMemoryInfoForLocation(const OrtDevice::DeviceType location
}
}
std::string MemoryInfo::MemoryInfoProfile::CreateMetadataEvent(const std::string& process_name, size_t process_id) {
std::string MemoryProfiler::CreateMetadataEvent(const std::string& process_name, size_t process_id) {
std::stringstream evt;
evt << "{";
evt << "\"ph\":\"M\",";
@ -185,9 +182,9 @@ std::string MemoryInfo::MemoryInfoProfile::CreateMetadataEvent(const std::string
return evt.str();
}
std::string MemoryInfo::MemoryInfoProfile::CreateMemoryEvent(size_t pid, size_t tid, const std::string& name,
size_t offset, size_t size,
const std::string& color_name) {
std::string MemoryProfiler::CreateMemoryEvent(size_t pid, size_t tid, const std::string& name,
size_t offset, size_t size,
const std::string& color_name) {
std::stringstream evt;
evt << "{";
evt << "\"ph\":\"X\",";
@ -206,8 +203,8 @@ std::string MemoryInfo::MemoryInfoProfile::CreateMemoryEvent(size_t pid, size_t
return evt.str();
}
std::string MemoryInfo::MemoryInfoProfile::CreateSummaryEvent(size_t pid, size_t tid, const AllocationSummary& summary,
size_t size, size_t bytes_for_pattern) {
std::string MemoryProfiler::CreateSummaryEvent(size_t pid, size_t tid, const MemoryInfo::AllocationSummary& summary,
size_t size, size_t bytes_for_pattern) {
const size_t total_bytes = summary.total_size;
const size_t used_bytes = summary.used_size;
const size_t free_bytes = total_bytes - used_bytes;
@ -251,36 +248,6 @@ static void UpdateSummary(MemoryInfo::AllocationSummary& summary, size_t alloc_o
summary.live_tensors.push_back(idx);
}
// The following colors are defined and accepted by Chrome Tracing/Edge Tracing.
const std::vector<std::string> MemoryInfo::MemoryInfoProfile::color_names = {
"good",
"bad",
"terrible",
"yellow",
"olive",
"generic_work",
"background_memory_dump",
"light_memory_dump",
"detailed_memory_dump",
"thread_state_uninterruptible",
"thread_state_iowait",
"thread_state_running",
"thread_state_runnable",
"thread_state_unknown",
"cq_build_running",
"cq_build_passed",
"cq_build_failed",
"cq_build_abandoned",
"cq_build_attempt_runnig",
"cq_build_attempt_passed",
"cq_build_attempt_failed",
};
size_t MemoryInfo::MemoryInfoProfile::pid_ = 0;
std::vector<std::string> MemoryInfo::MemoryInfoProfile::events;
std::unordered_map<size_t, std::unordered_map<size_t, MemoryInfo::AllocationSummary> >
MemoryInfo::MemoryInfoProfile::summary_;
// Create sessions in the profiler.
// p_name: session name
// pid: sessionid
@ -290,12 +257,12 @@ std::unordered_map<size_t, std::unordered_map<size_t, MemoryInfo::AllocationSumm
// Used for generating customized sessions for a group of tensors
// Top_k: The steps with the top-k highest memory consumptions are plot. When top_k == 0, we plot all the steps
// device_type: The type of the device where the tensors are.
void MemoryInfo::MemoryInfoProfile::CreateEvents(const std::string& p_name,
const size_t pid,
const MemoryInfo::MapType& map_type,
const std::string& group_name,
const size_t top_k,
const OrtDevice::DeviceType device_type) {
void MemoryProfiler::CreateEvents(const std::string& p_name,
const size_t pid,
const MemoryInfo::MapType& map_type,
const std::string& group_name,
const size_t top_k,
const OrtDevice::DeviceType device_type) {
// Metadata.
std::string pid_name_internal = "device_" + std::to_string(device_type) + "_" + p_name + group_name;
events.push_back(CreateMetadataEvent(pid_name_internal, pid));
@ -303,7 +270,8 @@ void MemoryInfo::MemoryInfoProfile::CreateEvents(const std::string& p_name,
std::hash<std::string> str_hash;
// Create Event for each tensor
for (const auto& location_map : tensors_memory_info_map_) {
auto& time_step_trace = GetMemoryInfo().time_step_trace_;
for (const auto& location_map : GetMemoryInfo().tensors_memory_info_map_) {
const OrtMemoryInfo& memory_info = location_map.first;
const auto& maptype_to_map_mapping = location_map.second;
@ -314,10 +282,10 @@ void MemoryInfo::MemoryInfoProfile::CreateEvents(const std::string& p_name,
auto summary_key = str_hash(memory_info.ToString() + std::to_string(map_type));
// Preprocessing
if (time_step_trace_[map_type].empty()) {
if (time_step_trace[map_type].empty()) {
for (const auto& item : maptype_to_map_mapping.at(map_type)) {
time_step_trace_[map_type].insert(AllocPlan(item.first)->lifetime_interval.first);
time_step_trace_[map_type].insert(AllocPlan(item.first)->lifetime_interval.second);
time_step_trace[map_type].insert(GetMemoryInfo().AllocPlan(item.first)->lifetime_interval.first);
time_step_trace[map_type].insert(GetMemoryInfo().AllocPlan(item.first)->lifetime_interval.second);
}
}
@ -326,14 +294,14 @@ void MemoryInfo::MemoryInfoProfile::CreateEvents(const std::string& p_name,
if (summary_.find(summary_key) == summary_.end()) {
for (const auto& item : map) {
const OrtValueIndex idx = item.first;
const auto info = AllocPlan(idx);
const auto info = GetMemoryInfo().AllocPlan(idx);
if (info->inplace_reuse) continue;
size_t offset = IsStaticType(map_type) ? map.GetPlannedAddress(idx) : map.GetAllocAddress(idx);
size_t size = IsStaticType(map_type) ? map.GetPlannedSize(idx) : map.GetAllocSize(idx);
size_t alloc_step = AllocPlan(idx)->lifetime_interval.first;
size_t dealloc_step = AllocPlan(idx)->lifetime_interval.second;
const auto& ts_map = time_step_trace_[map_type];
size_t alloc_step = GetMemoryInfo().AllocPlan(idx)->lifetime_interval.first;
size_t dealloc_step = GetMemoryInfo().AllocPlan(idx)->lifetime_interval.second;
const auto& ts_map = GetMemoryInfo().time_step_trace_[map_type];
const auto& start_itr = ts_map.find(alloc_step);
const auto& end_itr = ts_map.find(dealloc_step);
ORT_ENFORCE(start_itr != ts_map.end(),
@ -368,12 +336,12 @@ void MemoryInfo::MemoryInfoProfile::CreateEvents(const std::string& p_name,
if (top_k != 0 && item.second.total_size < top_kth_total_size) continue;
size_t alloc_size_for_pattern = 0;
for (const auto& live_tensor : item.second.live_tensors) {
const auto info = AllocPlan(live_tensor);
const auto info = GetMemoryInfo().AllocPlan(live_tensor);
if (info->inplace_reuse) continue;
const std::string& name = info->mlvalue_name;
// Filter out string without a certain name
if (!group_name.empty()) {
if (!InRecordingTensorGroup(group_name, name)) continue;
if (!GetMemoryInfo().InRecordingTensorGroup(group_name, name)) continue;
}
const std::string cname = color_names[str_hash(name) % color_names.size()];
// Sometimes a tensor can be both statically planned and dynamically allocated, so we need to use planned
@ -394,13 +362,15 @@ void MemoryInfo::MemoryInfoProfile::CreateEvents(const std::string& p_name,
}
}
void MemoryInfo::GenerateMemoryProfile() {
void MemoryProfiler::GenerateMemoryProfile() {
// Write memory profile .json
std::ofstream memory_profile("memory_profile_" + std::to_string(local_rank_) + ".json", std::ios::trunc);
std::stringstream ss;
ss << "memory_profile_" << GetMemoryInfo().GetLocalRank() << "_" << profiler_id_ << ".json";
std::ofstream memory_profile(ss.str(), std::ios::trunc);
memory_profile << "[" << std::endl;
for (size_t i = 0; i < MemoryInfoProfile::GetEvents().size(); i++) {
memory_profile << " " << MemoryInfoProfile::GetEvents().at(i);
if (i < MemoryInfoProfile::GetEvents().size() - 1) memory_profile << ",";
for (size_t i = 0; i < GetEvents().size(); i++) {
memory_profile << " " << GetEvents().at(i);
if (i < GetEvents().size() - 1) memory_profile << ",";
memory_profile << std::endl;
}
memory_profile << "]" << std::endl;

190
onnxruntime/core/framework/memory_info.h
View file

@ -1,3 +1,6 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
#pragma once
#include "core/graph/basic_types.h"
@ -9,6 +12,9 @@
#include "core/framework/ort_value_name_idx_map.h"
#include <iomanip>
#include <string>
#include <unordered_map>
#include <vector>
namespace onnxruntime {
using OrtValueIndex = int;
@ -90,11 +96,15 @@ struct MemoryInfoMap {
class MemoryInfo {
public:
friend struct MemoryProfiler;
MemoryInfo() = default;
// We separate the memory types into 3 categories.
//(1) Initializers: Which are planned and allocated before session run.
//(2) StaticActivation: Which are planned statically and allocatedly accordingly
//(3) DynamicActivation: Which are allocated during runtime.
// The reason with this separations is they start with different memory offsets, so diffcult to visualize them in the same session.
// The reason with this separations is they start with different memory offsets, so diffcult to visualize them in the
// same session.
enum MapType {
Initializer = 0,
StaticActivation,
@ -120,62 +130,23 @@ class MemoryInfo {
std::vector<OrtValueIndex> live_tensors;
};
struct MemoryInfoProfile {
public:
// Create sessions in the profiler.
// p_name: session name
// pid: sessionid
// map_type: Initializer, static_activation, dynamic_activation. We have this separtion because they are using different memory offsets.
// group_name: The group_name that is recorded previously using function "AddRecordingTensorGroup". Used for generating customized sessions for a group of tensors
// Top_k: The steps with the top-k highest memory consumptions are plot. When top_k == 0, we plot all the steps
// device_t: The type of the device where the tensors are.
static void CreateEvents(const std::string& p_name, const size_t pid, const MemoryInfo::MapType& map_type,
const std::string& group_name, const size_t top_k,
const OrtDevice::DeviceType device_t = OrtDevice::GPU);
void Init(const SequentialExecutionPlan* execution_plan,
const OrtValueNameIdxMap& value_name_idx_map);
void RecordPatternInfo(const MemoryPatternGroup& mem_patterns, const MapType& type);
static const std::vector<std::string>& GetEvents() { return events; }
void RecordInitializerAllocInfo(const std::unordered_map<int, OrtValue>& tensor_map);
void RecordActivationAllocInfo(const OrtValueIndex idx, const OrtValue& value);
void SetDynamicAllocation(const OrtValueIndex idx);
void SetLocalRank(const int rank) { local_rank_ = rank; }
static size_t GetAndIncreasePid() {
size_t val = pid_++;
return val;
}
void SetIteration(size_t iteration) { iteration_ = iteration; }
void IncreaseIteration() { ++iteration_; }
size_t GetIteration() { return iteration_; }
static void Clear() {
summary_.clear();
}
int GetLocalRank() const { return local_rank_; }
private:
static size_t pid_;
static const std::vector<std::string> color_names;
static std::vector<std::string> events;
// Key: the hash function of device+map_type. Value: (key: The time step. value: The allocation information)
static std::unordered_map<size_t, std::unordered_map<size_t, AllocationSummary> > summary_;
MemoryInfoProfile() = default;
static std::string CreateMetadataEvent(const std::string& process_name, size_t process_id);
static std::string CreateMemoryEvent(size_t pid, size_t tid, const std::string& name, size_t offset, size_t size,
const std::string& color_name);
static std::string CreateSummaryEvent(size_t pid, size_t tid, const AllocationSummary& summary, size_t size,
size_t bytes_for_pattern);
};
static void GenerateTensorMap(const SequentialExecutionPlan* execution_plan,
const OrtValueNameIdxMap& value_name_idx_map);
static void RecordPatternInfo(const MemoryPatternGroup& mem_patterns, const MapType& type);
static void RecordInitializerAllocInfo(const std::unordered_map<int, OrtValue>& tensor_map);
static void RecordActivationAllocInfo(const OrtValueIndex idx, const OrtValue& value);
static void SetDynamicAllocation(const OrtValueIndex idx);
static void SetLocalRank(const int rank) { local_rank_ = rank; }
static void SetIteration(size_t iteration) { iteration_ = iteration; }
static void IncreaseIteration() { ++iteration_; }
static size_t GetIteration() { return iteration_; }
static void PrintMemoryInfoForLocation(const OrtDevice::DeviceType location);
static void GenerateMemoryProfile();
static void ClearMemoryInfoPerExecution() {
void PrintMemoryInfoForLocation(const OrtDevice::DeviceType location);
void ClearMemoryInfoPerExecution() {
for (auto& location_map : tensors_memory_info_map_) {
location_map.second[MapType::DynamicActivation].clear();
location_map.second[MapType::StaticActivation].clear();
@ -183,7 +154,7 @@ class MemoryInfo {
}
// Get the Allocation Information for tensor with index = idx;
static const AllocInfoPerTensor* AllocPlan(const OrtValueIndex& idx) {
const AllocInfoPerTensor* AllocPlan(const OrtValueIndex& idx) {
if (tensor_alloc_info_map_.find(idx) != tensor_alloc_info_map_.end())
return &tensor_alloc_info_map_.at(idx);
else
@ -193,11 +164,11 @@ class MemoryInfo {
// Create or add a tensor to a group. This is used to display the trace of specific group of tensors of interest
// group_name: The name of the group of tensors
// tensor_name: The tensor_name which is generated by argdef when building the graph
static void AddRecordingTensorGroup(const std::string& group_name, const std::string& tensor_name) {
void AddRecordingTensorGroup(const std::string& group_name, const std::string& tensor_name) {
customized_recording_group_[group_name][tensor_name] = true;
}
static bool InRecordingTensorGroup(const std::string& group_name, const std::string& tensor_name) {
bool InRecordingTensorGroup(const std::string& group_name, const std::string& tensor_name) {
if (customized_recording_group_.find(group_name) == customized_recording_group_.end()) return false;
if (customized_recording_group_.at(group_name).find(tensor_name) ==
customized_recording_group_.at(group_name).end()) return false;
@ -205,30 +176,113 @@ class MemoryInfo {
}
private:
MemoryInfo() = default;
static void RecordMemoryPatternInfo(const MemoryPatternGroup& mem_patterns, MapType type);
static void RecordTensorDeviceAllocInfo(const OrtValueIndex idx, const OrtValue& value, const MapType& type);
void RecordMemoryPatternInfo(const MemoryPatternGroup& mem_patterns, MapType type);
void RecordTensorDeviceAllocInfo(const OrtValueIndex idx, const OrtValue& value, const MapType& type);
static bool IsInplaceReuse(const OrtValueIndex& idx) {
bool IsInplaceReuse(const OrtValueIndex& idx) {
return tensor_alloc_info_map_[idx].inplace_reuse;
}
// Key: The map type. E.g., initializer, activation. Value: A map from the tensor index to its memory information
static std::map<OrtMemoryInfo, std::map<MapType, MemoryInfoMap> > tensors_memory_info_map_;
std::map<OrtMemoryInfo, std::map<MapType, MemoryInfoMap> > tensors_memory_info_map_;
// Key: The tensor index. Value: The Allocation information per tensor
static std::map<OrtValueIndex, AllocInfoPerTensor> tensor_alloc_info_map_;
std::map<OrtValueIndex, AllocInfoPerTensor> tensor_alloc_info_map_;
// Key: The group name. Value: (key: The tensor name, value: not used (Use std::map for faster lookup).).
static std::map<const std::string, std::map<const std::string, bool> > customized_recording_group_;
std::map<const std::string, std::map<const std::string, bool> > customized_recording_group_;
// TODO: The dynamic and statically planned alignments may not be the same, need to check
static const int alignment_ = 256;
static size_t iteration_;
static size_t num_node_size_;
static int local_rank_;
const int alignment_ = 256;
size_t iteration_{0};
size_t num_node_size_{0};
int local_rank_{0};
// Memory Profile
static std::map<MapType, std::set<size_t> > time_step_trace_;
std::map<MapType, std::set<size_t> > time_step_trace_;
};
// A monotonically increasing profiler id for differentiating different sessions.
static std::atomic<uint32_t> global_mem_profiler_id;
struct MemoryProfiler {
public:
MemoryProfiler() = default;
void Init(const SequentialExecutionPlan* execution_plan,
const OrtValueNameIdxMap& value_name_idx_map) {
profiler_id_ = global_mem_profiler_id.fetch_add(1);
memory_info_for_profiler_.Init(execution_plan, value_name_idx_map);
}
// Create sessions in the profiler.
// p_name: session name
// pid: sessionid
// map_type: Initializer, static_activation, dynamic_activation. We have this separtion because they are using
// different memory offsets.
// group_name: The group_name that is recorded previously using function "AddRecordingTensorGroup". Used for
// generating customized sessions for a group of tensors
// Top_k: The steps with the top-k highest memory consumptions are plot. When top_k == 0, we plot all the steps
// device_t: The type of the device where the tensors are.
void CreateEvents(const std::string& p_name, const size_t pid, const MemoryInfo::MapType& map_type,
const std::string& group_name, const size_t top_k,
const OrtDevice::DeviceType device_t = OrtDevice::GPU);
const std::vector<std::string>& GetEvents() { return events; }
size_t GetAndIncreasePid() {
size_t val = pid_++;
return val;
}
void Clear() {
summary_.clear();
}
MemoryInfo& GetMemoryInfo() { return memory_info_for_profiler_; }
void GenerateMemoryProfile();
private:
size_t pid_{0};
uint32_t profiler_id_;
// The following colors are defined and accepted by Chrome Tracing/Edge Tracing.
const std::vector<std::string> color_names = {
"good",
"bad",
"terrible",
"yellow",
"olive",
"generic_work",
"background_memory_dump",
"light_memory_dump",
"detailed_memory_dump",
"thread_state_uninterruptible",
"thread_state_iowait",
"thread_state_running",
"thread_state_runnable",
"thread_state_unknown",
"cq_build_running",
"cq_build_passed",
"cq_build_failed",
"cq_build_abandoned",
"cq_build_attempt_runnig",
"cq_build_attempt_passed",
"cq_build_attempt_failed",
};
std::vector<std::string> events;
// Key: the hash function of device+map_type. Value: (key: The time step. value: The allocation information)
std::unordered_map<size_t, std::unordered_map<size_t, MemoryInfo::AllocationSummary> > summary_;
std::string CreateMetadataEvent(const std::string& process_name, size_t process_id);
std::string CreateMemoryEvent(size_t pid, size_t tid, const std::string& name, size_t offset, size_t size,
const std::string& color_name);
std::string CreateSummaryEvent(size_t pid, size_t tid, const MemoryInfo::AllocationSummary& summary, size_t size,
size_t bytes_for_pattern);
MemoryInfo memory_info_for_profiler_;
};
} // namespace onnxruntime

17
onnxruntime/core/framework/orttraining_partial_executor.cc
View file

@ -157,6 +157,7 @@ Status PartialExecutor::Execute(const SessionState& session_state, gsl::span<con
// Enable TRACE_EXECUTION compile flag to dump execution plan
#if defined(TRACE_EXECUTION)
std::cout << std::make_pair(&seq_exec_plan, &session_state) << std::endl;
#endif
const auto& graph_viewer = session_state.GetGraphViewer();
@ -357,9 +358,10 @@ Status PartialExecutor::Execute(const SessionState& session_state, gsl::span<con
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
if (partial_graph_index_ == 1) {
// Only record memory consumption after backward partial graph execution.
MemoryInfo::MemoryInfoProfile::CreateEvents("dynamic activations_" + std::to_string(MemoryInfo::GetIteration()),
MemoryInfo::MemoryInfoProfile::GetAndIncreasePid(),
MemoryInfo::MapType::DynamicActivation, "", 0);
session_state.GetMemoryProfiler()->CreateEvents(
"dynamic activations_" + std::to_string(session_state.GetMemoryProfiler()->GetMemoryInfo().GetIteration()),
session_state.GetMemoryProfiler()->GetAndIncreasePid(),
MemoryInfo::MapType::DynamicActivation, "", 0);
}
#endif
const auto msg_string = ss.str();
@ -480,10 +482,11 @@ Status PartialExecutor::Execute(const SessionState& session_state, gsl::span<con
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
if (partial_graph_index_ == 1) {
// Only record memory consumption after backward partial graph execution.
MemoryInfo::MemoryInfoProfile::CreateEvents("dynamic activations_" + std::to_string(MemoryInfo::GetIteration()),
MemoryInfo::MemoryInfoProfile::GetAndIncreasePid(),
MemoryInfo::MapType::DynamicActivation, "", 0);
MemoryInfo::MemoryInfoProfile::Clear();
session_state.GetMemoryProfiler()->CreateEvents(
"dynamic activations_" + std::to_string(session_state.GetMemoryProfiler()->GetMemoryInfo().GetIteration()),
session_state.GetMemoryProfiler()->GetAndIncreasePid(),
MemoryInfo::MapType::DynamicActivation, "", 0);
session_state.GetMemoryProfiler()->Clear();
}
#endif

28
onnxruntime/core/framework/sequential_executor.cc
View file

@ -72,10 +72,10 @@ static void CalculateTotalOutputSizes(OpKernelContextInternal* op_kernel_context
#if defined(TRACE_EXECUTION)
const TensorShape& tensor_shape = tensor.Shape();
std::cout << node_name << " output[" << i << "]"
<< " size=" << tensor_size
<< " shape=" << tensor_shape.ToString()
<< " element_size=" << tensor.DataType()->Size()
<< "\n";
<< " size=" << tensor_size
<< " shape=" << tensor_shape.ToString()
<< " element_size=" << tensor.DataType()->Size()
<< "\n";
#endif
total_output_sizes += tensor_size;
auto shape_str = tensor.Shape().ToString();
@ -208,11 +208,10 @@ Status SequentialExecutor::Execute(const SessionState& session_state, gsl::span<
#endif
#ifdef DEBUG_NODE_INPUTS_OUTPUTS
size_t program_counter = 0;
utils::NodeDumpContext dump_context { session_state.GetGraphExecutionCounter(), program_counter };
size_t program_counter = 0;
utils::NodeDumpContext dump_context{session_state.GetGraphExecutionCounter(), program_counter};
#endif
for (const auto& node_exec_plan : exec_plan_vec) {
if (terminate_flag_) {
LOGS(logger, WARNING) << "Exiting due to terminate flag being set to true.";
@ -359,10 +358,12 @@ Status SequentialExecutor::Execute(const SessionState& session_state, gsl::span<
std::ostringstream ss;
ss << "Non-zero status code returned while running " << node.OpType() << " node. Name:'" << node.Name()
<< "' Status Message: " << compute_status.ErrorMessage();
//If the computation failed, we still can record the memory consumption
// If the computation failed, we still can record the memory consumption
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
MemoryInfo::MemoryInfoProfile::CreateEvents("dynamic activations_" + std::to_string(MemoryInfo::GetIteration()),
MemoryInfo::MemoryInfoProfile::GetAndIncreasePid(), MemoryInfo::MapType::DynamicActivation, "", 0);
session_state.GetMemoryProfiler()->CreateEvents(
"dynamic activations_" + std::to_string(session_state.GetMemoryProfiler()->GetMemoryInfo().GetIteration()),
session_state.GetMemoryProfiler()->GetAndIncreasePid(),
MemoryInfo::MapType::DynamicActivation, "", 0);
#endif
const auto msg_string = ss.str();
LOGS(logger, ERROR) << msg_string;
@ -480,9 +481,10 @@ Status SequentialExecutor::Execute(const SessionState& session_state, gsl::span<
VLOGS(logger, 1) << "Done with execution.";
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
MemoryInfo::MemoryInfoProfile::CreateEvents("dynamic activations_" + std::to_string(MemoryInfo::GetIteration()),
MemoryInfo::MemoryInfoProfile::GetAndIncreasePid(), MemoryInfo::MapType::DynamicActivation, "", 0);
MemoryInfo::MemoryInfoProfile::Clear();
session_state.GetMemoryProfiler()->CreateEvents(
"dynamic activations_" + std::to_string(session_state.GetMemoryProfiler()->GetMemoryInfo().GetIteration()),
session_state.GetMemoryProfiler()->GetAndIncreasePid(), MemoryInfo::MapType::DynamicActivation, "", 0);
session_state.GetMemoryProfiler()->Clear();
#endif
if (frame.HasMemoryPatternPlanner()) {

27
onnxruntime/core/framework/session_state.cc
View file

@ -679,7 +679,6 @@ const MemoryPatternGroup* SessionState::GetMemoryPatternGroup(
gsl::span<const OrtValue> tensor_inputs,
gsl::span<const int> feed_mlvalue_idxs,
const InlinedHashMap<int, TensorShape>*& out_inferred_shapes) const {
out_inferred_shapes = nullptr;
int64_t key = CalculateMemoryPatternsKey(tensor_inputs);
std::lock_guard<OrtMutex> lock(mem_patterns_lock_);
@ -708,7 +707,6 @@ const MemoryPatternGroup* SessionState::GetMemoryPatternGroup(
return &it->second;
}
void SessionState::ResolveMemoryPatternFlag() {
if (enable_mem_pattern_) {
for (auto* input : graph_viewer_->GetInputs()) {
@ -1420,12 +1418,12 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_
subgraphs_kernel_create_info_maps,
outer_scope_node_arg_to_location_map,
ort_value_name_idx_map_, context, p_seq_exec_plan_));
// Record the allocation plan
// Record the allocation plan
// Uncomment the below to dump the allocation plan to std::cout
// LOGS(logger_, VERBOSE) << std::make_pair(p_seq_exec_plan_.get(), this);
// Uncomment the below to dump the allocation plan to std::cout
// LOGS(logger_, VERBOSE) << std::make_pair(p_seq_exec_plan_.get(), this);
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
MemoryInfo::GenerateTensorMap(GetExecutionPlan(), GetOrtValueNameIdxMap());
GetMemoryProfiler()->Init(GetExecutionPlan(), GetOrtValueNameIdxMap());
#endif
// Memory pattern tracer allocates all initializers on a single contiguous
@ -1453,6 +1451,18 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_
const auto& initializer_allocation_order = p_seq_exec_plan_->initializer_allocation_order;
// move initializers from TensorProto instances in Graph to OrtValue instances in SessionState
session_state_utils::MemoryProfileFunction memory_propfile_func = nullptr;
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
memory_propfile_func = [this](ITensorAllocator& planner) {
GetMemoryProfiler()->GetMemoryInfo().RecordPatternInfo(
planner.GetMemPatterns(), MemoryInfo::MapType::Initializer);
GetMemoryProfiler()->CreateEvents(
"initializer_" + std::to_string(GetMemoryProfiler()->GetMemoryInfo().GetIteration()),
GetMemoryProfiler()->GetAndIncreasePid(), MemoryInfo::MapType::Initializer, "", 0);
};
#endif
ORT_RETURN_IF_ERROR(
session_state_utils::SaveInitializedTensors(
Env::Default(), graph_location, *graph_viewer_,
@ -1461,10 +1471,11 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_
[this](int idx, const OrtValue& value, const OrtCallback& d, bool constant, bool sparse) -> Status {
return AddInitializedTensor(idx, value, &d, constant, sparse);
},
logger_, data_transfer_mgr_, *p_seq_exec_plan_, session_options));
logger_, data_transfer_mgr_, *p_seq_exec_plan_, session_options, memory_propfile_func));
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
// Record Weight allocation info on device
MemoryInfo::RecordInitializerAllocInfo(GetInitializedTensors());
GetMemoryProfiler()->GetMemoryInfo().RecordInitializerAllocInfo(GetInitializedTensors());
#endif
// remove weights from the graph now to save memory but in many cases it won't save memory, if the tensor was

14
onnxruntime/core/framework/session_state.h
View file

@ -203,6 +203,14 @@ class SessionState {
*/
profiling::Profiler& Profiler() const noexcept { return profiler_; }
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
MemoryProfiler* GetMemoryProfiler() const noexcept { return memory_profiler_; }
void SetMemoryProfiler(MemoryProfiler* memory_profiler) noexcept {
memory_profiler_ = memory_profiler;
}
#endif
/**
Get cached memory pattern based on input shapes
Must be called only when all values contain tensors
@ -386,7 +394,7 @@ class SessionState {
#ifdef ENABLE_TRAINING
Status GeneratePatternGroupCache(
gsl::span<const OrtValue> inputs,
gsl::span<const int > feed_mlvalue_idxs,
gsl::span<const int> feed_mlvalue_idxs,
MemoryPatternGroup& output,
InlinedHashMap<int, TensorShape>& inferred_shapes) const;
#endif
@ -480,6 +488,10 @@ class SessionState {
const logging::Logger& logger_;
profiling::Profiler& profiler_;
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
MemoryProfiler* memory_profiler_;
#endif
// switch for enable memory pattern optimization or not.
bool enable_mem_pattern_;

11
onnxruntime/core/framework/session_state_utils.cc
View file

@ -195,7 +195,8 @@ common::Status SaveInitializedTensors(
const SaveTensorFunction& save_tensor_func,
const logging::Logger& logger, const DataTransferManager& data_transfer_mgr,
const ExecutionPlanBase& exec_plan,
const SessionOptions& session_options) {
const SessionOptions& session_options,
const MemoryProfileFunction& memory_profile_func) {
LOGS(logger, INFO) << "Saving initialized tensors.";
ORT_ENFORCE(ort_value_name_idx_map.MaxIdx() > -1, "OrtValue indexes should have been populated.");
@ -276,11 +277,9 @@ common::Status SaveInitializedTensors(
InlinedHashMap<std::string, size_t> planned_initializers_memory_sizes_in_byte;
ORT_RETURN_IF_ERROR(
planner.FinalizePlan(planned_initializers_memory_sizes_in_byte));
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
MemoryInfo::RecordPatternInfo(planner.GetMemPatterns(), MemoryInfo::MapType::Initializer);
MemoryInfo::MemoryInfoProfile::CreateEvents("initializer_" + std::to_string(MemoryInfo::GetIteration()),
MemoryInfo::MemoryInfoProfile::GetAndIncreasePid(), MemoryInfo::MapType::Initializer, "", 0);
#endif
if (memory_profile_func)
memory_profile_func(planner);
for (auto i : planned_initializers_memory_sizes_in_byte) {
LOGS(logger, INFO) << "[Memory] SessionStateInitializer statically allocates "

4
onnxruntime/core/framework/session_state_utils.h
View file

@ -32,6 +32,7 @@ class Logger;
namespace session_state_utils {
using SaveTensorFunction = std::function<Status(int idx, const OrtValue& value, const OrtCallback& d,
bool constant, bool sparse)>;
using MemoryProfileFunction = std::function<void(ITensorAllocator& planner)>;
common::Status SaveInitializedTensors(
const Env& env, const std::basic_string<PATH_CHAR_TYPE>& graph_loc,
const GraphViewer& graph, const AllocatorPtr& default_cpu_memory_info,
@ -41,7 +42,8 @@ common::Status SaveInitializedTensors(
const logging::Logger& logger,
const DataTransferManager& data_transfer_mgr,
const ExecutionPlanBase& exec_plan,
const SessionOptions& session_options);
const SessionOptions& session_options,
const MemoryProfileFunction& memory_profile_func);
common::Status SaveInputOutputNamesToNodeMapping(const GraphViewer& graph,
SessionState& session_state,
gsl::span<const NodeArg* const> implicit_inputs);

7
onnxruntime/core/session/inference_session.cc
View file

@ -471,7 +471,7 @@ InferenceSession::~InferenceSession() {
TraceLoggingWriteStop(session_activity, "OrtInferenceSessionActivity");
#endif
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
MemoryInfo::GenerateMemoryProfile();
GetMemoryProfiler().GenerateMemoryProfile();
#endif
}
@ -1353,6 +1353,11 @@ common::Status InferenceSession::Initialize() {
session_options_.enable_mem_reuse,
prepacked_weights_container_);
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
// Don't want to pollute SessionState constructor since memory profile is enabled optionally.
session_state_->SetMemoryProfiler(&memory_profiler_);
#endif
// Collect the kernel registries from execution provider instances;
// There are 2 kinds of kernel registries with priority from high to low as below,
// 1. Custom execution provider type specific kernel registries.

10
onnxruntime/core/session/inference_session.h
View file

@ -443,6 +443,12 @@ class InferenceSession {
*/
const profiling::Profiler& GetProfiling() const;
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
MemoryProfiler& GetMemoryProfiler() {
return memory_profiler_;
}
#endif
/**
* Search registered execution providers for an allocator that has characteristics
* specified within mem_info
@ -657,6 +663,10 @@ class InferenceSession {
// Profiler for this session.
profiling::Profiler session_profiler_;
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
MemoryProfiler memory_profiler_;
#endif
// Immutable state for each op in the model. Shared by all executors.
// It has a dependency on execution_providers_.
std::unique_ptr<SessionState> session_state_;

6
orttraining/orttraining/core/agent/training_agent.cc
View file

@ -14,9 +14,9 @@ TrainingAgent::TrainingAgent(InferenceSession& session,
const std::vector<std::string>& bw_fetches_names,
const std::vector<OrtDevice>& bw_outputs_device_info,
int local_rank) : inference_session_(session) {
ORT_UNUSED_PARAMETER(local_rank);
ORT_UNUSED_PARAMETER(local_rank);
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
MemoryInfo::SetLocalRank(local_rank);
inference_session_.GetMemoryProfiler().GetMemoryInfo().SetLocalRank(local_rank);
#endif
auto& session_state = session.GetSessionState();
std::vector<std::string> fw_fetches_names;
@ -57,7 +57,7 @@ TrainingAgent::~TrainingAgent() = default;
common::Status TrainingAgent::RunForward(const std::vector<OrtValue>& feeds, std::vector<OrtValue>& fetches,
PartialGraphExecutionState& state, const OrtValueCachePtr& cache) {
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
MemoryInfo::SetIteration(profile_step_);
inference_session_.GetMemoryProfiler().GetMemoryInfo().SetIteration(profile_step_);
profile_step_ += 1;
#endif

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

@ -25,7 +25,7 @@
#include "orttraining/models/runner/training_util.h"
#include "orttraining/core/optimizer/megatron_transformer.h"
//Gist Encoding
// Gist Encoding
#include "orttraining/core/optimizer/gist_encode_decode.h"
#include "orttraining/training_ops/cpu/controlflow/event_pool.h"
@ -349,7 +349,7 @@ Status TrainingSession::ConfigureForTraining(
config.distributed_config.horizontal_parallel_size,
config.distributed_config.pipeline_parallel_size});
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
MemoryInfo::SetLocalRank(config.distributed_config.world_rank);
GetMemoryProfiler().GetMemoryInfo().SetLocalRank(config.distributed_config.world_rank);
#endif
#ifdef USE_MPI

42
orttraining/orttraining/models/runner/training_runner.cc
View file

@ -34,26 +34,26 @@ namespace training {
static std::vector<FreeDimensionOverride> overrides = {};
static SessionOptions SESSION_OPTION = {
ExecutionMode::ORT_SEQUENTIAL, //execution_mode
ExecutionOrder::PRIORITY_BASED, //execution_order
false, //enable_profiling
ORT_TSTR(""), //optimized_model_filepath
true, //enable_mem_pattern
true, //enable_mem_reuse
true, //enable_cpu_mem_arena
ORT_TSTR("onnxruntime_profile_"), //profile_file_prefix
"", //session_logid
-1, //session_log_severity_level
0, //session_log_verbosity_level
5, //max_num_graph_transformation_steps
TransformerLevel::Level1, //graph_optimization_level
{}, //intra_op_param
{}, //inter_op_param
overrides, //free_dimension_overrides
true, //use_per_session_threads
true, //thread_pool_allow_spinning
false, //use_deterministic_compute
{}, //config_options
ExecutionMode::ORT_SEQUENTIAL, // execution_mode
ExecutionOrder::PRIORITY_BASED, // execution_order
false, // enable_profiling
ORT_TSTR(""), // optimized_model_filepath
true, // enable_mem_pattern
true, // enable_mem_reuse
true, // enable_cpu_mem_arena
ORT_TSTR("onnxruntime_profile_"), // profile_file_prefix
"", // session_logid
-1, // session_log_severity_level
0, // session_log_verbosity_level
5, // max_num_graph_transformation_steps
TransformerLevel::Level1, // graph_optimization_level
{}, // intra_op_param
{}, // inter_op_param
overrides, // free_dimension_overrides
true, // use_per_session_threads
true, // thread_pool_allow_spinning
false, // use_deterministic_compute
{}, // config_options
{}, // initializers_to_share_map
{}, // external_initializers
};
@ -810,7 +810,7 @@ Status TrainingRunner::TrainingLoop(IDataLoader& training_data_loader, IDataLoad
size_t batch_num_cur_shard = training_data->TotalBatch(params_.batch_size);
for (size_t batch = 0; batch < batch_num_cur_shard && step_ < params_.num_train_steps; ++batch) {
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
MemoryInfo::SetIteration(step_);
GetSession().GetMemoryProfiler().GetMemoryInfo().SetIteration(step_);
#endif
const bool is_weight_update_step = (step_ + 1) % params_.gradient_accumulation_steps == 0;