mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
fix memory profile for partial graph run (#11911)
* fix mpi build for gcc8 or higher * fix memory profile for partial graph run * Revert "fix mpi build for gcc8 or higher" This reverts commit fb60beb05402cd380597a12fc25880c0c8652ed4. * remove debug code * fix build * fix build * fix cpplint and python black format
This commit is contained in:
pengwa
GitHub
parent
fa7f80c847
commit
0d6cbc6e57
13 changed files with 222 additions and 138 deletions
|
|
@ -8,6 +8,7 @@
|
|||
#include <queue>
|
||||
|
||||
namespace onnxruntime {
|
||||
|
||||
size_t MemoryInfo::iteration_ = 0;
|
||||
size_t MemoryInfo::num_node_size_ = 0;
|
||||
int MemoryInfo::local_rank_ = 0;
|
||||
|
|
@ -16,26 +17,31 @@ std::map<OrtValueIndex, MemoryInfo::AllocInfoPerTensor> MemoryInfo::tensor_alloc
|
|||
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) {
|
||||
// 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;
|
||||
}
|
||||
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()))
|
||||
// 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));
|
||||
mem_info.lifetime_interval = execution_plan->allocation_plan[value_idx].life_interval;
|
||||
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
|
||||
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;
|
||||
if (execution_plan->allocation_plan[mem_info.reused_buffer].alloc_kind == AllocKind::kAllocatedExternally) continue;
|
||||
mem_info.inplace_reuse = (execution_plan->allocation_plan[value_idx].inplace_reuse != -1 && execution_plan->allocation_plan[value_idx].inplace_reuse != value_idx);
|
||||
mem_info.inplace_reuse = (execution_plan->allocation_plan[value_idx].inplace_reuse != -1 &&
|
||||
execution_plan->allocation_plan[value_idx].inplace_reuse != value_idx);
|
||||
mem_info.alloc_kind = execution_plan->allocation_plan[value_idx].alloc_kind;
|
||||
mem_info.location = execution_plan->allocation_plan[value_idx].location;
|
||||
|
||||
|
|
@ -46,7 +52,7 @@ void MemoryInfo::GenerateTensorMap(const SequentialExecutionPlan* execution_plan
|
|||
return;
|
||||
}
|
||||
|
||||
//Record the planned memory information
|
||||
// Record the planned memory information
|
||||
void MemoryInfo::RecordMemoryPatternInfo(const MemoryPatternGroup& mem_patterns, MapType type) {
|
||||
for (const auto& location : mem_patterns.locations) {
|
||||
for (const auto& p : mem_patterns.GetPatterns(location)->GetPatternsMap()) {
|
||||
|
|
@ -70,12 +76,17 @@ void MemoryInfo::RecordPatternInfo(const MemoryPatternGroup& mem_patterns, const
|
|||
}
|
||||
}
|
||||
|
||||
//Record the actual allocated tensor in the device
|
||||
// Record the actual allocated tensor in the device
|
||||
void MemoryInfo::RecordTensorDeviceAllocInfo(const OrtValueIndex idx, const OrtValue& value, const MapType& type) {
|
||||
if (tensor_alloc_info_map_.find(idx) == tensor_alloc_info_map_.end()) return;
|
||||
ORT_ENFORCE(value.IsTensor(), "Memory profiler only supports tensor type.");
|
||||
auto& tensor = value.Get<Tensor>();
|
||||
auto sizeinbytes = tensor.SizeInBytes() % alignment ? tensor.SizeInBytes() + alignment - tensor.SizeInBytes() % alignment : tensor.SizeInBytes();
|
||||
MemoryBlock mb(size_t(tensor.DataRaw()), sizeinbytes);
|
||||
auto tensor_size_in_bytes = tensor.SizeInBytes();
|
||||
auto remainder = tensor_size_in_bytes % alignment_;
|
||||
auto size_in_bytes = tensor_size_in_bytes % alignment_ ? tensor_size_in_bytes + alignment_ - remainder
|
||||
: tensor_size_in_bytes;
|
||||
MemoryBlock mb((size_t)(tensor.DataRaw()), size_in_bytes);
|
||||
|
||||
if (type == MapType::StaticActivation) {
|
||||
bool valid = tensors_memory_info_map_.at(AllocPlan(idx)->location).at(type).Contain(idx);
|
||||
ORT_ENFORCE(valid);
|
||||
|
|
@ -91,6 +102,8 @@ void MemoryInfo::RecordInitializerAllocInfo(const std::unordered_map<int, OrtVal
|
|||
}
|
||||
|
||||
void MemoryInfo::RecordActivationAllocInfo(const OrtValueIndex idx, const OrtValue& value) {
|
||||
// Some activations are not tracked in GenerateTensorMap, for example some pre-existing tensors are re-used with
|
||||
// Alias by some operators, it will reuse this buffer and the buffer will not be tracked in GenerateTensorMap.
|
||||
if (!AllocPlan(idx)) return;
|
||||
auto reuse_buffer = AllocPlan(idx)->reused_buffer;
|
||||
MapType map_type;
|
||||
|
|
@ -102,7 +115,7 @@ void MemoryInfo::RecordActivationAllocInfo(const OrtValueIndex idx, const OrtVal
|
|||
else if (map[MapType::Initializer].Contain(reuse_buffer))
|
||||
map_type = MapType::Initializer;
|
||||
else
|
||||
ORT_ENFORCE(false, "Unsupported tensor type.");
|
||||
std::cout << "Find no map type for reuse_buffer: " << reuse_buffer << ", so skipping" << std::endl;
|
||||
|
||||
RecordTensorDeviceAllocInfo(idx, value, map_type);
|
||||
}
|
||||
|
|
@ -113,7 +126,8 @@ void MemoryInfo::SetDynamicAllocation(const OrtValueIndex idx) {
|
|||
if (!da_map.Contain(idx)) da_map[idx];
|
||||
}
|
||||
|
||||
void PrintInforPerTensor(const MemoryInfo::AllocInfoPerTensor& alloc_info, const MemoryInfoPerTensor& mem_info, const size_t& rel_addr) {
|
||||
void PrintInforPerTensor(const MemoryInfo::AllocInfoPerTensor& alloc_info, const MemoryInfoPerTensor& mem_info,
|
||||
const size_t& rel_addr) {
|
||||
std::cout << "Tensor name: " << alloc_info.mlvalue_name << ", ";
|
||||
std::cout << "Index: " << alloc_info.mlvalue_index << ", ";
|
||||
std::cout << "Reuse inplace: " << alloc_info.inplace_reuse << ", ";
|
||||
|
|
@ -171,7 +185,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 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::stringstream evt;
|
||||
evt << "{";
|
||||
evt << "\"ph\":\"X\",";
|
||||
|
|
@ -190,7 +206,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 MemoryInfo::MemoryInfoProfile::CreateSummaryEvent(size_t pid, size_t tid, const 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;
|
||||
|
|
@ -222,7 +239,8 @@ static bool IsStaticType(const MemoryInfo::MapType& map_type) {
|
|||
return map_type == MemoryInfo::MapType::Initializer || map_type == MemoryInfo::MapType::StaticActivation;
|
||||
}
|
||||
|
||||
static void UpdateSummary(MemoryInfo::AllocationSummary& summary, size_t alloc_offset, size_t alloc_size, const OrtValueIndex idx, const MemoryInfo::MapType& map_type) {
|
||||
static void UpdateSummary(MemoryInfo::AllocationSummary& summary, size_t alloc_offset, size_t alloc_size,
|
||||
const OrtValueIndex idx, const MemoryInfo::MapType& map_type) {
|
||||
summary.total_size = std::max(summary.total_size, alloc_offset + alloc_size);
|
||||
summary.used_size += alloc_size;
|
||||
if (!IsStaticType(map_type)) {
|
||||
|
|
@ -233,7 +251,7 @@ 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.
|
||||
// The following colors are defined and accepted by Chrome Tracing/Edge Tracing.
|
||||
const std::vector<std::string> MemoryInfo::MemoryInfoProfile::color_names = {
|
||||
"good",
|
||||
"bad",
|
||||
|
|
@ -260,59 +278,77 @@ const std::vector<std::string> MemoryInfo::MemoryInfoProfile::color_names = {
|
|||
|
||||
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_;
|
||||
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
|
||||
//map_type: Initalizer, 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 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_t) {
|
||||
// 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_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) {
|
||||
// Metadata.
|
||||
std::string pid_name_internal = "device_" + std::to_string(device_t) + "_" + p_name + group_name;
|
||||
std::string pid_name_internal = "device_" + std::to_string(device_type) + "_" + p_name + group_name;
|
||||
events.push_back(CreateMetadataEvent(pid_name_internal, pid));
|
||||
size_t summary_size = 10;
|
||||
std::hash<std::string> str_hash;
|
||||
|
||||
//Create Event for each tensor
|
||||
// Create Event for each tensor
|
||||
for (const auto& location_map : tensors_memory_info_map_) {
|
||||
if (location_map.first.device.Type() != device_t) continue;
|
||||
//If there is no tensor of a map_type, we skip creating event for that map_type
|
||||
if (location_map.second.find(map_type) == location_map.second.end()) continue;
|
||||
auto summary_key = str_hash(location_map.first.ToString() + std::to_string(map_type));
|
||||
//Preprocessing
|
||||
const OrtMemoryInfo& memory_info = location_map.first;
|
||||
const auto& maptype_to_map_mapping = location_map.second;
|
||||
|
||||
if (memory_info.device.Type() != device_type) continue;
|
||||
|
||||
// If there is no tensor of a map_type, we skip creating event for that map_type
|
||||
if (maptype_to_map_mapping.find(map_type) == maptype_to_map_mapping.end()) continue;
|
||||
|
||||
auto summary_key = str_hash(memory_info.ToString() + std::to_string(map_type));
|
||||
// Preprocessing
|
||||
if (time_step_trace_[map_type].empty()) {
|
||||
for (const auto& item : location_map.second.at(map_type)) {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
const auto& map = location_map.second.at(map_type);
|
||||
//Update summary
|
||||
const auto& map = maptype_to_map_mapping.at(map_type);
|
||||
// Update summary
|
||||
if (summary_.find(summary_key) == summary_.end()) {
|
||||
for (const auto& item : map) {
|
||||
const auto info = AllocPlan(item.first);
|
||||
const OrtValueIndex idx = item.first;
|
||||
const auto info = AllocPlan(idx);
|
||||
if (info->inplace_reuse) continue;
|
||||
size_t offset = IsStaticType(map_type) ? map.GetPlannedAddress(item.first) : map.GetAllocAddress(item.first);
|
||||
size_t size = IsStaticType(map_type) ? map.GetPlannedSize(item.first) : map.GetAllocSize(item.first);
|
||||
size_t alloc_step = AllocPlan(item.first)->lifetime_interval.first;
|
||||
size_t dealloc_step = AllocPlan(item.first)->lifetime_interval.second;
|
||||
|
||||
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];
|
||||
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() && end_itr != ts_map.end(), "The allocation step is not recorded");
|
||||
ORT_ENFORCE(start_itr != ts_map.end(),
|
||||
"The allocation step is not recorded: ", alloc_step, ", ortvalue_index: ", idx);
|
||||
ORT_ENFORCE(end_itr != ts_map.end(),
|
||||
"The deallocation step is not recorded: ", dealloc_step, ", ortvalue_index: ", idx);
|
||||
|
||||
for (auto itr = start_itr; itr != end_itr; ++itr) {
|
||||
UpdateSummary(summary_[summary_key][*itr], offset, size, item.first, map_type);
|
||||
UpdateSummary(summary_[summary_key][*itr], offset, size, idx, map_type);
|
||||
}
|
||||
UpdateSummary(summary_[summary_key][*end_itr], offset, size, item.first, map_type);
|
||||
UpdateSummary(summary_[summary_key][*end_itr], offset, size, idx, map_type);
|
||||
}
|
||||
}
|
||||
|
||||
//extract top_k total size
|
||||
// extract top_k total size
|
||||
size_t top_kth_total_size = 0;
|
||||
if (top_k != 0) {
|
||||
std::set<size_t> top_k_set;
|
||||
|
|
@ -335,19 +371,21 @@ void MemoryInfo::MemoryInfoProfile::CreateEvents(const std::string& p_name, cons
|
|||
const auto info = AllocPlan(live_tensor);
|
||||
if (info->inplace_reuse) continue;
|
||||
const std::string& name = info->mlvalue_name;
|
||||
//Filter out string without a certain name
|
||||
// Filter out string without a certain name
|
||||
if (!group_name.empty()) {
|
||||
if (!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 address/size in static_activation type
|
||||
// Sometimes a tensor can be both statically planned and dynamically allocated, so we need to use planned
|
||||
// address/size in static_activation type
|
||||
size_t offset = IsStaticType(map_type) ? map.GetPlannedAddress(live_tensor) : map.GetAllocAddress(live_tensor);
|
||||
size_t size = IsStaticType(map_type) ? map.GetPlannedSize(live_tensor) : map.GetAllocSize(live_tensor);
|
||||
alloc_size_for_pattern += size;
|
||||
events.push_back(CreateMemoryEvent(pid, item.first, name, offset, size, cname));
|
||||
has_other_tensors = true;
|
||||
}
|
||||
//If for that time steps, we have other tensors to plot other than just summary, add summary events. These will show up visually as 10% of the max width in a section.
|
||||
// If for that time steps, we have other tensors to plot other than just summary, add summary events.
|
||||
// These will show up visually as 10 % of the max width in a section.
|
||||
if (has_other_tensors) {
|
||||
summary_size = std::max(summary_size, item.second.total_size / 10);
|
||||
events.push_back(CreateSummaryEvent(pid, item.first, item.second, summary_size, alloc_size_for_pattern));
|
||||
|
|
|
|||
|
|
@ -85,16 +85,16 @@ struct MemoryInfoMap {
|
|||
|
||||
private:
|
||||
MemoryInfoMapT map_;
|
||||
size_t ptr_offset{0}; //The start ptr of the raw data
|
||||
size_t ptr_offset{0}; // The start ptr of the raw data
|
||||
};
|
||||
|
||||
class MemoryInfo {
|
||||
public:
|
||||
//We separate the memory types into 3 categories.
|
||||
// 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,
|
||||
|
|
@ -109,7 +109,7 @@ class MemoryInfo {
|
|||
|
||||
IntervalT lifetime_interval{0, 0};
|
||||
bool inplace_reuse{false};
|
||||
OrtValueIndex reused_buffer{0}; //The index of the reused tensor, if no reuse, it is its own tensor.
|
||||
OrtValueIndex reused_buffer{0}; // The index of the reused tensor, if no reuse, it is its own tensor.
|
||||
AllocKind alloc_kind{AllocKind::kAllocate};
|
||||
OrtMemoryInfo location;
|
||||
};
|
||||
|
|
@ -122,15 +122,16 @@ class MemoryInfo {
|
|||
|
||||
struct MemoryInfoProfile {
|
||||
public:
|
||||
//Create sessions in the profiler.
|
||||
//p_name: session name
|
||||
//pid: sessionid
|
||||
//map_type: Initalizer, 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);
|
||||
// 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);
|
||||
|
||||
static const std::vector<std::string>& GetEvents() { return events; }
|
||||
|
||||
|
|
@ -147,18 +148,21 @@ class MemoryInfo {
|
|||
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)
|
||||
// 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 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 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 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);
|
||||
|
|
@ -178,7 +182,7 @@ class MemoryInfo {
|
|||
}
|
||||
}
|
||||
|
||||
//Get the Allocation Information for tensor with index = idx;
|
||||
// Get the Allocation Information for tensor with index = idx;
|
||||
static 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);
|
||||
|
|
@ -186,16 +190,17 @@ class MemoryInfo {
|
|||
return nullptr;
|
||||
}
|
||||
|
||||
//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
|
||||
// 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) {
|
||||
customized_recording_group_[group_name][tensor_name] = true;
|
||||
}
|
||||
|
||||
static 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;
|
||||
if (customized_recording_group_.at(group_name).find(tensor_name) ==
|
||||
customized_recording_group_.at(group_name).end()) return false;
|
||||
return true;
|
||||
}
|
||||
|
||||
|
|
@ -208,21 +213,21 @@ class MemoryInfo {
|
|||
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
|
||||
// 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_;
|
||||
|
||||
//Key: The tensor index. Value: The Allocation information per tensor
|
||||
// Key: The tensor index. Value: The Allocation information per tensor
|
||||
static 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).).
|
||||
// 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_;
|
||||
|
||||
//TODO: The dynamic and statically planned alignments may not be the same, need to check
|
||||
static const int alignment = 256;
|
||||
// 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_;
|
||||
//Memory Profile
|
||||
// Memory Profile
|
||||
static std::map<MapType, std::set<size_t> > time_step_trace_;
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -157,7 +157,6 @@ Status PartialExecutor::Execute(const SessionState& session_state, const std::ve
|
|||
|
||||
// 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();
|
||||
|
|
@ -185,10 +184,9 @@ Status PartialExecutor::Execute(const SessionState& session_state, const std::ve
|
|||
#endif
|
||||
|
||||
#ifdef DEBUG_NODE_INPUTS_OUTPUTS
|
||||
utils::NodeDumpContext dump_context { session_state.GetGraphExecutionCounter(), 0 };
|
||||
utils::NodeDumpContext dump_context{session_state.GetGraphExecutionCounter(), 0};
|
||||
#endif
|
||||
|
||||
|
||||
for (size_t program_counter = state_.GetProgramCounterStart();
|
||||
program_counter < state_.GetProgramCounterEnd();
|
||||
program_counter += 1) {
|
||||
|
|
@ -355,11 +353,14 @@ Status PartialExecutor::Execute(const SessionState& session_state, const std::ve
|
|||
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);
|
||||
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);
|
||||
}
|
||||
#endif
|
||||
const auto msg_string = ss.str();
|
||||
LOGS(logger, ERROR) << msg_string;
|
||||
|
|
@ -449,7 +450,7 @@ Status PartialExecutor::Execute(const SessionState& session_state, const std::ve
|
|||
utils::DumpNodeOutputs(dump_context, op_kernel_context, p_op_kernel->Node(), session_state);
|
||||
#endif
|
||||
|
||||
// free ml-values corresponding to this node
|
||||
// Free ml-values corresponding to this node
|
||||
VLOGS(logger, 1) << "Releasing node ML values.";
|
||||
ORT_RETURN_IF_ERROR(ReleaseNodeMLValues(frame, seq_exec_plan, node_exec_plan, logger));
|
||||
}
|
||||
|
|
@ -477,10 +478,13 @@ Status PartialExecutor::Execute(const SessionState& session_state, const std::ve
|
|||
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();
|
||||
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();
|
||||
}
|
||||
#endif
|
||||
|
||||
if (frame.HasMemoryPatternPlanner()) {
|
||||
|
|
|
|||
|
|
@ -20,9 +20,13 @@ namespace onnxruntime {
|
|||
class PartialExecutor : public IExecutor {
|
||||
public:
|
||||
PartialExecutor(PartialGraphExecutionState& state,
|
||||
const OrtValueCachePtr& cache)
|
||||
const OrtValueCachePtr& cache,
|
||||
int32_t partial_graph_index = 0)
|
||||
: state_{state},
|
||||
cache_{cache} {}
|
||||
cache_{cache},
|
||||
partial_graph_index_{partial_graph_index} {
|
||||
ORT_UNUSED_PARAMETER(partial_graph_index_);
|
||||
}
|
||||
|
||||
common::Status Execute(const SessionState& session_state, const std::vector<int>& feed_mlvalue_idxs,
|
||||
const std::vector<OrtValue>& feeds, const std::vector<int>& fetch_mlvalue_idxs,
|
||||
|
|
@ -34,6 +38,7 @@ class PartialExecutor : public IExecutor {
|
|||
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(PartialExecutor);
|
||||
PartialGraphExecutionState& state_;
|
||||
const OrtValueCachePtr& cache_;
|
||||
int32_t partial_graph_index_{0};
|
||||
};
|
||||
} // namespace onnxruntime
|
||||
#endif
|
||||
#endif
|
||||
|
|
|
|||
|
|
@ -646,10 +646,11 @@ common::Status ExecuteGraph(const SessionState& session_state,
|
|||
common::Status ExecutePartialGraph(const SessionState& session_state, FeedsFetchesManager& feeds_fetches_manager,
|
||||
const std::vector<OrtValue>& feeds, std::vector<OrtValue>& fetches,
|
||||
const logging::Logger& logger, PartialGraphExecutionState& state,
|
||||
const OrtValueCachePtr& cache) {
|
||||
const OrtValueCachePtr& cache,
|
||||
int32_t partial_graph_index) {
|
||||
// finalize the copy info using the provided feeds and fetches. will update device_copy_checks in the background
|
||||
FinalizeFeedFetchCopyInfo(feeds_fetches_manager, feeds, fetches);
|
||||
PartialExecutor executor{state, cache};
|
||||
PartialExecutor executor{state, cache, partial_graph_index};
|
||||
const auto& feeds_fetches_info = feeds_fetches_manager.GetFeedsFetchesInfo();
|
||||
const auto& device_copy_checks = feeds_fetches_manager.GetDeviceCopyChecks();
|
||||
|
||||
|
|
|
|||
|
|
@ -95,7 +95,8 @@ common::Status ExecuteGraph(const SessionState& session_state, FeedsFetchesManag
|
|||
common::Status ExecutePartialGraph(const SessionState& session_state, FeedsFetchesManager& feeds_fetches_manager,
|
||||
const std::vector<OrtValue>& feeds, std::vector<OrtValue>& fetches,
|
||||
const logging::Logger& logger, PartialGraphExecutionState& state,
|
||||
const OrtValueCachePtr& cache);
|
||||
const OrtValueCachePtr& cache,
|
||||
int32_t partial_graph_index);
|
||||
#endif
|
||||
|
||||
// Execute a subgraph. The feeds_fetches_manager should have been finalized prior to calling this function.
|
||||
|
|
|
|||
|
|
@ -1769,7 +1769,8 @@ Status InferenceSession::PartialRun(onnxruntime::RunOptions& run_options,
|
|||
std::vector<OrtValue>& fetches,
|
||||
PartialGraphExecutionState& state,
|
||||
FeedsFetchesManager& feeds_fetches_manager,
|
||||
const OrtValueCachePtr& cache) {
|
||||
const OrtValueCachePtr& cache,
|
||||
int32_t partial_graph_index) {
|
||||
Status retval = Status::OK();
|
||||
std::vector<IExecutionProvider*> exec_providers_to_stop;
|
||||
exec_providers_to_stop.reserve(execution_providers_.NumProviders());
|
||||
|
|
@ -1815,7 +1816,7 @@ Status InferenceSession::PartialRun(onnxruntime::RunOptions& run_options,
|
|||
}
|
||||
#endif
|
||||
ORT_CHECK_AND_SET_RETVAL(utils::ExecutePartialGraph(*session_state_, feeds_fetches_manager, feeds, fetches,
|
||||
run_logger, state, cache));
|
||||
run_logger, state, cache, partial_graph_index));
|
||||
}
|
||||
ORT_CATCH(const std::exception& e) {
|
||||
ORT_HANDLE_EXCEPTION([&]() {
|
||||
|
|
|
|||
|
|
@ -349,13 +349,15 @@ class InferenceSession {
|
|||
* copy/checks.
|
||||
* @param cache Contains node arg name to OrtValue map stashed from previous run
|
||||
* for frontier tensors
|
||||
* @param partial_graph_index Index of the partial graph to run.
|
||||
*/
|
||||
common::Status PartialRun(onnxruntime::RunOptions& run_options,
|
||||
const std::vector<OrtValue>& feeds,
|
||||
std::vector<OrtValue>& fetches,
|
||||
PartialGraphExecutionState& state,
|
||||
FeedsFetchesManager& feeds_fetches_manager,
|
||||
const OrtValueCachePtr& cache);
|
||||
const OrtValueCachePtr& cache,
|
||||
int32_t partial_graph_index);
|
||||
#endif
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -12,7 +12,12 @@ TrainingAgent::TrainingAgent(InferenceSession& session,
|
|||
const std::vector<std::string>& fw_feed_names,
|
||||
const std::vector<OrtDevice>& fw_outputs_device_info,
|
||||
const std::vector<std::string>& bw_fetches_names,
|
||||
const std::vector<OrtDevice>& bw_outputs_device_info) : inference_session_(session) {
|
||||
const std::vector<OrtDevice>& bw_outputs_device_info,
|
||||
int local_rank) : inference_session_(session) {
|
||||
ORT_UNUSED_PARAMETER(local_rank);
|
||||
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
|
||||
MemoryInfo::SetLocalRank(local_rank);
|
||||
#endif
|
||||
auto& session_state = session.GetSessionState();
|
||||
std::vector<std::string> fw_fetches_names;
|
||||
std::vector<std::string> bw_feed_names;
|
||||
|
|
@ -51,25 +56,34 @@ 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_);
|
||||
profile_step_ += 1;
|
||||
#endif
|
||||
|
||||
state.SetProgramCounterStart(0);
|
||||
state.SetProgramCounterEnd(fw_program_counter_end_);
|
||||
return RunCore(feeds, fetches, state, *fw_feeds_fetches_manager_, cache);
|
||||
|
||||
const int32_t partial_graph_index = 0;
|
||||
return RunCore(feeds, fetches, state, *fw_feeds_fetches_manager_, cache, partial_graph_index);
|
||||
}
|
||||
|
||||
common::Status TrainingAgent::RunBackward(const std::vector<OrtValue>& feeds, std::vector<OrtValue>& fetches,
|
||||
PartialGraphExecutionState& state) {
|
||||
state.SetProgramCounterStart(fw_program_counter_end_);
|
||||
state.SetProgramCounterEnd(bw_program_counter_end_);
|
||||
return RunCore(feeds, fetches, state, *bw_feeds_fetches_manager_, nullptr);
|
||||
const int32_t partial_graph_index = 1;
|
||||
return RunCore(feeds, fetches, state, *bw_feeds_fetches_manager_, nullptr, partial_graph_index);
|
||||
}
|
||||
|
||||
common::Status TrainingAgent::RunCore(const std::vector<OrtValue>& feeds, std::vector<OrtValue>& fetches,
|
||||
PartialGraphExecutionState& state, FeedsFetchesManager& feeds_fetches_manager,
|
||||
const OrtValueCachePtr& cache) {
|
||||
const OrtValueCachePtr& cache, int32_t partial_graph_index) {
|
||||
auto fetches_size = feeds_fetches_manager.GetFeedsFetchesInfo().output_names.size();
|
||||
fetches.resize(fetches_size, {});
|
||||
RunOptions run_options;
|
||||
return inference_session_.PartialRun(run_options, feeds, fetches, state, feeds_fetches_manager, cache);
|
||||
return inference_session_.PartialRun(run_options, feeds, fetches, state, feeds_fetches_manager, cache,
|
||||
partial_graph_index);
|
||||
}
|
||||
|
||||
void TrainingAgent::CreateAndInitializeFeedsFetchesManager(const SessionState& session_state,
|
||||
|
|
|
|||
|
|
@ -21,7 +21,8 @@ class TrainingAgent {
|
|||
const std::vector<std::string>& fw_feed_names,
|
||||
const std::vector<OrtDevice>& fw_outputs_device_info,
|
||||
const std::vector<std::string>& bw_fetches_names,
|
||||
const std::vector<OrtDevice>& bw_outputs_device_info);
|
||||
const std::vector<OrtDevice>& bw_outputs_device_info,
|
||||
int local_rank = 0);
|
||||
~TrainingAgent();
|
||||
// For ORTModule.forward()
|
||||
common::Status RunForward(const std::vector<OrtValue>& feeds, std::vector<OrtValue>& fetches,
|
||||
|
|
@ -33,7 +34,7 @@ class TrainingAgent {
|
|||
|
||||
common::Status RunCore(const std::vector<OrtValue>& feeds, std::vector<OrtValue>& fetches,
|
||||
PartialGraphExecutionState& state, FeedsFetchesManager& feeds_fetches_manager,
|
||||
const OrtValueCachePtr& cache)
|
||||
const OrtValueCachePtr& cache, int32_t partial_graph_index)
|
||||
ORT_MUST_USE_RESULT;
|
||||
|
||||
void CreateAndInitializeFeedsFetchesManager(const SessionState& session_state,
|
||||
|
|
@ -49,6 +50,10 @@ class TrainingAgent {
|
|||
std::unique_ptr<FeedsFetchesManager> bw_feeds_fetches_manager_;
|
||||
size_t fw_program_counter_end_;
|
||||
size_t bw_program_counter_end_;
|
||||
|
||||
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
|
||||
size_t profile_step_{0};
|
||||
#endif
|
||||
};
|
||||
|
||||
} // namespace training
|
||||
|
|
|
|||
|
|
@ -137,7 +137,7 @@ struct PyGradientGraphBuilder {
|
|||
// TODO: this method does not handle parallel optimization.
|
||||
TrainingConfigurationResult ConfigureSessionForTraining(
|
||||
training::PipelineTrainingSession* sess, TrainingParameters& parameters) {
|
||||
//TODO tix, refactor the mpi related code to populate all fields correctly by default.
|
||||
// TODO tix, refactor the mpi related code to populate all fields correctly by default.
|
||||
ORT_ENFORCE(parameters.data_parallel_size <= parameters.world_size, "data_parallel_size: ", parameters.data_parallel_size, ", world_size: ", parameters.world_size);
|
||||
ORT_ENFORCE(parameters.horizontal_parallel_size <= parameters.world_size, "horizontal_parallel_size: ", parameters.horizontal_parallel_size, ", world_size: ", parameters.world_size);
|
||||
ORT_ENFORCE(parameters.pipeline_parallel_size <= parameters.world_size, "pipeline_parallel_size: ", parameters.pipeline_parallel_size, ", world_size: ", parameters.world_size);
|
||||
|
|
@ -552,7 +552,7 @@ void addObjectMethodsForTraining(py::module& m, ExecutionProviderRegistrationFn
|
|||
NameMLValMap state_tensors;
|
||||
ORT_THROW_IF_ERROR(static_cast<PipelineTrainingSession*>(sess->GetSessionHandle())->GetStateTensors(state_tensors));
|
||||
auto& data_transfer_manager = sess->GetSessionHandle()->GetDataTransferManager();
|
||||
//convert to numpy array
|
||||
// convert to numpy array
|
||||
std::map<std::string, py::object> rmap;
|
||||
for (auto& kv : state_tensors) {
|
||||
if (kv.second.IsTensor()) {
|
||||
|
|
@ -610,9 +610,10 @@ void addObjectMethodsForTraining(py::module& m, ExecutionProviderRegistrationFn
|
|||
.def(py::init([](PyInferenceSession* session, const std::vector<std::string>& fw_feed_names,
|
||||
const std::vector<OrtDevice>& fw_outputs_device_info,
|
||||
const std::vector<std::string>& bw_fetches_names,
|
||||
const std::vector<OrtDevice>& bw_outputs_device_info) {
|
||||
const std::vector<OrtDevice>& bw_outputs_device_info,
|
||||
int local_rank) {
|
||||
return std::make_unique<TrainingAgent>(*session->GetSessionHandle(), fw_feed_names, fw_outputs_device_info,
|
||||
bw_fetches_names, bw_outputs_device_info);
|
||||
bw_fetches_names, bw_outputs_device_info, local_rank);
|
||||
}))
|
||||
.def("run_forward", [](TrainingAgent* agent, const std::vector<OrtValue>& feeds, std::vector<OrtValue>& fetches, PartialGraphExecutionState* state, OrtValueCachePtr cache) -> void {
|
||||
Status status = agent->RunForward(feeds, fetches, *state, cache);
|
||||
|
|
@ -742,28 +743,28 @@ void addObjectMethodsForTraining(py::module& m, ExecutionProviderRegistrationFn
|
|||
const std::unordered_set<std::string>& y_node_arg_names,
|
||||
const std::unordered_set<std::string>& x_node_arg_names,
|
||||
const std::string loss_node_arg_name) {
|
||||
std::shared_ptr<Model> model;
|
||||
auto logger_ptr = std::make_unique<logging::Logger>(logging::LoggingManager::DefaultLogger());
|
||||
logger_ptr->SetSeverity(logging::Severity::kINFO);
|
||||
ONNX_NAMESPACE::ModelProto model_proto;
|
||||
std::istringstream model_istream(serialized_model);
|
||||
ORT_THROW_IF_ERROR(Model::Load(model_istream, &model_proto));
|
||||
ORT_THROW_IF_ERROR(Model::Load(model_proto, model, nullptr, *logger_ptr));
|
||||
GradientGraphConfiguration gradient_graph_config{};
|
||||
gradient_graph_config.set_gradients_as_graph_outputs = true;
|
||||
// Save some objects, otherwise they get lost.
|
||||
auto gradient_graph_config_ptr = std::make_unique<GradientGraphConfiguration>(gradient_graph_config);
|
||||
std::shared_ptr<Model> model;
|
||||
auto logger_ptr = std::make_unique<logging::Logger>(logging::LoggingManager::DefaultLogger());
|
||||
logger_ptr->SetSeverity(logging::Severity::kINFO);
|
||||
ONNX_NAMESPACE::ModelProto model_proto;
|
||||
std::istringstream model_istream(serialized_model);
|
||||
ORT_THROW_IF_ERROR(Model::Load(model_istream, &model_proto));
|
||||
ORT_THROW_IF_ERROR(Model::Load(model_proto, model, nullptr, *logger_ptr));
|
||||
GradientGraphConfiguration gradient_graph_config{};
|
||||
gradient_graph_config.set_gradients_as_graph_outputs = true;
|
||||
// Save some objects, otherwise they get lost.
|
||||
auto gradient_graph_config_ptr = std::make_unique<GradientGraphConfiguration>(gradient_graph_config);
|
||||
|
||||
auto builder = std::make_unique<GradientGraphBuilder>(
|
||||
&model->MainGraph(),
|
||||
y_node_arg_names,
|
||||
x_node_arg_names,
|
||||
loss_node_arg_name,
|
||||
*gradient_graph_config_ptr,
|
||||
*logger_ptr);
|
||||
auto builder = std::make_unique<GradientGraphBuilder>(
|
||||
&model->MainGraph(),
|
||||
y_node_arg_names,
|
||||
x_node_arg_names,
|
||||
loss_node_arg_name,
|
||||
*gradient_graph_config_ptr,
|
||||
*logger_ptr);
|
||||
|
||||
return std::make_unique<PyGradientGraphBuilder>(std::move(builder), std::move(model), std::move(logger_ptr), std::move(gradient_graph_config_ptr));
|
||||
}))
|
||||
return std::make_unique<PyGradientGraphBuilder>(std::move(builder), std::move(model), std::move(logger_ptr), std::move(gradient_graph_config_ptr));
|
||||
}))
|
||||
.def("build", [](PyGradientGraphBuilder* gradient_graph_builder) {
|
||||
ORT_THROW_IF_ERROR(gradient_graph_builder->builder->Build());
|
||||
})
|
||||
|
|
|
|||
|
|
@ -5,8 +5,8 @@
|
|||
|
||||
import onnxruntime
|
||||
from onnxruntime.capi import _pybind_state as C
|
||||
from onnxruntime.capi.onnxruntime_inference_collection import IOBinding, OrtValue
|
||||
from onnxruntime.capi._pybind_state import TrainingAgent as C_TrainingAgent
|
||||
from onnxruntime.capi.onnxruntime_inference_collection import IOBinding, OrtValue
|
||||
|
||||
|
||||
class ExecutionAgentOutput: # pylint: disable=R0903
|
||||
|
|
@ -96,6 +96,7 @@ class TrainingAgent(object):
|
|||
session_options=None,
|
||||
providers=None,
|
||||
provider_options=None,
|
||||
local_rank=None,
|
||||
):
|
||||
"""
|
||||
:param path_or_bytes: filename or serialized ONNX or ORT format model in a byte string
|
||||
|
|
@ -110,6 +111,8 @@ class TrainingAgent(object):
|
|||
providers are used with the default precedence.
|
||||
:param provider_options: Optional sequence of options dicts corresponding
|
||||
to the providers listed in 'providers'.
|
||||
:param local_rank: Optional rank of current device, used for memory profiling only.
|
||||
Default rank is 0 if not specified.
|
||||
|
||||
The model type will be inferred unless explicitly set in the SessionOptions.
|
||||
To explicitly set:
|
||||
|
|
@ -137,6 +140,7 @@ class TrainingAgent(object):
|
|||
fw_outputs_device_info,
|
||||
bw_fetches_names,
|
||||
bw_outputs_device_info,
|
||||
local_rank,
|
||||
)
|
||||
|
||||
def run_forward(self, feeds, fetches, state, cache=None):
|
||||
|
|
|
|||
|
|
@ -3,17 +3,18 @@
|
|||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
from . import _utils, _io, _logger, _are_deterministic_algorithms_enabled, _use_deterministic_algorithms
|
||||
from ._graph_execution_manager import GraphExecutionManager, _RunStateInfo, _SkipCheck
|
||||
from ._execution_agent import TrainingAgent
|
||||
from .debug_options import DebugOptions
|
||||
from ._fallback import ORTModuleFallbackException, _FallbackPolicy, _FallbackManager
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
|
||||
from onnxruntime.capi import _pybind_state as C
|
||||
from onnxruntime.capi.onnxruntime_inference_collection import get_ort_device_type
|
||||
|
||||
import torch
|
||||
import warnings
|
||||
from . import _are_deterministic_algorithms_enabled, _io, _logger, _use_deterministic_algorithms, _utils
|
||||
from ._execution_agent import TrainingAgent
|
||||
from ._fallback import ORTModuleFallbackException, _FallbackManager, _FallbackPolicy
|
||||
from ._graph_execution_manager import GraphExecutionManager, _RunStateInfo, _SkipCheck
|
||||
from .debug_options import DebugOptions
|
||||
|
||||
|
||||
class TrainingManager(GraphExecutionManager):
|
||||
|
|
@ -322,6 +323,7 @@ class TrainingManager(GraphExecutionManager):
|
|||
)
|
||||
] * len(bw_fetches_names)
|
||||
|
||||
local_device_rank = self._device.index if device_type == "ort" else _utils.get_device_index(self._device)
|
||||
self._execution_agent = TrainingAgent(
|
||||
self._onnx_models.optimized_model.SerializeToString(),
|
||||
fw_feed_names,
|
||||
|
|
@ -331,6 +333,7 @@ class TrainingManager(GraphExecutionManager):
|
|||
session_options,
|
||||
providers,
|
||||
provider_options,
|
||||
local_device_rank,
|
||||
)
|
||||
|
||||
def _reinitialize_graph_builder(self, input_info):
|
||||
|
|
|
|||
Loading…
Reference in a new issue