TE with threading build (NVIDIA#1092)

* added threading build back * integrating threading for pytorch and paddle extensions * added messages --------- Signed-off-by: Phuong Nguyen <[email protected]> Co-authored-by: Tim Moon <[email protected]>
mgoldfarb-nvidia · Aug 14, 2024 · b2b77fd · b2b77fd
1 parent c32f69a
commit b2b77fd
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 4 deletions.
diff --git a/build_tools/paddle.py b/build_tools/paddle.py
@@ -6,6 +6,7 @@
 from pathlib import Path
 
 import setuptools
+import os
 
 from .utils import cuda_version
 
@@ -62,7 +63,7 @@ def setup_paddle_extension(
         print("Could not determine CUDA Toolkit version")
     else:
         if version >= (11, 2):
-            nvcc_flags.extend(["--threads", "4"])
+            nvcc_flags.extend(["--threads", os.getenv("NVTE_BUILD_THREADS_PER_JOB", "1")])
         if version >= (11, 0):
             nvcc_flags.extend(["-gencode", "arch=compute_80,code=sm_80"])
         if version >= (11, 8):

diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py
@@ -68,7 +68,7 @@ def setup_pytorch_extension(
         print("Could not determine CUDA Toolkit version")
     else:
         if version >= (11, 2):
-            nvcc_flags.extend(["--threads", "4"])
+            nvcc_flags.extend(["--threads", os.getenv("NVTE_BUILD_THREADS_PER_JOB", "1")])
         if version >= (11, 0):
             nvcc_flags.extend(["-gencode", "arch=compute_80,code=sm_80"])
         if version >= (11, 8):

diff --git a/build_tools/utils.py b/build_tools/utils.py
@@ -37,8 +37,8 @@ def get_max_jobs_for_parallel_build() -> int:
     num_jobs = 0
 
     # Check environment variable
-    if os.getenv("NVTE_MAX_BUILD_JOBS"):
-        num_jobs = int(os.getenv("NVTE_MAX_BUILD_JOBS"))
+    if os.getenv("NVTE_BUILD_MAX_JOBS"):
+        num_jobs = int(os.getenv("NVTE_BUILD_MAX_JOBS"))
     elif os.getenv("MAX_JOBS"):
         num_jobs = int(os.getenv("MAX_JOBS"))
 

diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt
@@ -14,6 +14,22 @@ set(CMAKE_CUDA_STANDARD_REQUIRED ON)
 
 project(transformer_engine LANGUAGES CUDA CXX)
 
+set(BUILD_THREADS_PER_JOB $ENV{NVTE_BUILD_THREADS_PER_JOB})
+if (NOT BUILD_THREADS_PER_JOB)
+  set(BUILD_THREADS_PER_JOB 1)
+endif()
+set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --threads ${BUILD_THREADS_PER_JOB}")
+
+if(DEFINED ENV{MAX_JOBS})
+  set(JOBS $ENV{MAX_JOBS})
+elseif(DEFINED ENV{NVTE_BUILD_MAX_JOBS})
+  set(JOBS $ENV{NVTE_BUILD_MAX_JOBS})
+else()
+  set(JOBS "max number of")
+endif()
+
+message(STATUS "Parallel build with ${JOBS} jobs and ${BUILD_THREADS_PER_JOB} threads per job")
+
 if (CMAKE_BUILD_TYPE STREQUAL "Debug")
   set(CMAKE_CUDA_FLAGS_DEBUG "${CMAKE_CUDA_FLAGS_DEBUG} -G")
 endif()