[c10d][fr] Decouple the core logic of FR with the entry and event type (#152585)

fduwjj · pytorchmergebot · commit 5c878d4b04be · 2025-05-07T06:21:33.000Z
We want to make FR generic enough so the first step is to make the FR a template struct so that most of common code logic can be reused. The reason for this is that CudaEvent does not inherit c10::Event and we just want to swap the event part so that for NCCL we use CudaEvent and for the rest of backends, we use c10::event. Differential Revision: [D74262695](https://our.internmc.facebook.com/intern/diff/D74262695) Pull Request resolved: #152585 Approved by: https://github.com/kwen2501, https://github.com/d4l3k
diff --git a/torch/csrc/distributed/c10d/FlightRecorder.cpp b/torch/csrc/distributed/c10d/FlightRecorder.cpp
@@ -174,7 +174,8 @@ void DebugInfoWriter::registerWriter(std::unique_ptr<DebugInfoWriter> writer) {
 // Note: `getTraceback` invokes `torch::symbolize`, which may need to acquire
 // the GIL. If you don't want to block the current thread or take the risk of a
 // GIL deadlock, you can use an asynchronous calling mechanism like std::async.
-std::string FlightRecorder::Entry::getTraceback() {
+template <typename EventType>
+std::string FlightRecorder<EventType>::Entry::getTraceback() {
   torch::CapturedTraceback* traceback = traceback_.get();
   torch::SymbolizedTracebacks s_tbs = torch::symbolize({traceback});
   // We use 0 because we only have one traceback here.
@@ -197,7 +198,8 @@ std::string FlightRecorder::Entry::getTraceback() {
   return oss.str();
 }
 
-std::optional<size_t> FlightRecorder::record(
+template <typename EventType>
+std::optional<size_t> FlightRecorder<EventType>::record(
     size_t pg_id,
     const std::tuple<std::string, std::string>& pg_name,
     size_t collective_seq_id,
@@ -206,8 +208,8 @@ std::optional<size_t> FlightRecorder::record(
     std::string profiling_name,
     const std::vector<at::Tensor>& inputs,
     const std::vector<at::Tensor>& outputs,
-    Event* start,
-    Event* end,
+    EventType* start,
+    EventType* end,
     std::chrono::milliseconds timeout_ms,
     std::shared_ptr<ProcessGroupStatus> pg_status,
     bool isP2P) {
@@ -271,7 +273,8 @@ std::optional<size_t> FlightRecorder::record(
   return id_++;
 }
 
-void FlightRecorder::record_pg_ranks(
+template <typename EventType>
+void FlightRecorder<EventType>::record_pg_ranks(
     const std::tuple<std::string, std::string>& pg_name,
     std::vector<uint64_t> ranks) {
   if (!enabled_) {
@@ -281,7 +284,8 @@ void FlightRecorder::record_pg_ranks(
   pg_name_to_ranks_[pg_name] = std::move(ranks);
 }
 
-void FlightRecorder::record_accelerator_version(
+template <typename EventType>
+void FlightRecorder<EventType>::record_accelerator_version(
     const std::string nccl_version) {
   if (!enabled_) {
     return;
@@ -290,7 +294,8 @@ void FlightRecorder::record_accelerator_version(
   nccl_version_ = std::move(nccl_version);
 }
 
-void FlightRecorder::update_state(Entry& r) {
+template <typename EventType>
+void FlightRecorder<EventType>::update_state(Entry& r) {
   if (r.start_ != nullptr) {
     bool started = r.start_->query();
     if (started && !r.time_discovered_started_) {
@@ -305,7 +310,9 @@ void FlightRecorder::update_state(Entry& r) {
   }
 }
 
-std::vector<FlightRecorder::Entry> FlightRecorder::dump_entries() {
+template <typename EventType>
+std::vector<typename FlightRecorder<EventType>::Entry> FlightRecorder<
+    EventType>::dump_entries() {
   std::lock_guard<std::mutex> guard(mutex_);
   std::vector<Entry> result;
   result.reserve(entries_.size());
@@ -325,10 +332,11 @@ std::vector<FlightRecorder::Entry> FlightRecorder::dump_entries() {
   return result;
 }
 
+template <typename EventType>
 // Returns the entry with the given id, if it exists. Otherwise, returns
 // std::nullopt.
-std::optional<FlightRecorder::Entry> FlightRecorder::getEntry(
-    std::optional<size_t> id) {
+std::optional<typename FlightRecorder<EventType>::Entry> FlightRecorder<
+    EventType>::getEntry(std::optional<size_t> id) {
   if (!enabled_ || !id) {
     return std::nullopt;
   }
@@ -342,16 +350,17 @@ std::optional<FlightRecorder::Entry> FlightRecorder::getEntry(
   }
 }
 
-void FlightRecorder::retire_id(
+template <typename EventType>
+void FlightRecorder<EventType>::retire_id(
     std::optional<size_t> id,
     bool compute_duration) {
   if (!enabled_ || !id) {
     return;
   }
 
   bool can_compute_duration = false;
-  Event* startEvent = nullptr;
-  Event* endEvent = nullptr;
+  EventType* startEvent = nullptr;
+  EventType* endEvent = nullptr;
   std::optional<float> duration = std::nullopt;
 
   std::unique_lock<std::mutex> guard(mutex_);
@@ -391,7 +400,8 @@ void FlightRecorder::retire_id(
   }
 }
 
-const c10::List<c10::IValue> FlightRecorder::getCollectiveTrace(
+template <typename EventType>
+const c10::List<c10::IValue> FlightRecorder<EventType>::getCollectiveTrace(
     bool includeStacktraces,
     bool onlyActive) {
   auto entries = new_list();
@@ -495,7 +505,9 @@ const c10::List<c10::IValue> FlightRecorder::getCollectiveTrace(
   return entries;
 }
 
-const c10::Dict<c10::IValue, c10::IValue> FlightRecorder::getPgConfig() {
+template <typename EventType>
+const c10::Dict<c10::IValue, c10::IValue> FlightRecorder<
+    EventType>::getPgConfig() {
   auto pg_config = new_dict();
   for (const auto& [pg_name, ranks] : pg_name_to_ranks_) {
     auto pg_info = new_dict();
@@ -507,8 +519,9 @@ const c10::Dict<c10::IValue, c10::IValue> FlightRecorder::getPgConfig() {
   return pg_config;
 }
 
-const std::map<std::string, std::map<std::string, std::string>> FlightRecorder::
-    getPgConfigJson() {
+template <typename EventType>
+const std::map<std::string, std::map<std::string, std::string>> FlightRecorder<
+    EventType>::getPgConfigJson() {
   std::map<std::string, std::map<std::string, std::string>> result;
   for (const auto& [pg_name, ranks] : pg_name_to_ranks_) {
     auto pg_info = std::map<std::string, std::string>();
@@ -520,7 +533,9 @@ const std::map<std::string, std::map<std::string, std::string>> FlightRecorder::
   return result;
 }
 
-const c10::Dict<c10::IValue, c10::IValue> FlightRecorder::getPgStatus() {
+template <typename EventType>
+const c10::Dict<c10::IValue, c10::IValue> FlightRecorder<
+    EventType>::getPgStatus() {
   auto all_pg_status = new_dict();
   for (const auto& [pg_id, status] : all_pg_status_) {
     auto pg_status = new_dict();
@@ -532,8 +547,9 @@ const c10::Dict<c10::IValue, c10::IValue> FlightRecorder::getPgStatus() {
   return all_pg_status;
 }
 
-const std::map<std::string, std::map<std::string, std::string>> FlightRecorder::
-    getPgStatusJson() {
+template <typename EventType>
+const std::map<std::string, std::map<std::string, std::string>> FlightRecorder<
+    EventType>::getPgStatusJson() {
   std::map<std::string, std::map<std::string, std::string>> result;
   for (const auto& [pg_id, status] : all_pg_status_) {
     auto pg_status = std::map<std::string, std::string>();
@@ -548,7 +564,8 @@ const std::map<std::string, std::map<std::string, std::string>> FlightRecorder::
   return result;
 }
 
-std::string FlightRecorder::dump_json(
+template <typename EventType>
+std::string FlightRecorder<EventType>::dump_json(
     const std::optional<std::unordered_map<
         std::string,
         std::unordered_map<std::string, std::string>>>& ncclDumpMap,
@@ -640,7 +657,8 @@ std::string FlightRecorder::dump_json(
   return result.dump();
 }
 
-std::string FlightRecorder::dump(
+template <typename EventType>
+std::string FlightRecorder<EventType>::dump(
     const std::optional<std::unordered_map<
         std::string,
         std::unordered_map<std::string, std::string>>>& ncclDumpMap,
@@ -677,6 +695,8 @@ std::string FlightRecorder::dump(
   return pickle_str(result);
 }
 
+template struct FlightRecorder<at::cuda::CUDAEvent>;
+
 std::unique_ptr<DebugInfoWriter> DebugInfoWriter::writer_ = nullptr;
 std::atomic<bool> DebugInfoWriter::hasWriterRegistered_(false);
 
diff --git a/torch/csrc/distributed/c10d/FlightRecorder.hpp b/torch/csrc/distributed/c10d/FlightRecorder.hpp
@@ -93,19 +93,20 @@ float getDurationFromEvent(
     at::cuda::CUDAEvent& ncclStartEvent,
     at::cuda::CUDAEvent& ncclEndEvent);
 
+template <typename EventType>
 struct FlightRecorder {
-  static FlightRecorder* get() {
+  static FlightRecorder<EventType>* get() {
     // intentionally leak on exit
     // because this will hold python state that may get destructed
-    static FlightRecorder* instance = new FlightRecorder();
+    static FlightRecorder<EventType>* instance =
+        new FlightRecorder<EventType>();
     return instance;
   }
   FlightRecorder() {
     max_entries_ = getCvarInt({"TORCH_NCCL_TRACE_BUFFER_SIZE"}, 0);
     capture_cpp_stack_ = getCvarBool({"TORCH_NCCL_TRACE_CPP_STACK"}, false);
     enabled_ = max_entries_ > 0;
   }
-  using Event = at::cuda::CUDAEvent;
   struct Entry {
     size_t id_; // incremented id in the trace buffer
                 // used to figure out where in the circular entries
@@ -129,7 +130,7 @@ struct FlightRecorder {
     // we borrow pointers to start_ and end_ so we can query the state
     // on reporting. However, once the event is completed, the call
     // to `complete` will clear these.
-    Event *start_, *end_;
+    EventType *start_, *end_;
 
     // timestamp when the entry was created, likely close to the time the work
     // was 'enqueued'- not necessarily started
@@ -164,6 +165,10 @@ struct FlightRecorder {
                            // a retired but not completed event has timed out
 
     // Returns the traceback of current entry, in string form.
+    // Note: `getTraceback` invokes `torch::symbolize`, which may need to
+    // acquire the GIL. If you don't want to block the current thread or take
+    // the risk of a GIL deadlock, you can use an asynchronous calling mechanism
+    // like std::async.
     std::string getTraceback();
   };
 
@@ -188,8 +193,8 @@ struct FlightRecorder {
       std::string profiling_name,
       const std::vector<at::Tensor>& inputs,
       const std::vector<at::Tensor>& outputs,
-      Event* start,
-      Event* end,
+      EventType* start,
+      EventType* end,
       std::chrono::milliseconds timeout_ms,
       std::shared_ptr<ProcessGroupStatus> pg_status,
       bool isP2P);
diff --git a/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp b/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp
@@ -37,6 +37,7 @@
 namespace c10d {
 
 constexpr const char* const kNCCLAbortedCommStoreKey = "NCCLABORTEDCOMM";
+using FlightRecorderCUDA = FlightRecorder<at::cuda::CUDAEvent>;
 
 namespace {
 
@@ -439,13 +440,13 @@ std::string dump_nccl_trace(
     printNcclCommProxyTrace("Received dump signal " + ncclUniqueIDStr, dump);
   }
 #endif // defined(USE_ROCM) && defined(NCCL_COMM_DUMP)
-  return FlightRecorder::get()->dump(
+  return FlightRecorderCUDA::get()->dump(
       ncclDumpMap, includeCollectives, includeStackTraces, onlyActive);
 }
 
 std::string dump_nccl_trace_json(bool includeCollectives, bool onlyActive) {
   auto ncclDumpMap = getNCCLCommDumpMap();
-  return FlightRecorder::get()->dump_json(
+  return FlightRecorderCUDA::get()->dump_json(
       ncclDumpMap, includeCollectives, onlyActive);
 }
 
@@ -740,8 +741,8 @@ bool ProcessGroupNCCL::WorkNCCL::checkTimeout(
 void ProcessGroupNCCL::WorkNCCL::printTraceback() const {
   // First step we get the corresponding record entry from FR, based on work's
   // trace_id_
-  std::optional<FlightRecorder::Entry> entry =
-      FlightRecorder::get()->getEntry(trace_id_);
+  std::optional<FlightRecorderCUDA::Entry> entry =
+      FlightRecorderCUDA::get()->getEntry(trace_id_);
   if (entry.has_value()) {
     auto entryVal = entry.value();
     // Get stack trace from FR entry, in string format
@@ -2475,7 +2476,7 @@ void ProcessGroupNCCL::watchdogHandler() {
         pgStatus_->lastCompletedWorkName = opTypeToString(work.opType_);
         pgStatus_->lastCompletedNumelIn = work.numelIn_;
         pgStatus_->lastCompletedNumelOut = work.numelOut_;
-        FlightRecorder::get()->retire_id(work.trace_id_, true);
+        FlightRecorderCUDA::get()->retire_id(work.trace_id_, true);
         if (onCompletionHook_) {
           // Move Work object to completedWorkList_ to be consumed by the hook
           // thread
@@ -3007,9 +3008,9 @@ std::shared_ptr<NCCLComm> ProcessGroupNCCL::initNCCLComm(
     inInitializationCommMap_.emplace(deviceKey, ncclComm);
   }
 
-  FlightRecorder::get()->record_pg_ranks(
+  FlightRecorderCUDA::get()->record_pg_ranks(
       std::make_tuple(pg_uid_, pg_desc_), groupRanks());
-  FlightRecorder::get()->record_accelerator_version(getNcclVersion());
+  FlightRecorderCUDA::get()->record_accelerator_version(getNcclVersion());
 
   VLOG(2) << logPrefix() << "ProcessGroupNCCL created ncclComm_ "
           << ncclComm->repr()
@@ -3223,7 +3224,7 @@ c10::intrusive_ptr<ProcessGroupNCCL::WorkNCCL> ProcessGroupNCCL::initWork(
     //   these objects to the Work becuase it has implications for keeping those
     //   tensors alive longer and adds overhead when copying Work objects
     //   between threads
-    r->trace_id_ = FlightRecorder::get()->record(
+    r->trace_id_ = FlightRecorderCUDA::get()->record(
         local_id_,
         std::make_tuple(pg_uid_, pg_desc_),
         seqCollective_,
@@ -3510,7 +3511,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::collective(
     // later in endCoalescing we record a 'coalesced' Work which has
     // timing/state updates via watchdog thread, but lacks op metadata such as
     // input/output sizes and profilingTitle per-op in the group.
-    FlightRecorder::get()->record(
+    FlightRecorderCUDA::get()->record(
         local_id_,
         std::make_tuple(pg_uid_, pg_desc_),
         seqCollective_,
@@ -3908,7 +3909,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::pointToPoint(
     // later in endCoalescing we record a 'coalesced' Work which has
     // timing/state updates via watchdog thread, but lacks op metadata such as
     // input/output sizes and profilingTitle per-op in the group.
-    FlightRecorder::get()->record(
+    FlightRecorderCUDA::get()->record(
         local_id_,
         std::make_tuple(pg_uid_, pg_desc_),
         seqCollective_,
@@ -3949,7 +3950,7 @@ c10::intrusive_ptr<Work> ProcessGroupNCCL::pointToPoint(
     // TODO(whc) because we don't pass output {tensor} to initWork, we tell
     // initWork to not record, and then we manually call record passing all the
     // information it wants.
-    work->trace_id_ = FlightRecorder::get()->record(
+    work->trace_id_ = FlightRecorderCUDA::get()->record(
         local_id_,
         std::make_tuple(pg_uid_, pg_desc_),
         seqCollective_,
