Fixes for rocm 7.0 and dev

Author: alextmagro

examples/pytorch/comm_gemm_overlap/te_layer_with_overlap_profile.py

@@ -185,32 +185,29 @@ def _get_layer_args(config, tp_group, tp_size, reference=False):
 
     return args, kwargs, input_shape
 
-def create_ub_cfgs(config_file:str, tp_size: int = 8):
+def create_ub_cfgs(config_file: str, tp_size: int = 8):
     import json
     with open(config_file, 'r') as f:
         data = json.load(f)
     cfgs = {}
     _MIN_STREAM_PRIORITY, _MAX_STREAM_PRIORITY = None, None
     layers_reduce_scatter_overlap = ["proj_fprop", "fc2_fprop", "qkv_wgrad", "fc1_wgrad"]
-
-    for name, method in data.items():
-        is_reduce_scatter = name in layers_reduce_scatter_overlap
-
-        layers_all_gather_overlap = [
+    layers_all_gather_overlap = [
                 "qkv_fprop",
                 "qkv_dgrad",
                 "proj_dgrad",
                 "fc1_fprop",
                 "fc1_dgrad",
                 "fc2_dgrad",
             ]
+
+    for name, method in data.items():
         if _MIN_STREAM_PRIORITY is None or _MAX_STREAM_PRIORITY is None:
             _MIN_STREAM_PRIORITY, _MAX_STREAM_PRIORITY = tex.get_stream_priority_range()
 
-        
         cfg = {
             "method": method,
-            "is_reduce_scatter": is_reduce_scatter,
+            "is_reduce_scatter": name in layers_reduce_scatter_overlap,
             "num_sm": 1 if method == "ring_exchange" else 16,
             "cga_size": 1 if method == "ring_exchange" else 2,
             "set_sm_margin": False,

transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp

@@ -21,6 +21,12 @@
 #define HALF_BYTES 2
 #define UB_MAX_SM 32
 
+#ifdef __HIP_PLATFORM_AMD__
+#define half_dtype hip_bfloat16
+#define __nv_fp8_e5m2 te_hip_fp8_e5m2
+#define __nv_fp8_e4m3 te_hip_fp8_e4m3
+#endif
+
 using namespace std::placeholders;
 
 namespace transformer_engine {
@@ -448,7 +454,7 @@ void CommOverlapBase::split_overlap_rs(const TensorWrapper &A, bool transa, cons
                                        TensorWrapper &pre_gelu_out, TensorWrapper &workspace,
                                        bool grad, bool accumulate, bool use_split_accumulator,
                                        TensorWrapper &rs_output, cudaStream_t stream_main) {
-  printf("split_overlap_rs_pipeline");
+  printf("split_overlap_rs_pipeline\n");
   // Get GEMM dimensions
   int ori_sms = _ub_comm->sms;
   _ub_comm->use_ce = _use_ce;
@@ -596,7 +602,7 @@ CommOverlapP2PBase::CommOverlapP2PBase(const std::vector<size_t> &buffer_shape,
                                        CommOverlapType comm_type, int num_max_streams,
                                        int comm_cga_size, int gemm_priority, int comm_priority,
                                        int num_comm_sm, bool set_sm_margin, bool use_ce,
-                                       bool atomic_gemm, bool aggregate, bool use_rd = false)
+                                       bool atomic_gemm, bool aggregate, bool use_rd)
     : CommOverlapCore(myrank, numranks, mylocal, numlocal, mynode, numnodes, tp_size,
                       allgather_handle, barrier_handle, tp_size, num_max_streams, comm_cga_size,
                       gemm_priority, comm_priority, num_comm_sm, set_sm_margin, use_ce,
@@ -798,7 +804,7 @@ void CommOverlapP2PBase::split_overlap_ag(const TensorWrapper &A, bool transa,
                                           TensorWrapper &workspace, bool grad, bool accumulate,
                                           bool use_split_accumulator, TensorWrapper &B_copy,
                                           cudaStream_t stream_main) {
-  printf("split_overlap_ag");
+  printf("split_overlap_ag\n");
   int ori_sms = _ub_comm->sms;
   _ub_comm->use_ce = _use_ce;
   _ub_comm->sms = _num_comm_sm;
@@ -960,12 +966,12 @@ void CommOverlapP2PBase::split_overlap_ag(const TensorWrapper &A, bool transa,
 ** This function assumes the input_b is pre-copied to _ubufs[rank_id]. This is needed to have AG
 ** outputs in each rank to be in the contiguous memory space after all ring exchange phases.
 */
-void CommOverlapP2PBase::split_overlap_ag_rd(TensorWrapper &A, bool transa, TensorWrapper &B,
-                                          bool transb, TensorWrapper &D, TensorWrapper &bias,
-                                          TensorWrapper &pre_gelu_out, TensorWrapper &workspace,
-                                          bool grad, bool accumulate, bool use_split_accumulator,
-                                          TensorWrapper &B_copy, cudaStream_t stream_main) {
-  printf("split_overlap_ag_rd");
+void CommOverlapP2PBase::split_overlap_ag_rd(const TensorWrapper &A, bool transa, const TensorWrapper &B,
+                                bool transb, TensorWrapper &D, TensorWrapper &bias,
+                                TensorWrapper &pre_gelu_out, TensorWrapper &workspace, bool grad,
+                                bool accumulate, bool use_split_accumulator, TensorWrapper &B_copy,
+                                cudaStream_t stream_main) {
+  printf("split_overlap_ag_rd\n");
   int ori_sms = _ub_comm->sms;
   _ub_comm->use_ce = _use_ce;
   _ub_comm->sms = _num_comm_sm;
@@ -1025,12 +1031,12 @@ void CommOverlapP2PBase::split_overlap_ag_rd(TensorWrapper &A, bool transa, Tens
       // GEMM
       char *input_b_chunk_ptr = input_b_ptr + send_offset;
       auto input_b_chunk =
-          TensorWrapper(reinterpret_cast<void *>(input_b_chunk_ptr), {n_chunk * 2, k}, B.dtype(),
+          TensorWrapper(reinterpret_cast<void *>(input_b_chunk_ptr), std::vector<size_t>{n_chunk * 2, k}, B.dtype(),
                         nullptr, nullptr, B.scale_inv());
 
       char *output_chunk_ptr = output_ptr + (send_chunk_id * output_chunk_bytes);
       auto output_chunk = TensorWrapper(reinterpret_cast<void *>(output_chunk_ptr),
-                                        {n_chunk * 2, m}, D.dtype(), D.amax(), D.scale(), nullptr);
+                                        std::vector<size_t>{n_chunk * 2, m}, D.dtype(), D.amax(), D.scale(), nullptr);
 
       char *aux_chunk_ptr =
           (do_gelu) ? pre_gelu_out_ptr + (send_chunk_id * aux_chunk_bytes) : nullptr;
@@ -1084,12 +1090,12 @@ void CommOverlapP2PBase::split_overlap_ag_rd(TensorWrapper &A, bool transa, Tens
         cudaStream_t compute_stream = _stream_compute[chunk_id % _stream_compute.size()];
 
         auto input_b_chunk = TensorWrapper(_ubufs[chunk_id].dptr(),
-                                          {n_chunk, k}, B.dtype(),
+                                          std::vector<size_t>{n_chunk, k}, B.dtype(),
                                           nullptr, nullptr, B.scale_inv());
 
         char* output_chunk_ptr = output_ptr + (chunk_id * output_chunk_bytes);
         auto output_chunk = TensorWrapper(reinterpret_cast<void *>(output_chunk_ptr),
-                                          {n_chunk, m},
+                                          std::vector<size_t>{n_chunk, m},
                                           D.dtype(), D.amax(), D.scale(), nullptr);
 
         char *aux_chunk_ptr =
@@ -1140,12 +1146,12 @@ void CommOverlapP2PBase::split_overlap_ag_rd(TensorWrapper &A, bool transa, Tens
           cudaStream_t compute_stream = _stream_compute[new_chunk_id % _stream_compute.size()];
 
           auto input_b_chunk = TensorWrapper(_ubufs[new_chunk_id].dptr(),
-                                            {n_chunk, k}, B.dtype(),
+                                            std::vector<size_t>{n_chunk, k}, B.dtype(),
                                             nullptr, nullptr, B.scale_inv());
 
           char* output_chunk_ptr = output_ptr + (new_chunk_id * output_chunk_bytes);
           auto output_chunk = TensorWrapper(reinterpret_cast<void *>(output_chunk_ptr),
-                                            {n_chunk, m},
+                                            std::vector<size_t>{n_chunk, m},
                                             D.dtype(), D.amax(), D.scale(), nullptr);
 
           char *aux_chunk_ptr =
@@ -1271,7 +1277,7 @@ void CommOverlapP2PBase::split_overlap_rs(const TensorWrapper &A, bool transa,
                                           TensorWrapper &workspace, bool grad, bool accumulate,
                                           bool use_split_accumulator, TensorWrapper &rs_output,
                                           cudaStream_t stream_main) {
-  printf("split_overlap_rs_p2p");
+  printf("split_overlap_rs_p2p\n");
   int ori_sms = _ub_comm->sms;
   _ub_comm->use_ce = _use_ce;
   _ub_comm->sms = _num_comm_sm;

transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers.cu

@@ -8,16 +8,17 @@
 #include <cuda_fp8.h>
 #include <cuda_runtime.h>
 
-#if __CUDA_ARCH__ >= 800
-#define half_dtype nv_bfloat16
-#else
-#define half_dtype half
-#endif
 
 #ifdef __HIP_PLATFORM_AMD__
 #define half_dtype hip_bfloat16
 #define __nv_fp8_e5m2 te_hip_fp8_e5m2
 #define __nv_fp8_e4m3 te_hip_fp8_e4m3
+#else
+#if __CUDA_ARCH__ >= 800
+#define half_dtype nv_bfloat16
+#else
+#define half_dtype half
+#endif
 #endif
 
 #include <assert.h>
@@ -2094,7 +2095,8 @@ void allgather2_userbuff_inplace(const int handler, const int offset, const int
     }
   }
 #else
-  if (comm->use_mc && (comm->memflags[handler] & UB_MEM_MC_CREATED)) {
+  int threads = comm->threads;
+  if (comm->use_mc && (comm->memflags[handler] & NVTE_UB_MEM_MC_CREATED)) {
     callranks_agMC(2) callranks_agMC(4) callranks_agMC(8) callranks_agMC(16) callranks_agMC(32)
   } else {
     callranks_ag(2) callranks_ag(4) callranks_ag(8) callranks_ag(16) callranks_ag(32)
@@ -2150,7 +2152,7 @@ void reducescatter2_userbuff_inplace(const int handler, const int offset, const
   }
 #else
   int threads = comm->threads;
-  if (comm->use_mc && (comm->memflags[handler] & UB_MEM_MC_CREATED)) {
+  if (comm->use_mc && (comm->memflags[handler] & NVTE_UB_MEM_MC_CREATED)) {
     callranks_rsMC(2) callranks_rsMC(4) callranks_rsMC(8) callranks_rsMC(16) callranks_rsMC(32)
   } else {
     callranks_rs(2) callranks_rs(4) callranks_rs(8) callranks_rs(16) callranks_rs(32)
@@ -2666,6 +2668,7 @@ void userbuffers_send(const int srchandler, const size_t srcoffset, const int ds
         cudaLaunchKernelExC(&cfg, reinterpret_cast<void *>(kuserbuffers_pushsend), kernelArgs));
 #else
         cudaLaunchKernel(reinterpret_cast<void *>(kuserbuffers_pushsend), sms, threads, kernelArgs, 0, stream));
+#endif
   }
 }
 
@@ -2812,7 +2815,7 @@ void userbuffers_sendrecv_multiatomic(const int srchandler, const int dsthandler
   void *flagptr_send = GET_SEND_PTR_BY_INDEX(send_peerlocal, comm, dsthandler, 0);
   void *flagptr_recv = GET_RECV_PTR_BY_INDEX(recv_peer, comm, dsthandler, 0);
 
-#ifndef
+#ifndef __HIP_PLATFORM_AMD__
   SETUP_LAUNCH_CONFIG(comm->sms, 1024, stream);
 #else
   int sms = comm->sms;

transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h

@@ -58,6 +58,7 @@ class CommOverlapCore {
   int _comm_priority;
   bool _atomic_gemm{false};
   bool _is_p2p{false};
+  bool _use_rd{false};
 
   TensorWrapper _ubuf;
   TensorWrapper _counter;
@@ -93,6 +94,8 @@ class CommOverlapCore {
 
   bool is_p2p_overlap() { return _is_p2p; }
 
+  bool is_use_rd() { return _use_rd; }
+
   bool is_fp8_ubuf() { return _ubuf.element_size() == 1; }
 
   virtual void bulk_overlap(const TensorWrapper &A, bool transa, const TensorWrapper &B,
@@ -199,7 +202,9 @@ class CommOverlapBase : public CommOverlapCore {
                         TensorWrapper &D, TensorWrapper &bias, TensorWrapper &pre_gelu_out,
                         TensorWrapper &workspace, bool grad, bool accumulate,
                         bool use_split_accumulator, TensorWrapper &B_copy,
-                        cudaStream_t stream_main) override;
+                        cudaStream_t stream_main) override {
+    NVTE_ERROR("Operation not supported.");                      
+  };
 
   /*
   ** Split FPROP GEMM + ReduceScatter

transformer_engine/common/util/pybind_helper.h

@@ -14,7 +14,11 @@
 #include <transformer_engine/fused_attn.h>
 #include <transformer_engine/transformer_engine.h>
 
+#ifdef __HIP_PLATFORM_AMD__
+#include "hip_runtime.h"
+#else
 #include "cuda_runtime.h"
+#endif
 
 // Define fused-attention handles separately for USE_ROCM
 #ifndef USE_ROCM

transformer_engine/pytorch/csrc/extensions.h

@@ -13,14 +13,6 @@
 
 #include "common.h"
 
-#ifdef USE_ROCM
-namespace transformer_engine {
-//dummy CommOverlapCore, CommOverlapType in rocm
-class CommOverlapCore{};
-class CommOverlapType{};
-}
-#endif
-
 namespace transformer_engine::pytorch {
 
 /***************************************************************************************************
@@ -456,7 +448,7 @@ class CommOverlapP2P : torch::CustomClassHolder, public transformer_engine::Comm
                  int num_max_streams = NVTE_COMM_OVERLAP_MAX_STREAMS, int comm_cga_size = 2,
                  int gemm_priority = 0, int comm_priority = 0, int num_comm_sm = 3,
                  bool set_sm_margin = true, bool atomic_gemm = false, bool use_ce = true,
-                 bool aggregate = false);
+                 bool aggregate = false, bool use_rd = false);
 
   ~CommOverlapP2P() {}
 

transformer_engine/pytorch/csrc/extensions/gemm.cpp

@@ -187,7 +187,6 @@ std::vector<py::object> gemm(py::handle A, bool transa, py::handle B, bool trans
         std::move(swizzle_scaling_factors(B_tensor, !transb)));
 
     if (comm_overlap) {
-#ifndef USE_ROCM
       // Prepare extra output tensor
       TensorWrapper extra_output_tensor;
       if (extra_output.has_value()) {
@@ -213,6 +212,13 @@ std::vector<py::object> gemm(py::handle A, bool transa, py::handle B, bool trans
                                                  accumulate, use_split_accumulator,
                                                  extra_output_tensor, main_stream);
           });
+        } else if (comm_overlap->is_use_rd()) {
+          NVTE_SCOPED_GIL_RELEASE({
+            comm_overlap->split_overlap_ag_rd(A_tensor, transa, B_tensor, transb, D_tensor,
+                                                 bias_tensor, te_pre_gelu_out, te_workspace, grad,
+                                                 accumulate, use_split_accumulator,
+                                                 extra_output_tensor, main_stream);
+          });
         } else {
           NVTE_SCOPED_GIL_RELEASE({
             comm_overlap->split_overlap_ag(A_tensor, transa, B_tensor, transb, D_tensor,
@@ -238,9 +244,6 @@ std::vector<py::object> gemm(py::handle A, bool transa, py::handle B, bool trans
           });
         }
       }
-#else
-    NVTE_ERROR("ROCm TE does not support comm_overlap\n");
-#endif //!USE_ROCM
     } else {
       // Launch GEMM
       NVTE_SCOPED_GIL_RELEASE({

transformer_engine/pytorch/csrc/extensions/pybind.cpp

@@ -403,9 +403,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
            py::arg("num_max_streams") = NVTE_COMM_OVERLAP_MAX_STREAMS, py::arg("comm_cga_size") = 1,
            py::arg("gemm_priority") = 0, py::arg("comm_priority") = 0, py::arg("num_comm_sm") = 1,
            py::arg("set_sm_margin") = false, py::arg("atomic_gemm") = false,
-           py::arg("use_ce") = true, py::arg("aggregate") = false, py::arg("use_rd" = false))
+           py::arg("use_ce") = true, py::arg("aggregate") = false, py::arg("use_rd") = false)
       .def("copy_into_buffer", &CommOverlapP2P::copy_into_buffer, py::arg("input"),
            py::arg("local_chunk") = false)
       .def("get_buffer", &CommOverlapP2P::get_buffer, py::arg("local_chunk") = false,
-           py::arg("shape") = std::nullopt,);
+           py::arg("shape") = std::nullopt);
 }

transformer_engine/pytorch/module/base.py

@@ -308,6 +308,7 @@ def get_default_config(name):
             "comm_priority": _MAX_STREAM_PRIORITY,
             "gemm_priority": _MIN_STREAM_PRIORITY,
             "pipeline_rs_overlap_first_gemm": False,
+            "use_rd": False,
         }
         return default_cfg
 
@@ -326,6 +327,7 @@ def add_ub(
         comm_priority: int = 0,
         gemm_priority: int = 0,
         pipeline_rs_overlap_first_gemm: bool = False,
+        use_rd: bool = False,
     ) -> None:
         if atomic_gemm:
             warnings.warn(