setup.py 6.82 KB
Newer Older
1
# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Przemek Tredak's avatar
Przemek Tredak committed
2
3
4
#
# See LICENSE for license information.

5
6
"""Installation script."""

7
from importlib import metadata
Przemek Tredak's avatar
Przemek Tredak committed
8
import os
Phuong Nguyen's avatar
Phuong Nguyen committed
9
import time
Tim Moon's avatar
Tim Moon committed
10
from pathlib import Path
11
from typing import List, Tuple
Przemek Tredak's avatar
Przemek Tredak committed
12

Tim Moon's avatar
Tim Moon committed
13
import setuptools
Phuong Nguyen's avatar
Phuong Nguyen committed
14
from wheel.bdist_wheel import bdist_wheel
Przemek Tredak's avatar
Przemek Tredak committed
15

16
from build_tools.build_ext import CMakeExtension, get_build_ext
17
from build_tools.te_version import te_version
18
from build_tools.utils import (
19
    cuda_archs,
20
    cuda_version,
21
    get_frameworks,
22
    remove_dups,
23
)
Przemek Tredak's avatar
Przemek Tredak committed
24

25
26
frameworks = get_frameworks()
current_file_path = Path(__file__).parent.resolve()
Przemek Tredak's avatar
Przemek Tredak committed
27
28


29
from setuptools.command.build_ext import build_ext as BuildExtension
30

31
32
os.environ["NVTE_PROJECT_BUILDING"] = "1"

33
34
35
36
if "pytorch" in frameworks:
    from torch.utils.cpp_extension import BuildExtension
elif "jax" in frameworks:
    from pybind11.setup_helpers import build_ext as BuildExtension
Tim Moon's avatar
Tim Moon committed
37
38


39
CMakeBuildExtension = get_build_ext(BuildExtension)
40
archs = cuda_archs()
41

Tim Moon's avatar
Tim Moon committed
42

Phuong Nguyen's avatar
Phuong Nguyen committed
43
44
45
46
47
48
49
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
50
        print(f"Total time for bdist_wheel: {total_time:.2f} seconds")
Phuong Nguyen's avatar
Phuong Nguyen committed
51
52


53
def setup_common_extension() -> CMakeExtension:
54
    """Setup CMake extension for common library"""
55
    cmake_flags = ["-DCMAKE_CUDA_ARCHITECTURES={}".format(archs)]
56
57
58
59
60
61
    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")

62
63
64
65
66
67
    if bool(int(os.getenv("NVTE_ENABLE_NVSHMEM", "0"))):
        assert (
            os.getenv("NVSHMEM_HOME") is not None
        ), "NVSHMEM_HOME must be set when compiling with NVTE_ENABLE_NVSHMEM=1"
        cmake_flags.append("-DNVTE_ENABLE_NVSHMEM=ON")

68
69
70
    if bool(int(os.getenv("NVTE_BUILD_ACTIVATION_WITH_FAST_MATH", "0"))):
        cmake_flags.append("-DNVTE_BUILD_ACTIVATION_WITH_FAST_MATH=ON")

71
72
73
    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(
74
75
            f"nvidia-cublasmp-cu{cuda_version()[0]}"
        ).locate_file(f"nvidia/cublasmp/cu{cuda_version()[0]}")
76
77
        cmake_flags.append(f"-DCUBLASMP_DIR={cublasmp_dir}")
        nvshmem_dir = os.getenv("NVSHMEM_HOME") or metadata.distribution(
78
            f"nvidia-nvshmem-cu{cuda_version()[0]}"
79
80
81
82
        ).locate_file("nvidia/nvshmem")
        cmake_flags.append(f"-DNVSHMEM_DIR={nvshmem_dir}")
        print("CMAKE_FLAGS:", cmake_flags[-2:])

83
84
85
86
87
    # Add custom CMake arguments from environment variable
    nvte_cmake_extra_args = os.getenv("NVTE_CMAKE_EXTRA_ARGS")
    if nvte_cmake_extra_args:
        cmake_flags.extend(nvte_cmake_extra_args.split())

88
89
    # Project directory root
    root_path = Path(__file__).resolve().parent
90

91
92
    return CMakeExtension(
        name="transformer_engine",
93
        cmake_path=root_path / Path("transformer_engine/common"),
94
        cmake_flags=cmake_flags,
Tim Moon's avatar
Tim Moon committed
95
96
97
    )


98
def setup_requirements() -> Tuple[List[str], List[str]]:
Tim Moon's avatar
Tim Moon committed
99
100
    """Setup Python dependencies

101
    Returns dependencies for runtime and testing.
Tim Moon's avatar
Tim Moon committed
102
103
104
    """

    # Common requirements
105
106
    install_reqs: List[str] = [
        "pydantic",
107
        "importlib-metadata>=1.0",
108
        "packaging",
109
    ]
110
    test_reqs: List[str] = ["pytest>=8.2.1"]
Tim Moon's avatar
Tim Moon committed
111

112
113
114
    # Framework-specific requirements
    if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))):
        if "pytorch" in frameworks:
115
116
117
118
            from build_tools.pytorch import install_requirements, test_requirements

            install_reqs.extend(install_requirements())
            test_reqs.extend(test_requirements())
119
        if "jax" in frameworks:
120
121
122
123
            from build_tools.jax import install_requirements, test_requirements

            install_reqs.extend(install_requirements())
            test_reqs.extend(test_requirements())
124

125
    return [remove_dups(reqs) for reqs in [install_reqs, test_reqs]]
Tim Moon's avatar
Tim Moon committed
126

127

128
129
if __name__ == "__main__":
    __version__ = te_version()
Przemek Tredak's avatar
Przemek Tredak committed
130

131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
    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 = []
        package_data = {}
        include_package_data = False
        install_requires = ([f"transformer_engine_cu12=={__version__}"],)
        extras_require = {
            "pytorch": [f"transformer_engine_torch=={__version__}"],
            "jax": [f"transformer_engine_jax=={__version__}"],
        }
    else:
148
        install_requires, test_requires = setup_requirements()
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
        ext_modules = [setup_common_extension()]
        package_data = {"": ["VERSION.txt"]}
        include_package_data = True
        extras_require = {"test": test_requires}

        if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))):
            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",
                    )
164
                )
165
166
167
168
169
170
171
172
173
            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",
                    )
174
                )
175

Tim Moon's avatar
Tim Moon committed
176
177
178
    # Configure package
    setuptools.setup(
        name="transformer_engine",
179
180
        version=__version__,
        packages=setuptools.find_packages(
181
182
183
184
185
            include=[
                "transformer_engine",
                "transformer_engine.*",
                "transformer_engine/build_tools",
            ],
186
        ),
187
        extras_require=extras_require,
Tim Moon's avatar
Tim Moon committed
188
        description="Transformer acceleration library",
189
190
        long_description=long_description,
        long_description_content_type="text/x-rst",
Tim Moon's avatar
Tim Moon committed
191
        ext_modules=ext_modules,
Phuong Nguyen's avatar
Phuong Nguyen committed
192
        cmdclass={"build_ext": CMakeBuildExtension, "bdist_wheel": TimedBdist},
193
194
        python_requires=">=3.8",
        classifiers=["Programming Language :: Python :: 3"],
Tim Moon's avatar
Tim Moon committed
195
196
        install_requires=install_requires,
        license_files=("LICENSE",),
197
198
        include_package_data=include_package_data,
        package_data=package_data,
Tim Moon's avatar
Tim Moon committed
199
    )