# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Installation script.""" # NVTE_FRAMEWORK=pytorch NVTE_USE_ROCM=1 NVTE_USE_HIPBLASLT=1 NVTE_USE_ROCBLAS=1 CMAKE_PREFIX_PATH=/opt/dtk/lib/cmake/amd_comgr/ MPI_HOME=/opt/mpi/ NVTE_UB_WITH_MPI=1 CXX=hipcc pip3 install . -v # NVTE_FRAMEWORK=pytorch NVTE_USE_ROCM=1 NVTE_USE_HIPBLASLT=1 NVTE_USE_ROCBLAS=1 CMAKE_PREFIX_PATH=/opt/dtk/lib/cmake/amd_comgr/ MPI_HOME=/opt/mpi/ NVTE_UB_WITH_MPI=1 CXX=hipcc PYTHONPATH=/home/TransformerEngine/3rdparty/hipify_torch:$PYTHONPATH python3 setup.py bdist_wheel import os import sys import time from pathlib import Path from typing import List, Tuple import setuptools from wheel.bdist_wheel import bdist_wheel from build_tools.build_ext import CMakeExtension, get_build_ext from build_tools.te_version import te_version from build_tools.utils import ( rocm_build, cuda_archs, found_cmake, found_ninja, found_pybind11, get_frameworks, install_and_import, remove_dups, uninstall_te_wheel_packages, ) frameworks = get_frameworks() current_file_path = Path(__file__).parent.resolve() from setuptools.command.build_ext import build_ext as BuildExtension os.environ["NVTE_PROJECT_BUILDING"] = "1" if "pytorch" in frameworks: from torch.utils.cpp_extension import BuildExtension elif "jax" in frameworks: install_and_import("pybind11[global]") from pybind11.setup_helpers import build_ext as BuildExtension CMakeBuildExtension = get_build_ext(BuildExtension) if rocm_build(): archs = None else: archs = cuda_archs() class TimedBdist(bdist_wheel): """Helper class to measure build time""" def run(self): start_time = time.perf_counter() super().run() total_time = time.perf_counter() - start_time print(f"Total time for bdist_wheel: {total_time:.2f} seconds") def setup_common_extension() -> CMakeExtension: """Setup CMake extension for common library""" if rocm_build(): cmake_flags = [] if bool(int(os.getenv("NVTE_BUILD_SUPPRESS_UNUSED_WARNING", "1"))): cmake_flags.append("-DNVTE_BUILD_SUPPRESS_UNUSED_WARNING=ON") if bool(int(os.getenv("NVTE_BUILD_SUPPRESS_RETURN_TYPE_WARNING", "0"))): cmake_flags.append("-DNVTE_BUILD_SUPPRESS_RETURN_TYPE_WARNING=ON") if bool(int(os.getenv("NVTE_BUILD_SUPPRESS_SIGN_COMPARE", "0"))): cmake_flags.append("-DNVTE_BUILD_SUPPRESS_SIGN_COMPARE_WARNING=ON") else: cmake_flags = ["-DCMAKE_CUDA_ARCHITECTURES={}".format(archs)] if bool(int(os.getenv("NVTE_UB_WITH_MPI", "0"))): assert ( os.getenv("MPI_HOME") is not None ), "MPI_HOME must be set when compiling with NVTE_UB_WITH_MPI=1" cmake_flags.append("-DNVTE_UB_WITH_MPI=ON") if bool(int(os.getenv("NVTE_BUILD_ACTIVATION_WITH_FAST_MATH", "0"))): cmake_flags.append("-DNVTE_BUILD_ACTIVATION_WITH_FAST_MATH=ON") # Project directory root root_path = Path(__file__).resolve().parent if rocm_build(): if os.getenv("NVTE_USE_HIPBLASLT") is not None: cmake_flags.append("-DUSE_HIPBLASLT=ON") if os.getenv("NVTE_USE_ROCBLAS") is not None: cmake_flags.append("-DUSE_ROCBLAS=ON") return CMakeExtension( name="transformer_engine", cmake_path=root_path / Path("transformer_engine/common"), cmake_flags=cmake_flags, ) def setup_requirements() -> Tuple[List[str], List[str], List[str]]: """Setup Python dependencies Returns dependencies for build, runtime, and testing. """ # Common requirements setup_reqs: List[str] = [] install_reqs: List[str] = [ "pydantic", "importlib-metadata>=1.0", "packaging", ] test_reqs: List[str] = ["pytest>=8.2.1"] # Requirements that may be installed outside of Python if not found_cmake(): setup_reqs.append("cmake>=3.21") if not found_ninja(): setup_reqs.append("ninja") if not found_pybind11(): setup_reqs.append("pybind11") # Framework-specific requirements if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))): if "pytorch" in frameworks: install_reqs.extend(["torch>=2.1"]) # Blackwell is not supported as of Triton 3.2.0, need custom internal build # install_reqs.append("triton") test_reqs.extend(["numpy", "torchvision", "prettytable", "PyYAML"]) if "jax" in frameworks: install_reqs.extend(["jax", "flax>=0.7.1"]) # test_reqs.extend(["numpy", "praxis"]) test_reqs.extend(["numpy"]) return [remove_dups(reqs) for reqs in [setup_reqs, install_reqs, test_reqs]] if __name__ == "__main__": __version__ = te_version() with open("README.rst", encoding="utf-8") as f: long_description = f.read() # Settings for building top level empty package for dependency management. if bool(int(os.getenv("NVTE_BUILD_METAPACKAGE", "0"))): assert bool( int(os.getenv("NVTE_RELEASE_BUILD", "0")) ), "NVTE_RELEASE_BUILD env must be set for metapackage build." ext_modules = [] cmdclass = {} package_data = {} include_package_data = False setup_requires = [] install_requires = ([f"transformer_engine_cu12=={__version__}"],) extras_require = { "pytorch": [f"transformer_engine_torch=={__version__}"], "jax": [f"transformer_engine_jax=={__version__}"], } else: setup_requires, install_requires, test_requires = setup_requirements() ext_modules = [setup_common_extension()] cmdclass = {"build_ext": CMakeBuildExtension, "bdist_wheel": TimedBdist} package_data = {"": ["VERSION.txt"]} include_package_data = True extras_require = {"test": test_requires} if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))): # Remove residual FW packages since compiling from source # results in a single binary with FW extensions included. uninstall_te_wheel_packages() if "pytorch" in frameworks: from build_tools.pytorch import setup_pytorch_extension ext_modules.append( setup_pytorch_extension( "transformer_engine/pytorch/csrc", current_file_path / "transformer_engine" / "pytorch" / "csrc", current_file_path / "transformer_engine", ) ) if "jax" in frameworks: from build_tools.jax import setup_jax_extension ext_modules.append( setup_jax_extension( "transformer_engine/jax/csrc", current_file_path / "transformer_engine" / "jax" / "csrc", current_file_path / "transformer_engine", ) ) # Configure package setuptools.setup( name="transformer_engine", version=__version__, packages=setuptools.find_packages( include=[ "transformer_engine", "transformer_engine.*", "transformer_engine/build_tools", ], ), extras_require=extras_require, description="Transformer acceleration library", long_description=long_description, long_description_content_type="text/x-rst", ext_modules=ext_modules, cmdclass={"build_ext": CMakeBuildExtension, "bdist_wheel": TimedBdist}, python_requires=">=3.8, <3.13", classifiers=[ "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", ], setup_requires=setup_requires, install_requires=install_requires, license_files=("LICENSE",), include_package_data=include_package_data, package_data=package_data, )