Add NVTX to TE modules (#50)

* Add NVTX to TE modules
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fix
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fix pylint
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fix NVTX in _prepare_backward
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Add NVTX to C API
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fix cpplint and link nvToolsExt
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Add NVTX to GeGlu
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

qa/L0_lint/pylintrc

@@ -22,7 +22,8 @@ disable=too-many-locals,
         attribute-defined-outside-init,
         global-statement,
         too-many-branches,
-        global-variable-not-assigned
+        global-variable-not-assigned,
+        redefined-argument-from-local
 
 [TYPECHECK]
 ignored-modules=torch

transformer_engine/common/CMakeLists.txt

@@ -40,9 +40,9 @@ add_library(transformer_engine SHARED
 
 target_include_directories(transformer_engine PUBLIC "${PROJECT_SOURCE_DIR}/include")
 
-find_package(CUDAToolkit REQUIRED cublas)
+find_package(CUDAToolkit REQUIRED cublas nvToolsExt)
 
-list(APPEND transformer_engine_LINKER_LIBS CUDA::cublas CUDA::cudart)
+list(APPEND transformer_engine_LINKER_LIBS CUDA::cublas CUDA::cudart CUDA::nvToolsExt)
 target_link_libraries(transformer_engine PUBLIC ${transformer_engine_LINKER_LIBS})
 
 target_include_directories(transformer_engine PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})

transformer_engine/common/activation/gelu.cu

@@ -116,6 +116,7 @@ void dgeglu(const Tensor &grad,
 void nvte_gelu(const NVTETensor input,
                NVTETensor output,
                cudaStream_t stream) {
+  NVTE_API_CALL(nvte_gelu);
   using namespace transformer_engine;
   gelu_cast(*reinterpret_cast<const Tensor*>(input),
             reinterpret_cast<Tensor*>(output),
@@ -125,6 +126,7 @@ void nvte_gelu(const NVTETensor input,
 void nvte_geglu(const NVTETensor input,
                 NVTETensor output,
                 cudaStream_t stream) {
+  NVTE_API_CALL(nvte_geglu);
   using namespace transformer_engine;
   geglu_cast(*reinterpret_cast<const Tensor*>(input),
              reinterpret_cast<Tensor*>(output),
@@ -135,6 +137,7 @@ void nvte_dgeglu(const NVTETensor grad,
                  const NVTETensor input,
                  NVTETensor output,
                  cudaStream_t stream) {
+  NVTE_API_CALL(nvte_dgeglu);
   using namespace transformer_engine;
   dgeglu(*reinterpret_cast<const Tensor*>(grad),
          *reinterpret_cast<const Tensor*>(input),

transformer_engine/common/common.h

@@ -20,6 +20,7 @@
 #include <string>
 #include <tuple>
 #include <vector>
+#include "nvtx.h"
 
 namespace transformer_engine {
 
@@ -285,6 +286,9 @@ void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empt
 
 bool is_fp8_dtype(const DType t);
 
+#define NVTE_API_CALL(api_name) \
+  transformer_engine::nvtx::NVTXWrapper _ ## api_name ## _nvtx_wrapper(#api_name);
+
 }  // namespace transformer_engine
 
 #endif  // TRANSFORMER_ENGINE_COMMON_COMMON_H_

transformer_engine/common/fused_softmax/scaled_masked_softmax.cu

@@ -1060,12 +1060,13 @@ void nvte_scaled_softmax_forward(
     float scale_factor,
     cudaStream_t stream
 ) {
-    using namespace transformer_engine;
-    scaled_softmax_forward(
-        *reinterpret_cast<const Tensor*>(input),
-        reinterpret_cast<Tensor*>(softmax_results),
-        scale_factor,
-        stream);
+  NVTE_API_CALL(nvte_scaled_softmax_forward);
+  using namespace transformer_engine;
+  scaled_softmax_forward(
+      *reinterpret_cast<const Tensor*>(input),
+      reinterpret_cast<Tensor*>(softmax_results),
+      scale_factor,
+      stream);
 }
 
 
@@ -1076,13 +1077,14 @@ void nvte_scaled_softmax_backward(
     float scale_factor,
     cudaStream_t stream
 ) {
-    using namespace transformer_engine;
-    scaled_softmax_backward(
-        *reinterpret_cast<Tensor*>(output_grads),
-        *reinterpret_cast<const Tensor*>(incoming_grads),
-        *reinterpret_cast<const Tensor*>(softmax_results),
-        scale_factor,
-        stream);
+  NVTE_API_CALL(nvte_scaled_softmax_backward);
+  using namespace transformer_engine;
+  scaled_softmax_backward(
+      *reinterpret_cast<Tensor*>(output_grads),
+      *reinterpret_cast<const Tensor*>(incoming_grads),
+      *reinterpret_cast<const Tensor*>(softmax_results),
+      scale_factor,
+      stream);
 }
 
 
@@ -1093,13 +1095,14 @@ void nvte_scaled_masked_softmax_forward(
     float scale_factor,
     cudaStream_t stream
 ) {
-    using namespace transformer_engine;
-    scaled_masked_softmax_forward(
-        *reinterpret_cast<const Tensor*>(input),
-        *reinterpret_cast<const Tensor*>(mask),
-        reinterpret_cast<Tensor*>(softmax_results),
-        scale_factor,
-        stream);
+  NVTE_API_CALL(nvte_scaled_masked_softmax_forward);
+  using namespace transformer_engine;
+  scaled_masked_softmax_forward(
+      *reinterpret_cast<const Tensor*>(input),
+      *reinterpret_cast<const Tensor*>(mask),
+      reinterpret_cast<Tensor*>(softmax_results),
+      scale_factor,
+      stream);
 }
 
 
@@ -1110,11 +1113,12 @@ void nvte_scaled_masked_softmax_backward(
     float scale_factor,
     cudaStream_t stream
 ) {
-    using namespace transformer_engine;
-    scaled_masked_softmax_backward(
-        *reinterpret_cast<Tensor*>(output_grads),
-        *reinterpret_cast<const Tensor*>(incoming_grads),
-        *reinterpret_cast<const Tensor*>(softmax_results),
-        scale_factor,
-        stream);
+  NVTE_API_CALL(nvte_scaled_masked_softmax_backward);
+  using namespace transformer_engine;
+  scaled_masked_softmax_backward(
+      *reinterpret_cast<Tensor*>(output_grads),
+      *reinterpret_cast<const Tensor*>(incoming_grads),
+      *reinterpret_cast<const Tensor*>(softmax_results),
+      scale_factor,
+      stream);
 }

transformer_engine/common/gemm/cublaslt_gemm.cu

@@ -236,6 +236,7 @@ void nvte_cublas_gemm(const NVTETensor A,
                       bool accumulate,
                       bool use_split_accumulator,
                       cudaStream_t stream) {
+  NVTE_API_CALL(nvte_cublas_gemm);
   using namespace transformer_engine;
   const Tensor *inputA = reinterpret_cast<const Tensor*>(A);
   const Tensor *inputB = reinterpret_cast<const Tensor*>(B);

transformer_engine/common/layer_norm/ln_api.cpp

@@ -375,6 +375,7 @@ void nvte_layernorm_fwd(const NVTETensor x,       // BxSxhidden_size
                         const int multiprocessorCount,
                         NVTETensor workspace,
                         NVTETensor barrier) {
+  NVTE_API_CALL(nvte_layernorm_fwd);
   using namespace transformer_engine;
   layernorm_fwd(*reinterpret_cast<const Tensor*>(x),
                 *reinterpret_cast<const Tensor*>(gamma),
@@ -403,6 +404,7 @@ void nvte_layernorm_bwd(const NVTETensor dz,       // BxSxhidden_size
                         const int multiprocessorCount,
                         NVTETensor workspace,
                         NVTETensor barrier) {
+  NVTE_API_CALL(nvte_layernorm_fwd);
   using namespace transformer_engine;
   layernorm_bwd(*reinterpret_cast<const Tensor*>(dz),
                 *reinterpret_cast<const Tensor*>(x),

transformer_engine/common/nvtx.h

+/*************************************************************************
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See LICENSE for license information.
+ ************************************************************************/
+
+#ifndef TRANSFORMER_ENGINE_COMMON_NVTX_H_
+#define TRANSFORMER_ENGINE_COMMON_NVTX_H_
+
+#include <string>
+#include <nvToolsExt.h>
+
+namespace transformer_engine::nvtx {
+
+struct NVTXWrapper {
+  explicit NVTXWrapper(const std::string &name) {
+    nvtxRangePush(name.c_str());
+  }
+
+  ~NVTXWrapper() {
+    nvtxRangePop();
+  }
+};
+
+}  // namespace transformer_engine::nvtx
+
+#endif  // TRANSFORMER_ENGINE_COMMON_NVTX_H_

transformer_engine/common/rmsnorm/rmsnorm_api.cpp

@@ -287,11 +287,12 @@ void nvte_rmsnorm_fwd(const NVTETensor x,      // Nxhidden_size
                       const NVTETensor gamma,  // hidden_size
                       const float epsilon, NVTETensor z, NVTETensor rsigma, cudaStream_t stream,
                       const int multiprocessorCount, NVTETensor workspace, NVTETensor barrier) {
-    using namespace transformer_engine;
-    rmsnorm_fwd(*reinterpret_cast<const Tensor *>(x), *reinterpret_cast<const Tensor *>(gamma),
-                epsilon, reinterpret_cast<Tensor *>(z), reinterpret_cast<Tensor *>(rsigma), stream,
-                multiprocessorCount, reinterpret_cast<Tensor *>(workspace),
-                reinterpret_cast<Tensor *>(barrier));
+  NVTE_API_CALL(nvte_rmsnorm_fwd);
+  using namespace transformer_engine;
+  rmsnorm_fwd(*reinterpret_cast<const Tensor *>(x), *reinterpret_cast<const Tensor *>(gamma),
+              epsilon, reinterpret_cast<Tensor *>(z), reinterpret_cast<Tensor *>(rsigma), stream,
+              multiprocessorCount, reinterpret_cast<Tensor *>(workspace),
+              reinterpret_cast<Tensor *>(barrier));
 }
 
 void nvte_rmsnorm_bwd(const NVTETensor dz,      // Nxhidden_size
@@ -300,10 +301,11 @@ void nvte_rmsnorm_bwd(const NVTETensor dz,      // Nxhidden_size
                       const NVTETensor gamma,   // hidden_size
                       NVTETensor dx, NVTETensor dgamma, NVTETensor dgamma_part, cudaStream_t stream,
                       const int multiprocessorCount, NVTETensor workspace, NVTETensor barrier) {
-    using namespace transformer_engine;
-    rmsnorm_bwd(*reinterpret_cast<const Tensor *>(dz), *reinterpret_cast<const Tensor *>(x),
-                *reinterpret_cast<const Tensor *>(rsigma), *reinterpret_cast<const Tensor *>(gamma),
-                reinterpret_cast<Tensor *>(dx), reinterpret_cast<Tensor *>(dgamma),
-                reinterpret_cast<Tensor *>(dgamma_part), stream, multiprocessorCount,
-                reinterpret_cast<Tensor *>(workspace), reinterpret_cast<Tensor *>(barrier));
+  NVTE_API_CALL(nvte_rmsnorm_bwd);
+  using namespace transformer_engine;
+  rmsnorm_bwd(*reinterpret_cast<const Tensor *>(dz), *reinterpret_cast<const Tensor *>(x),
+              *reinterpret_cast<const Tensor *>(rsigma), *reinterpret_cast<const Tensor *>(gamma),
+              reinterpret_cast<Tensor *>(dx), reinterpret_cast<Tensor *>(dgamma),
+              reinterpret_cast<Tensor *>(dgamma_part), stream, multiprocessorCount,
+              reinterpret_cast<Tensor *>(workspace), reinterpret_cast<Tensor *>(barrier));
 }

transformer_engine/common/transpose/cast_transpose.cu

@@ -392,6 +392,7 @@ void nvte_cast_transpose(const NVTETensor input,
                          NVTETensor cast_output,
                          NVTETensor transposed_output,
                          cudaStream_t stream) {
+  NVTE_API_CALL(nvte_cast_transpose);
   using namespace transformer_engine;
   cast_transpose(*reinterpret_cast<const Tensor*>(input),
                  reinterpret_cast<Tensor*>(cast_output),

transformer_engine/common/transpose/cast_transpose_fusion.cu

@@ -1566,6 +1566,7 @@ void nvte_cast_transpose_dbias(const NVTETensor input,
                                NVTETensor dbias,
                                NVTETensor workspace,
                                cudaStream_t stream) {
+  NVTE_API_CALL(nvte_cast_transpose_dbias);
   using namespace transformer_engine;
   cast_transpose_dbias(*reinterpret_cast<const Tensor*>(input),
                        reinterpret_cast<Tensor*>(cast_output),
@@ -1582,6 +1583,7 @@ void nvte_cast_transpose_dbias_dgelu(const NVTETensor input,
                                      NVTETensor dbias,
                                      NVTETensor workspace,
                                      cudaStream_t stream) {
+  NVTE_API_CALL(nvte_cast_transpose_dbias_dgelu);
   using namespace transformer_engine;
   cast_transpose_dbias_dgelu(*reinterpret_cast<const Tensor*>(input),
                              *reinterpret_cast<const Tensor*>(gelu_input),
@@ -1597,6 +1599,7 @@ void nvte_dgeglu_cast_transpose(const NVTETensor input,
                                 NVTETensor cast_output,
                                 NVTETensor transposed_output,
                                 cudaStream_t stream) {
+  NVTE_API_CALL(nvte_dgeglu_cast_transpose);
   using namespace transformer_engine;
   dgeglu_cast_transpose(*reinterpret_cast<const Tensor*>(input),
                         *reinterpret_cast<const Tensor*>(geglu_input),

transformer_engine/common/transpose/multi_cast_transpose.cu

@@ -339,6 +339,7 @@ void nvte_multi_cast_transpose(size_t num_tensors,
                                NVTETensor* cast_output_list,
                                NVTETensor* transposed_output_list,
                                cudaStream_t stream) {
+  NVTE_API_CALL(nvte_multi_cast_transpose);
   using namespace transformer_engine;
   std::vector<Tensor*> input_list_,
     cast_output_list_, transposed_output_list_;

transformer_engine/common/transpose/transpose.cu

@@ -308,6 +308,7 @@ void transpose(const Tensor &input,
 void nvte_transpose(const NVTETensor input,
                     NVTETensor transposed_output,
                     cudaStream_t stream) {
+  NVTE_API_CALL(nvte_transpose);
   using namespace transformer_engine;
   transpose(*reinterpret_cast<const Tensor*>(input),
             reinterpret_cast<Tensor*>(transposed_output),

transformer_engine/common/util/cast.cu

@@ -95,6 +95,7 @@ void fp8_dequantize(const Tensor &input,
 void nvte_fp8_quantize(const NVTETensor input,
                        NVTETensor output,
                        cudaStream_t stream) {
+  NVTE_API_CALL(nvte_fp8_quantize);
   using namespace transformer_engine;
   fp8_quantize(*reinterpret_cast<const Tensor*>(input),
                reinterpret_cast<Tensor*>(output),
@@ -104,6 +105,7 @@ void nvte_fp8_quantize(const NVTETensor input,
 void nvte_fp8_dequantize(const NVTETensor input,
                          NVTETensor output,
                          cudaStream_t stream) {
+  NVTE_API_CALL(nvte_fp8_dequantize);
   using namespace transformer_engine;
   fp8_dequantize(*reinterpret_cast<const Tensor*>(input),
                  reinterpret_cast<Tensor*>(output),