[UIL] Enable flagcxBackend to receive custom backend options (flagos-ai#309)

Author: mikethegoblin

plugin/torch/example/example.py

@@ -27,6 +27,7 @@
 
 FLAGCX_GROUP1 = None
 FLAGCX_GROUP2 = None
+FLAGCX_GROUP3 = None
 MY_RANK = None
 WORLD_SIZE = None
 PREV_RANK = None
@@ -50,7 +51,7 @@ def get_args():
     return parser.parse_args()
 
 def init_pg():
-    global FLAGCX_GROUP1, FLAGCX_GROUP2, MY_RANK, WORLD_SIZE, PREV_RANK, NEXT_RANK
+    global FLAGCX_GROUP1, FLAGCX_GROUP2, FLAGCX_GROUP3, MY_RANK, WORLD_SIZE, PREV_RANK, NEXT_RANK
 
     # Get rank and world_size from environment
     MY_RANK = int(os.environ["RANK"])
@@ -66,6 +67,13 @@ def init_pg():
     FLAGCX_GROUP2 = dist.new_group(ranks=ranks, backend=f"cpu:gloo,{dev_name}:flagcx")
     print(f"ranks_flagcx: {dist.get_process_group_ranks(FLAGCX_GROUP1)}")
 
+    # Create a group with options; this only works when flagcxBackend has Options defined
+    # TODO: confirm with all vendors to see if their torch implementation support backend options
+    if flagcx._C is not None and hasattr(flagcx._C, 'ProcessGroupFlagCX') and hasattr(flagcx._C.ProcessGroupFlagCX, 'Options'):
+        flagcx_options = flagcx._C.ProcessGroupFlagCX.Options(enable_tuner=True)
+        FLAGCX_GROUP3 = dist.new_group(ranks=ranks, backend=f"{dev_name}:flagcx", pg_options=flagcx_options)
+        print(f"ranks_flagcx with options: {dist.get_process_group_ranks(FLAGCX_GROUP3)}")
+
     # Get prev_rank and next_rank
     PREV_RANK = (MY_RANK - 1 + WORLD_SIZE) % WORLD_SIZE
     NEXT_RANK = (MY_RANK + 1) % WORLD_SIZE
@@ -77,6 +85,7 @@ def init_pg():
 
 def destroy_pg():
     dist.destroy_process_group()
+    
 
 def test_broadcast():
     if torch.cuda.is_available():
@@ -351,3 +360,4 @@ def test_all():
     dict_op_to_test.get(args.op, test_all)()
 
     destroy_pg()
+    

plugin/torch/flagcx/include/backend_flagcx.hpp

@@ -19,6 +19,7 @@
 
 #include "event_flagcx.hpp"
 #include "stream_guard_flagcx.hpp"
+#include "utils_flagcx.hpp"
 
 namespace c10d {
 
@@ -74,8 +75,38 @@ class flagcxWork : public Work {
 
 class flagcxBackend : public Backend {
 public:
+// TODO: check with all vendors to make sure their torch implementation support
+// backend options
+#if defined(USE_NVIDIA_ADAPTOR) || defined(USE_METAX_ADAPTOR)
+  struct TuneObjectKey {
+    std::string commOp;
+    size_t nBytes;
+
+    bool operator<(const TuneObjectKey &other) const noexcept {
+      if (other.commOp == commOp) {
+        return nBytes < other.nBytes;
+      }
+      return commOp < other.commOp;
+    }
+  };
+
+  struct Options : Backend::Options {
+    explicit Options(bool enableTuner = false);
+
+    static c10::intrusive_ptr<Options> create(bool enableTuner = false) {
+      return c10::make_intrusive<Options>(enableTuner);
+    }
+
+    bool enableTuner{false};
+  };
+
+  explicit flagcxBackend(
+      const c10::intrusive_ptr<::c10d::Store> &store, int rank = -1,
+      int size = -1, c10::intrusive_ptr<Options> options = Options::create());
+#else
   explicit flagcxBackend(const c10::intrusive_ptr<::c10d::Store> &store,
                          int rank = -1, int size = -1);
+#endif
 
   ~flagcxBackend() override;
 
@@ -167,9 +198,18 @@ class flagcxBackend : public Backend {
   c10::intrusive_ptr<Work> recvAnysource(std::vector<at::Tensor> &tensors,
                                          int tag) override;
 
+#if defined(USE_NVIDIA_ADAPTOR) || defined(USE_METAX_ADAPTOR)
+  void checkRecordingEnded();
+  void recordTuneObject(flagcxCommOp_t commOp, flagcxDataType_t dataType,
+                        size_t count);
+  static c10::intrusive_ptr<Backend> createFlagcxBackend(
+      c10d::DistributedBackendOptions backendOptions,
+      c10::intrusive_ptr<Options> extraOptions = Options::create());
+#else
   static c10::intrusive_ptr<Backend>
   createFlagcxBackend(const c10::intrusive_ptr<::c10d::Store> &store, int rank,
                       int size, const std::chrono::duration<float> &timeout);
+#endif
 
   static void flagcxBackendConstructor() __attribute__((constructor)) {
     std::string devName = "cuda";
@@ -193,8 +233,14 @@ class flagcxBackend : public Backend {
     py::object module = py::module::import("torch.distributed");
     py::object registerBackend =
         module.attr("Backend").attr("register_backend");
+#if defined(USE_NVIDIA_ADAPTOR) || defined(USE_METAX_ADAPTOR)
+    registerBackend("flagcx", py::cpp_function(createFlagcxBackend),
+                    py::arg("extended_api") = true,
+                    py::arg("devices") = py::make_tuple(devName));
+#else
     registerBackend("flagcx", py::cpp_function(createFlagcxBackend),
                     py::arg("devices") = py::make_tuple(devName));
+#endif
   }
 
 protected:
@@ -214,6 +260,15 @@ class flagcxBackend : public Backend {
   std::unordered_map<int, flagcxStream_t> flagcxStreams_;
   std::unordered_map<int, std::unique_ptr<flagcxEvent>> flagcxEvents_;
   flagcxHandlerGroup_t handler_ = nullptr;
+#if defined(USE_NVIDIA_ADAPTOR) || defined(USE_METAX_ADAPTOR)
+  const c10::intrusive_ptr<Options> options_;
+  std::set<TuneObjectKey> tuneObjectSet_;
+  // whether we finished recording tuning objects
+  // a tuning object is a (commOp, nBytes) pair
+  // we record the tuning objects that will occur in this communicator so that
+  // flagcxTuner knows which communicator it is tuning
+  bool recordingEnded = false;
+#endif
 #ifdef USE_ASCEND_ADAPTOR
   aclrtStream acl_stream;
 #endif