diff --git a/build_tools/jax.py b/build_tools/jax.py index d61c26c128..92081625e7 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -15,6 +15,7 @@ all_files_in_dir, cudnn_frontend_include_path, debug_build_enabled, + get_bolt_build_flags, setup_mpi_flags, nccl_ep_enabled, ) @@ -109,6 +110,9 @@ def setup_jax_extension( else: cxx_flags.append("-g0") + bolt_cxx_flags, linker_flags = get_bolt_build_flags() + cxx_flags.extend(bolt_cxx_flags) + setup_mpi_flags(include_dirs, cxx_flags) if bool(int(os.getenv("NVTE_WITH_CUBLASMP", 0))): @@ -125,5 +129,6 @@ def setup_jax_extension( sources=[str(path) for path in sources], include_dirs=[str(path) for path in include_dirs], extra_compile_args=cxx_flags, + extra_link_args=linker_flags, libraries=["nccl"], ) diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index 2bb238c522..9865c8a657 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -15,6 +15,7 @@ cuda_version, get_cuda_include_dirs, debug_build_enabled, + get_bolt_build_flags, setup_mpi_flags, ) from typing import List @@ -75,6 +76,9 @@ def setup_pytorch_extension( else: cxx_flags.append("-g0") + bolt_cxx_flags, linker_flags = get_bolt_build_flags() + cxx_flags.extend(bolt_cxx_flags) + # Version-dependent CUDA options try: version = cuda_version() @@ -121,6 +125,7 @@ def setup_pytorch_extension( sources=[str(src) for src in sources], include_dirs=[str(inc) for inc in include_dirs], extra_compile_args={"cxx": cxx_flags}, + extra_link_args=linker_flags, libraries=[str(lib) for lib in libraries], library_dirs=[str(lib_dir) for lib_dir in library_dirs], ) diff --git a/build_tools/utils.py b/build_tools/utils.py index 6a8168e7d0..a6d5d8834b 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -43,6 +43,42 @@ def debug_build_enabled() -> bool: return bool(int(os.getenv("NVTE_BUILD_DEBUG", "0"))) +@functools.lru_cache(maxsize=None) +def bolt_compatible_build_enabled() -> bool: + """Whether to build host ELF libraries with BOLT-compatible options.""" + configured = os.getenv("NVTE_ENABLE_BOLT_COMPATIBLE") + if configured is None: + enabled = platform.system() == "Linux" and platform.machine().lower() in ( + "aarch64", + "arm64", + ) + else: + enabled = bool(int(configured)) + + if enabled and platform.system() != "Linux": + raise RuntimeError("NVTE_ENABLE_BOLT_COMPATIBLE is only supported on Linux") + return enabled + + +def get_bolt_build_flags() -> Tuple[List[str], List[str]]: + """BOLT-compatible host compiler and linker flags.""" + if not bolt_compatible_build_enabled(): + return [], [] + + compiler_flags = ["-fno-reorder-blocks-and-partition", "-fno-jump-tables"] + linker_flags = ["-Wl,--emit-relocs", "-Wl,-z,now"] + if platform.machine().lower() in ("aarch64", "arm64"): + compiler_flags.extend( + [ + "-mno-fix-cortex-a53-835769", + "-mno-fix-cortex-a53-843419", + ] + ) + # The Cortex-A53 843419 workaround is applied by the linker. + linker_flags.append("-mno-fix-cortex-a53-843419") + return compiler_flags, linker_flags + + @functools.lru_cache(maxsize=None) def get_max_jobs_for_parallel_build() -> int: """Number of parallel jobs for Nina build""" diff --git a/docs/envvars.rst b/docs/envvars.rst index 97eaed5ddc..9c43bad4e3 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -60,6 +60,12 @@ Build Configuration :Default: None :Description: Path to the CMake build directory for incremental builds. If set, CMake will use this directory for build artifacts. +.. envvar:: NVTE_ENABLE_BOLT_COMPATIBLE + + :Type: ``int`` (0 or 1) + :Default: ``1`` on Linux Arm64 (AArch64), ``0`` otherwise + :Description: Build the core, PyTorch, and JAX host ELFs with LLVM BOLT-compatible compiler and linker flags. On Arm64, this also disables the Cortex-A53 835769 and 843419 errata workarounds as required by BOLT. Set to ``0`` to disable the default on Arm64, or ``1`` to opt in on another supported Linux architecture. When using :envvar:`NVTE_CMAKE_BUILD_DIR`, use a fresh build directory after changing this setting. + .. envvar:: NVTE_RELEASE_BUILD :Type: ``int`` (0 or 1) diff --git a/setup.py b/setup.py index 2a8a9d7688..1b52e49bc3 100644 --- a/setup.py +++ b/setup.py @@ -77,6 +77,11 @@ def setup_common_extension() -> CMakeExtension: if bool(int(os.getenv("NVTE_BUILD_ACTIVATION_WITH_FAST_MATH", "0"))): cmake_flags.append("-DNVTE_BUILD_ACTIVATION_WITH_FAST_MATH=ON") + bolt_compatible = os.getenv("NVTE_ENABLE_BOLT_COMPATIBLE") + if bolt_compatible is not None: + bolt_compatible = "ON" if bool(int(bolt_compatible)) else "OFF" + cmake_flags.append(f"-DNVTE_ENABLE_BOLT_COMPATIBLE={bolt_compatible}") + if bool(int(os.getenv("NVTE_WITH_CUBLASMP", "0"))): cmake_flags.append("-DNVTE_WITH_CUBLASMP=ON") cublasmp_dir = os.getenv("CUBLASMP_HOME") or metadata.distribution( diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 6ceadc7405..81866eb774 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -15,6 +15,81 @@ endif() # Transformer Engine library project(transformer_engine LANGUAGES CUDA CXX) +# BOLT-compatible builds are enabled by default on Linux Arm64, where function +# layout has the largest impact on Grace CPUs. Other platforms remain opt-in. +string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" _nvte_system_processor) +set(NVTE_TARGET_IS_ARM64 OFF) +if(_nvte_system_processor MATCHES "^(aarch64|arm64)$") + set(NVTE_TARGET_IS_ARM64 ON) +endif() +set(_nvte_bolt_compatible_default OFF) +if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND NVTE_TARGET_IS_ARM64) + set(_nvte_bolt_compatible_default ON) +endif() +option(NVTE_ENABLE_BOLT_COMPATIBLE + "Build host ELF objects with LLVM BOLT-compatible compile and link flags" + ${_nvte_bolt_compatible_default}) +unset(_nvte_bolt_compatible_default) +unset(_nvte_system_processor) + +if(NVTE_ENABLE_BOLT_COMPATIBLE) + if(NOT CMAKE_SYSTEM_NAME STREQUAL "Linux") + message(FATAL_ERROR + "NVTE_ENABLE_BOLT_COMPATIBLE is only supported for Linux ELF builds") + endif() + + include(CheckCXXCompilerFlag) + check_cxx_compiler_flag("-fno-reorder-blocks-and-partition" + NVTE_CXX_SUPPORTS_FNO_REORDER_BLOCKS_AND_PARTITION) + check_cxx_compiler_flag("-fno-jump-tables" + NVTE_CXX_SUPPORTS_FNO_JUMP_TABLES) + if(NVTE_TARGET_IS_ARM64) + check_cxx_compiler_flag("-mno-fix-cortex-a53-835769" + NVTE_CXX_SUPPORTS_MNO_FIX_CORTEX_A53_835769) + check_cxx_compiler_flag("-mno-fix-cortex-a53-843419" + NVTE_CXX_SUPPORTS_MNO_FIX_CORTEX_A53_843419) + endif() + if(NOT NVTE_CXX_SUPPORTS_FNO_REORDER_BLOCKS_AND_PARTITION + OR NOT NVTE_CXX_SUPPORTS_FNO_JUMP_TABLES + OR (NVTE_TARGET_IS_ARM64 + AND (NOT NVTE_CXX_SUPPORTS_MNO_FIX_CORTEX_A53_835769 + OR NOT NVTE_CXX_SUPPORTS_MNO_FIX_CORTEX_A53_843419))) + message(FATAL_ERROR + "The host C++ compiler does not support the flags required for " + "NVTE_ENABLE_BOLT_COMPATIBLE") + endif() + + # BOLT consumes the regular symbol table and emitted relocations. Shadow the + # cached strip tool for this configure without changing it for later builds. + set(CMAKE_STRIP "") +endif() + +function(nvte_enable_bolt_compatible_compile_options TARGET_NAME) + if(NOT NVTE_ENABLE_BOLT_COMPATIBLE) + return() + endif() + + target_compile_options( + ${TARGET_NAME} + PRIVATE + $<$:-fno-reorder-blocks-and-partition> + $<$:-fno-jump-tables> + $<$:-Xcompiler=-fno-reorder-blocks-and-partition> + $<$:-Xcompiler=-fno-jump-tables> + ) + + if(NVTE_TARGET_IS_ARM64) + target_compile_options( + ${TARGET_NAME} + PRIVATE + $<$:-mno-fix-cortex-a53-835769> + $<$:-mno-fix-cortex-a53-843419> + $<$:-Xcompiler=-mno-fix-cortex-a53-835769> + $<$:-Xcompiler=-mno-fix-cortex-a53-843419> + ) + endif() +endfunction() + # CUDA Toolkit find_package(CUDAToolkit REQUIRED) if (CUDAToolkit_VERSION VERSION_LESS 12.1) @@ -137,6 +212,7 @@ function(get_nccl_version OUT_VERSION INCLUDE_DIR) PARENT_SCOPE) endfunction() +set(NVTE_BOLT_MIN_NCCL_VERSION "2.30.4") get_nccl_version(NCCL_VERSION "${NCCL_INCLUDE_DIR}") function(find_cublasmp_version OUT_VERSION OUT_INCLUDE_DIR SEARCH_DIR) @@ -332,6 +408,7 @@ foreach(cuda_source IN LISTS transformer_engine_cuda_arch_specific_sources) endforeach() add_library(transformer_engine SHARED ${transformer_engine_SOURCES}) +nvte_enable_bolt_compatible_compile_options(transformer_engine) # This is TE-specific and should not apply to all targets target_link_options( @@ -339,6 +416,22 @@ target_link_options( PRIVATE "LINKER:--version-script=${CMAKE_CURRENT_SOURCE_DIR}/libtransformer_engine.version" ) +if(NVTE_ENABLE_BOLT_COMPATIBLE) + target_link_options( + transformer_engine + PRIVATE + "LINKER:--emit-relocs" + "LINKER:-z,now" + ) + if(NVTE_TARGET_IS_ARM64) + # The Cortex-A53 843419 workaround is applied by the linker. + target_link_options( + transformer_engine + PRIVATE + "LINKER:--no-fix-cortex-a53-843419" + ) + endif() +endif() # Disable CMake's automatic architecture flag injection. # All architectures are handled explicitly via per-source COMPILE_OPTIONS @@ -395,6 +488,7 @@ endif() option(NVTE_ENABLE_NVSHMEM "Compile with NVSHMEM library" OFF) if (NVTE_ENABLE_NVSHMEM) add_subdirectory(nvshmem_api) + nvte_enable_bolt_compatible_compile_options(nvshmemapi) target_link_libraries(transformer_engine PUBLIC nvshmemapi) target_include_directories(transformer_engine PUBLIC ${NVSHMEMAPI_INCLUDE_DIR}) endif() @@ -452,11 +546,19 @@ endif() # -- NCCL EP (on by default, HT mode only) --------------------------------- # Set -DNVTE_WITH_NCCL_EP=OFF (or NVTE_WITH_NCCL_EP=0 in setup.py) to -# skip NCCL EP entirely - useful on older images whose system NCCL is below -# the 2.30.4 EP minimum. +# skip NCCL EP entirely. option(NVTE_WITH_NCCL_EP "Build NCCL EP into libtransformer_engine.so" ON) if(NVTE_WITH_NCCL_EP) -# SM>=90 and NCCL>=2.30.4 are gated at runtime in EPBackend::initialize. +# SM>=90 and the NCCL version are gated at runtime in EPBackend::initialize. +# BOLT-compatible builds additionally require the matching NCCL symbols at load +# time because they use eager symbol binding. +if(NVTE_ENABLE_BOLT_COMPATIBLE + AND NCCL_VERSION VERSION_LESS NVTE_BOLT_MIN_NCCL_VERSION) + message(FATAL_ERROR + "BOLT-compatible NCCL EP builds require NCCL >= " + "${NVTE_BOLT_MIN_NCCL_VERSION}, but found NCCL ${NCCL_VERSION} in " + "${NCCL_INCLUDE_DIR}/nccl.h") +endif() # -- NCCL EP headers -------------------------------------------------------- # Headers + libs are produced by the in-tree 3rdparty/nccl-extensions submodule build # (auto-built by setup.py via build_nccl_ep_submodule). @@ -471,10 +573,9 @@ endif() message(STATUS "NCCL EP headers: ${NCCL_EP_INCLUDE_DIR}") # -- libnccl_ep.a ----------------------------------------------------------- -# Statically linked into libtransformer_engine.so. EPBackend::initialize checks -# NCCL >= 2.30.4 before any nccl_ep call, so the newer NCCL symbols nccl_ep -# imports stay unresolved (and harmless) under default ELF lazy binding when -# the gate trips. LD_BIND_NOW environments lose this property. +# Statically linked into libtransformer_engine.so. BOLT-compatible builds use +# eager symbol binding, so the build-time and runtime NCCL minimum must remain +# synchronized with the symbols imported by this archive. set(NCCL_EP_LIB_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/lib") find_file(NCCL_EP_LIB NAMES libnccl_ep.a