diff --git a/onnxruntime/core/framework/execution_frame.cc b/onnxruntime/core/framework/execution_frame.cc
index 4df1c8335a..0dcae9a1c0 100644
--- a/onnxruntime/core/framework/execution_frame.cc
+++ b/onnxruntime/core/framework/execution_frame.cc
@@ -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();
diff --git a/onnxruntime/core/framework/memory_info.cc b/onnxruntime/core/framework/memory_info.cc
index e2464e1373..ebdef0e2fe 100644
--- a/onnxruntime/core/framework/memory_info.cc
+++ b/onnxruntime/core/framework/memory_info.cc
@@ -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;
diff --git a/onnxruntime/core/framework/memory_info.h b/onnxruntime/core/framework/memory_info.h
index 348cb291a8..1369371506 100644
--- a/onnxruntime/core/framework/memory_info.h
+++ b/onnxruntime/core/framework/memory_info.h
@@ -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
diff --git a/onnxruntime/core/framework/orttraining_partial_executor.cc b/onnxruntime/core/framework/orttraining_partial_executor.cc
index 5756ddfca0..f3ba26b2be 100644
--- a/onnxruntime/core/framework/orttraining_partial_executor.cc
+++ b/onnxruntime/core/framework/orttraining_partial_executor.cc
@@ -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
 
diff --git a/onnxruntime/core/framework/sequential_executor.cc b/onnxruntime/core/framework/sequential_executor.cc
index 2afe0596c9..0a24da2d72 100644
--- a/onnxruntime/core/framework/sequential_executor.cc
+++ b/onnxruntime/core/framework/sequential_executor.cc
@@ -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()) {
diff --git a/onnxruntime/core/framework/session_state.cc b/onnxruntime/core/framework/session_state.cc
index 2352a74883..0fc7616d65 100644
--- a/onnxruntime/core/framework/session_state.cc
+++ b/onnxruntime/core/framework/session_state.cc
@@ -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
diff --git a/onnxruntime/core/framework/session_state.h b/onnxruntime/core/framework/session_state.h
index 547f76b642..83b9407005 100644
--- a/onnxruntime/core/framework/session_state.h
+++ b/onnxruntime/core/framework/session_state.h
@@ -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_;
 
diff --git a/onnxruntime/core/framework/session_state_utils.cc b/onnxruntime/core/framework/session_state_utils.cc
index db88d79c65..d305d8463d 100644
--- a/onnxruntime/core/framework/session_state_utils.cc
+++ b/onnxruntime/core/framework/session_state_utils.cc
@@ -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 "
diff --git a/onnxruntime/core/framework/session_state_utils.h b/onnxruntime/core/framework/session_state_utils.h
index fde53dc682..a30dc9585c 100644
--- a/onnxruntime/core/framework/session_state_utils.h
+++ b/onnxruntime/core/framework/session_state_utils.h
@@ -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);
diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc
index 4f3d92c630..87b4e30a1e 100644
--- a/onnxruntime/core/session/inference_session.cc
+++ b/onnxruntime/core/session/inference_session.cc
@@ -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.
diff --git a/onnxruntime/core/session/inference_session.h b/onnxruntime/core/session/inference_session.h
index b1cb1a6841..a41e1c38fd 100644
--- a/onnxruntime/core/session/inference_session.h
+++ b/onnxruntime/core/session/inference_session.h
@@ -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_;
diff --git a/orttraining/orttraining/core/agent/training_agent.cc b/orttraining/orttraining/core/agent/training_agent.cc
index a9f87e83a4..055c78929f 100644
--- a/orttraining/orttraining/core/agent/training_agent.cc
+++ b/orttraining/orttraining/core/agent/training_agent.cc
@@ -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
 
diff --git a/orttraining/orttraining/core/session/training_session.cc b/orttraining/orttraining/core/session/training_session.cc
index 60b278ac7c..84c8b5e9f2 100644
--- a/orttraining/orttraining/core/session/training_session.cc
+++ b/orttraining/orttraining/core/session/training_session.cc
@@ -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
diff --git a/orttraining/orttraining/models/runner/training_runner.cc b/orttraining/orttraining/models/runner/training_runner.cc
index b97ffc6f41..ef875a5f1c 100644
--- a/orttraining/orttraining/models/runner/training_runner.cc
+++ b/orttraining/orttraining/models/runner/training_runner.cc
@@ -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;