Add collective latency profiler (#1785)

* [LatencyProfiler] Initial commit * [LatencyProfiler] Add unit tests * [LatencyProfiler] add more * [LatencyProfiler] Pass unit tests * [LatencyProfiler] Add hooks to integrate with meta internal tools * [LatencyProfiler] Restore install.sh * [LatencyProfiler] Resolved comments 1. add proper license 2. use proper namespace * [LatencyProfiler] Add header
2025-07-30 14:59:28 -07:00
parent 4ce3df8d3a
commit 874cd657ef
15 changed files with 891 additions and 0 deletions
@@ -19,6 +19,7 @@
 #include "graph.h"
 #include "nvmlwrap.h"
 #include "profiler.h"
+#include "latency_profiler/CollTrace.h"
 #include "rccl_common.h"
 #include "recorder.h"

@@ -634,6 +635,7 @@ struct ncclComm {

  hipEvent_t doneEvent;
  hipStream_t lastStream;
+  std::unique_ptr<latency_profiler::CollTrace> ctrace;

 #ifdef ENABLE_COLLTRACE
  struct ncclCollTrace* collTrace;
@@ -0,0 +1,61 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+#pragma once
+
+#include <thread>
+#include <atomic>
+#include "CollTraceEvent.h"
+#include "CollTraceUtils.h"
+#include "EventQueue.h"
+
+struct ncclComm;
+
+namespace latency_profiler {
+
+struct CollStats;
+
+// CollTrace Workflow
+// 1) measure collective latency by adding cuda/hip event before and after
+// a kernel launch in RCCL.
+
+// 2) Started a separate worker thread to collect latency data into
+// its local ring buffer.
+
+// 3) Perform CPU level all gather to exchange latency data
+// between threads (200 to 300us per 100 latency data exchange)
+
+// 4) Report to scuba when results buffer is full or every few minutes
+
+class CollTrace {
+ public:
+  CollTrace(ncclComm* comm);
+  __attribute__((visibility("default"))) ~CollTrace();
+
+  void enqueueEvent(std::unique_ptr<CollTraceEvent> event);
+
+  std::unique_ptr<CollTraceEvent> createEvent(
+      CollTraceEvent::EventType type = CollTraceEvent::EventType::COMM);
+
+  void recordCurCollResult(int rank, float latencyMs);
+
+  void reportIfNeeded(bool checkInterval);
+
+ private:
+  EventQueue<CollTraceEvent> eventQueue_;
+  std::thread profilingWorkerThread_;
+  void* collTraceThreadFn(int cudaDev);
+  std::atomic<uint64_t> curCollId_{0};
+  std::unique_ptr<CollTraceEvent> curEvent_;
+  std::deque<std::unique_ptr<CollTraceInfo>> pastColls_;
+
+  ncclComm* comm_{nullptr};
+  std::string commHash_;
+  int rank_{-1};
+  std::deque<std::vector<CollStats>> stats_;
+  std::chrono::time_point<std::chrono::steady_clock> lastReportTime_;
+};
+} // namespace latency_profiler
@@ -0,0 +1,68 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+#pragma once
+
+#include <cuda_runtime.h>
+#include <functional>
+#include <memory>
+#include <queue>
+#include "checks.h"
+#include "CollTraceUtils.h"
+
+namespace latency_profiler {
+
+// CUDA event pointer w/ deleter
+struct CudaEventDeleter {
+  void operator()(cudaEvent_t e) {
+    // Ignore error at destroy
+    cudaEventDestroy(e);
+  }
+};
+using CudaEventPtr = std::unique_ptr<
+    std::pointer_traits<cudaEvent_t>::element_type,
+    CudaEventDeleter>;
+
+// a wrapper class for cuda event
+class CudaWaitEvent {
+ public:
+  CudaWaitEvent(CudaEventPtr e) : event_(std::move(e)) {}
+  ~CudaWaitEvent() {}
+
+  cudaEvent_t getCudaEvent() {
+    return event_.get();
+  }
+
+  std::shared_ptr<float> getElapsedTimeSinceEvent(CudaWaitEvent* start);
+  ncclResult_t waitEventFinish();
+
+ private:
+  CudaEventPtr event_;
+};
+
+// Event data structure
+struct CollTraceEvent {
+  enum class EventType { COMM, TERMINATE };
+
+  CollTraceInfo coll;
+  std::unique_ptr<CudaWaitEvent> start{nullptr};
+  std::unique_ptr<CudaWaitEvent> stop{nullptr};
+  EventType eventType = EventType::COMM;
+
+  CollTraceEvent(EventType type) : eventType(type) {}
+  CollTraceEvent() = default;
+
+  ~CollTraceEvent() {}
+
+  // CollTraceEvent is not copyable
+  CollTraceEvent(const CollTraceEvent&) = delete;
+  CollTraceEvent& operator=(const CollTraceEvent&) = delete;
+
+  // CollTraceEvent is movable
+  CollTraceEvent(CollTraceEvent&&) = default;
+  CollTraceEvent& operator=(CollTraceEvent&&) = default;
+};
+} // namespace latency_profiler
@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+#pragma once
+
+#include "CollTraceEvent.h"
+#include "comm.h"
+#include "CollTraceUtils.h"
+
+namespace latency_profiler {
+class CollTraceError : public std::runtime_error {
+ public:
+  explicit CollTraceError(const std::string& what) : std::runtime_error(what) {}
+};
+
+ncclResult_t collTraceInit(ncclComm* comm);
+
+ncclResult_t collTraceDestroy(ncclComm* comm);
+
+std::unique_ptr<CollTraceEvent> collTraceAquireEventBaseline(
+    ncclKernelPlan* plan,
+    cudaStream_t stream);
+
+ncclResult_t collTraceRecordStartEvent(
+    ncclComm* comm,
+    cudaStream_t launchStream,
+    CollTraceEvent* event);
+
+ncclResult_t collTraceRecordEndEvent(
+    ncclComm* comm,
+    ncclKernelPlan* plan,
+    cudaStream_t launchStream,
+    std::unique_ptr<CollTraceEvent> event);
+
+CollTraceInfo parseCollInfoFromCollTask(const ncclTaskColl& collTask);
+} // namespace latency_profiler
@@ -0,0 +1,49 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+#pragma once
+
+#include <deque>
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace latency_profiler {
+
+struct CollStats {
+  const int collId;
+  const int percent;
+  const float minLatencyUs;
+  const float maxLatencyUs;
+  const std::string opName;
+  const std::string dataType;
+  int64_t count;
+  CollStats(const int collId, const int percent, const float minLatencyUs, const float maxLatencyUs, const std::string& opName, const std::string& dataType, const int64_t count) : collId(collId), percent(percent), minLatencyUs(minLatencyUs), maxLatencyUs(maxLatencyUs), opName(opName), dataType(dataType), count(count) {
+
+  }
+};
+
+struct CollTraceInfo {
+  int64_t collId;
+  std::string opName;
+  std::string dataType;
+  int64_t count;
+  float latencyMs{-1};
+};
+
+void reportToFile(
+    const std::deque<std::vector<CollStats>>& stats,
+    const std::string& commHash);
+
+std::vector<CollStats> aggregateResults(
+    const std::deque<std::unique_ptr<CollTraceInfo>>& info,
+    const std::vector<float>& latencyAllGather,
+    int RANKS_PER_HOST,
+    int NCCL_COLLTRACE_RECORD_MAX);
+
+float getSizeMb(const std::string& dataType, int count);
+
+} // namespace latency_profiler
@@ -0,0 +1,46 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+#pragma once
+
+#include <condition_variable>
+#include <deque>
+
+namespace latency_profiler {
+
+// A multi-producer, single consumer queue.
+// This queue is designed to be used in scenarios where multiple producers
+// are pushing items into the queue, and a single consumer is waiting for
+// items to be available.
+
+template <class Element>
+class EventQueue {
+ private:
+  std::deque<std::unique_ptr<Element>> queue_;
+  std::condition_variable cv_;
+  mutable std::mutex mutex_;
+
+ public:
+  void push(std::unique_ptr<Element> item) {
+    {
+      std::lock_guard<std::mutex> lock(mutex_);
+      queue_.push_back(std::move(item));
+    }
+    cv_.notify_one();
+  }
+
+  std::unique_ptr<Element> waitPop() {
+    std::unique_lock<std::mutex> lock(mutex_);
+    if (queue_.empty()) {
+      cv_.wait(lock, [this] { return !queue_.empty(); });
+    }
+    std::unique_ptr<Element> item = std::move(queue_.front());
+    queue_.pop_front();
+
+    return item;
+  }
+};
+} // namespace latency_profiler