[kineto] populate src/dst rank for p2p (pytorch#130812)

Summary: as title populate src/dst rank (global rank) for p2p kernel Differential Revision: D59794535 Pull Request resolved: pytorch#130812 Approved by: https://github.com/aaronenyeshi
xuanzhang816 · Jul 25, 2024 · 89bdd9c · 89bdd9c
1 parent 1c58aac
commit 89bdd9c
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 3 deletions.
diff --git a/third_party/kineto b/third_party/kineto
diff --git a/torch/csrc/profiler/util.cpp b/torch/csrc/profiler/util.cpp
@@ -376,8 +376,8 @@ std::unordered_map<std::string, std::string> saveNcclMeta(
     return map;
   }
 
-  map.emplace(
-      kCommsName, fmt::format("\"{}\"", debugInfo->getCollectiveName()));
+  auto& collective_name = debugInfo->getCollectiveName();
+  map.emplace(kCommsName, fmt::format("\"{}\"", collective_name));
   map.emplace(
       kDtype, fmt::format("\"{}\"", c10::toString(debugInfo->getDType())));
   map.emplace(kInMsgNelems, std::to_string(debugInfo->getInMessageNelems()));
@@ -411,6 +411,16 @@ std::unordered_map<std::string, std::string> saveNcclMeta(
 
   auto rank = debugInfo->getRank();
   map.emplace(kRank, std::to_string(rank));
+  int nRanks = static_cast<int>(groupRanks.size());
+  if (collective_name == "send") {
+    if (rank >= 0 && rank < nRanks) {
+      map.emplace(kP2pDst, std::to_string(groupRanks[rank]));
+    }
+  } else if (collective_name == "recv") {
+    if (rank >= 0 && rank < nRanks) {
+      map.emplace(kP2pSrc, std::to_string(groupRanks[rank]));
+    }
+  }
 #endif // USE_DISTRIBUTED
   return map;
 }

diff --git a/torch/csrc/profiler/util.h b/torch/csrc/profiler/util.h
@@ -170,6 +170,8 @@ constexpr auto kProcessGroupName = "Process Group Name";
 constexpr auto kProcessGroupDesc = "Process Group Description";
 constexpr auto kGroupRanks = "Process Group Ranks";
 constexpr auto kRank = "Rank";
+constexpr auto kP2pSrc = "Src Rank";
+constexpr auto kP2pDst = "Dst Rank";
 #endif // USE_DISTRIBUTED
 
 } // namespace torch::profiler::impl