diff --git a/onnxruntime/core/framework/memory_info.cc b/onnxruntime/core/framework/memory_info.cc
index ce69cd75eb..e2464e1373 100644
--- a/onnxruntime/core/framework/memory_info.cc
+++ b/onnxruntime/core/framework/memory_info.cc
@@ -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));
diff --git a/onnxruntime/core/framework/memory_info.h b/onnxruntime/core/framework/memory_info.h
index 0c40cd90e5..348cb291a8 100644
--- a/onnxruntime/core/framework/memory_info.h
+++ b/onnxruntime/core/framework/memory_info.h
@@ -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_;
 };
 
diff --git a/onnxruntime/core/framework/orttraining_partial_executor.cc b/onnxruntime/core/framework/orttraining_partial_executor.cc
index 40d5e2a15f..9650c99532 100644
--- a/onnxruntime/core/framework/orttraining_partial_executor.cc
+++ b/onnxruntime/core/framework/orttraining_partial_executor.cc
@@ -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()) {
diff --git a/onnxruntime/core/framework/orttraining_partial_executor.h b/onnxruntime/core/framework/orttraining_partial_executor.h
index 0a17a0937b..b34c4e67f2 100644
--- a/onnxruntime/core/framework/orttraining_partial_executor.h
+++ b/onnxruntime/core/framework/orttraining_partial_executor.h
@@ -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
\ No newline at end of file
+#endif
diff --git a/onnxruntime/core/framework/utils.cc b/onnxruntime/core/framework/utils.cc
index 1287e879ce..19d0f3431d 100644
--- a/onnxruntime/core/framework/utils.cc
+++ b/onnxruntime/core/framework/utils.cc
@@ -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();
 
diff --git a/onnxruntime/core/framework/utils.h b/onnxruntime/core/framework/utils.h
index c9bd4cedab..f320c75f34 100644
--- a/onnxruntime/core/framework/utils.h
+++ b/onnxruntime/core/framework/utils.h
@@ -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.
diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc
index 10078eaf7c..1d5a279f95 100644
--- a/onnxruntime/core/session/inference_session.cc
+++ b/onnxruntime/core/session/inference_session.cc
@@ -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([&]() {
diff --git a/onnxruntime/core/session/inference_session.h b/onnxruntime/core/session/inference_session.h
index 3539f5a4ca..963eeddb3b 100644
--- a/onnxruntime/core/session/inference_session.h
+++ b/onnxruntime/core/session/inference_session.h
@@ -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
 
   /**
diff --git a/orttraining/orttraining/core/agent/training_agent.cc b/orttraining/orttraining/core/agent/training_agent.cc
index 107795c944..a9f87e83a4 100644
--- a/orttraining/orttraining/core/agent/training_agent.cc
+++ b/orttraining/orttraining/core/agent/training_agent.cc
@@ -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,
diff --git a/orttraining/orttraining/core/agent/training_agent.h b/orttraining/orttraining/core/agent/training_agent.h
index 5645cc1059..31f396aeb0 100644
--- a/orttraining/orttraining/core/agent/training_agent.h
+++ b/orttraining/orttraining/core/agent/training_agent.h
@@ -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
diff --git a/orttraining/orttraining/python/orttraining_pybind_state.cc b/orttraining/orttraining/python/orttraining_pybind_state.cc
index 3defb5b7e3..08c90ee350 100644
--- a/orttraining/orttraining/python/orttraining_pybind_state.cc
+++ b/orttraining/orttraining/python/orttraining_pybind_state.cc
@@ -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());
       })
diff --git a/orttraining/orttraining/python/training/ortmodule/_execution_agent.py b/orttraining/orttraining/python/training/ortmodule/_execution_agent.py
index c3b623a327..be2cf01b1f 100644
--- a/orttraining/orttraining/python/training/ortmodule/_execution_agent.py
+++ b/orttraining/orttraining/python/training/ortmodule/_execution_agent.py
@@ -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):
diff --git a/orttraining/orttraining/python/training/ortmodule/_training_manager.py b/orttraining/orttraining/python/training/ortmodule/_training_manager.py
index 496eb584ef..8908c0968b 100644
--- a/orttraining/orttraining/python/training/ortmodule/_training_manager.py
+++ b/orttraining/orttraining/python/training/ortmodule/_training_manager.py
@@ -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):