diff --git a/.ci/gitlab/configs/linux/ci.yaml b/.ci/gitlab/configs/linux/ci.yaml index 9add0f89153..ed8ea2be4a4 100644 --- a/.ci/gitlab/configs/linux/ci.yaml +++ b/.ci/gitlab/configs/linux/ci.yaml @@ -13,6 +13,8 @@ ci: - wrf build-job: tags: [ "spack", "huge" ] + # 12h for long ROCm/ML rebuilds (maximum allowed for AWS runners) + timeout: 720 minutes variables: CI_JOB_SIZE: huge SPACK_BUILD_JOBS: "12" diff --git a/repos/spack_repo/builtin/packages/py_torch/PR152569-Update-spack-includes-2.5.patch b/repos/spack_repo/builtin/packages/py_torch/PR152569-Update-spack-includes-2.5.patch index 2e7a80bcbe8..2c35aafac2f 100644 --- a/repos/spack_repo/builtin/packages/py_torch/PR152569-Update-spack-includes-2.5.patch +++ b/repos/spack_repo/builtin/packages/py_torch/PR152569-Update-spack-includes-2.5.patch @@ -25,10 +25,10 @@ index 9be7f37..39d0f24 100644 endif() diff --git a/cmake/public/LoadHIP.cmake b/cmake/public/LoadHIP.cmake -index 1c0d3a2..e0de4b1 100644 +index 1c0d3a2..83f9f9d 100644 --- a/cmake/public/LoadHIP.cmake +++ b/cmake/public/LoadHIP.cmake -@@ -167,6 +167,10 @@ if(HIP_FOUND) +@@ -167,6 +167,11 @@ if(HIP_FOUND) find_package_and_print_version(hipsolver REQUIRED) find_package_and_print_version(hiprtc REQUIRED) @@ -36,6 +36,7 @@ index 1c0d3a2..e0de4b1 100644 + list(APPEND ROCM_INCLUDE ${rocprim_INCLUDE_DIR}) + list(APPEND ROCM_INCLUDE ${hipcub_INCLUDE_DIR}) + list(APPEND ROCM_INCLUDE ${rocRAND_INCLUDE_DIR}) ++ list(APPEND ROCM_INCLUDE $ENV{AOTRITON_INSTALLED_PREFIX}/include) find_library(PYTORCH_HIP_LIBRARIES amdhip64 HINTS ${ROCM_PATH}/lib) # TODO: miopen_LIBRARIES should return fullpath to the library file, diff --git a/repos/spack_repo/builtin/packages/py_torch/PR152569-Update-spack-includes-2.7.patch b/repos/spack_repo/builtin/packages/py_torch/PR152569-Update-spack-includes-2.7.patch index 4392e00d76a..af876735594 100644 --- a/repos/spack_repo/builtin/packages/py_torch/PR152569-Update-spack-includes-2.7.patch +++ b/repos/spack_repo/builtin/packages/py_torch/PR152569-Update-spack-includes-2.7.patch @@ -1,5 +1,5 @@ diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt -index d2d23b7..620a89f 100644 +index d2d23b7ab65..620a89f65cb 100644 --- a/caffe2/CMakeLists.txt +++ b/caffe2/CMakeLists.txt @@ -1379,13 +1379,6 @@ if(USE_ROCM) @@ -26,7 +26,7 @@ index d2d23b7..620a89f 100644 endif() diff --git a/cmake/public/LoadHIP.cmake b/cmake/public/LoadHIP.cmake -index 58c74dd..d3e1ad4 100644 +index 58c74ddda35..54f96871372 100644 --- a/cmake/public/LoadHIP.cmake +++ b/cmake/public/LoadHIP.cmake @@ -26,12 +26,6 @@ else() @@ -78,7 +78,15 @@ index 58c74dd..d3e1ad4 100644 find_package_and_print_version(amd_comgr REQUIRED) find_package_and_print_version(rocrand REQUIRED) find_package_and_print_version(hiprand REQUIRED) -@@ -171,7 +168,11 @@ if(HIP_FOUND) +@@ -157,6 +154,7 @@ if(HIP_FOUND) + find_package_and_print_version(hipcub REQUIRED) + find_package_and_print_version(rocthrust REQUIRED) + find_package_and_print_version(hipsolver REQUIRED) ++ list(APPEND ROCM_INCLUDE_DIRS $ENV{AOTRITON_INSTALLED_PREFIX}/include) + # workaround cmake 4 build issue + if(CMAKE_VERSION VERSION_GREATER_EQUAL "4.0.0") + message(WARNING "Work around hiprtc cmake failure for cmake >= 4") +@@ -171,7 +169,11 @@ if(HIP_FOUND) if(UNIX) find_package_and_print_version(rccl) find_package_and_print_version(hsa-runtime64 REQUIRED) @@ -89,4 +97,3 @@ index 58c74dd..d3e1ad4 100644 + if(UNIX) # roctx is part of roctracer find_library(ROCM_ROCTX_LIB roctx64 HINTS ${ROCM_PATH}/lib) - diff --git a/repos/spack_repo/builtin/packages/py_torch/package.py b/repos/spack_repo/builtin/packages/py_torch/package.py index 25da2c23fda..80aa14c82e5 100644 --- a/repos/spack_repo/builtin/packages/py_torch/package.py +++ b/repos/spack_repo/builtin/packages/py_torch/package.py @@ -97,7 +97,16 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage): variant("rocm", default=False, description="Use ROCm") variant("cudnn", default=not is_darwin, description="Use cuDNN", when="+cuda") variant("fbgemm", default=True, description="Use FBGEMM (quantized 8-bit server operators)") - variant("kineto", default=True, description="Use Kineto profiling library", when="@1.8:") + variant("kineto", default=True, description="Use Kineto profiling library", when="@1.8:2.10") + variant( + "kineto", default=True, description="Use Kineto profiling library", when="@2.11: ~rocm" + ) + variant( + "kineto", + default=False, + description="Disable Kineto from 2.11 on ROCm", + when="@2.11: +rocm", + ) variant("magma", default=not is_darwin, description="Use MAGMA", when="+cuda") variant("metal", default=is_darwin, description="Use Metal for Caffe2 iOS build") variant( @@ -107,7 +116,13 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage): when="@1.12: platform=darwin", ) variant("nccl", default=True, description="Use NCCL", when="+cuda platform=linux") - variant("nccl", default=True, description="Use NCCL", when="+rocm platform=linux") + variant("nccl", default=True, description="Use NCCL", when="@:2.10 +rocm platform=linux") + variant( + "nccl", + default=False, + description="Disable NCCL from 2.11 on rocm", + when="@2.11: +rocm platform=linux", + ) # Requires AVX2: https://discuss.pytorch.org/t/107518 variant("nnpack", default=True, description="Use NNPACK", when="target=x86_64_v3:") variant("numa", default=True, description="Use NUMA", when="platform=linux") @@ -144,6 +159,9 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage): conflicts("+gloo+rocm") conflicts("+rocm", when="@2.3", msg="Rocm doesn't support py-torch 2.3 release") conflicts("+rocm", when="@2.4", msg="Rocm doesn't support py-torch 2.4 release") + conflicts("+rocm", when="@2.8", msg="Rocm doesn't support py-torch 2.8 release") + conflicts("+rocm", when="@2.12", msg="Rocm doesn't support py-torch 2.12 release") + conflicts("+rocm", when="@2.13", msg="Rocm doesn't support py-torch 2.13 release") conflicts("+tensorpipe", when="+rocm ^hip@:5.1", msg="TensorPipe not supported until ROCm 5.2") conflicts("+breakpad", when="target=ppc64:") conflicts("+breakpad", when="target=ppc64le:") @@ -344,7 +362,8 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage): depends_on("valgrind", when="+valgrind") with when("+rocm"): depends_on("hsa-rocr-dev") - depends_on("hip") + depends_on("hip@7.0:", when="@2.9:") + depends_on("hip@:6.4", when="@:2.7") depends_on("rccl", when="+nccl") depends_on("rocprim") depends_on("hipcub") @@ -356,14 +375,22 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage): depends_on("hiprand") depends_on("hipsolver") depends_on("rocm-core") + depends_on("amdsmi", when="@2.12") depends_on("rocfft") depends_on("rocblas") depends_on("miopen-hip") + depends_on("composable-kernel") + depends_on("hipblaslt") + depends_on("rocm-smi-lib") + # Ensure hipblaslt version for 2.9+ + depends_on("hipblaslt@7.0:", when="@2.9:") depends_on("rocminfo") - depends_on("aotriton@0.8.1b", when="@2.5:2.6") - depends_on("aotriton@0.9.1b", when="@2.7:") - depends_on("composable-kernel@:6.3.2", when="@2.5") - depends_on("composable-kernel@6.3.2:", when="@2.6:") + depends_on("hipsparselt@7.0:", when="@2.9:") + depends_on("aotriton@0.8b", when="@2.5:2.6") + depends_on("aotriton@0.9.2b", when="@2.7") + depends_on("aotriton@0.10b", when="@2.8:2.10") + depends_on("aotriton@0.11b", when="@2.11:") + depends_on("mpi", when="+mpi") depends_on("ucc", when="+ucc") depends_on("ucx", when="+ucc") @@ -627,6 +654,17 @@ def patch(self): "torch_global_deps PROPERTIES LINKER_LANGUAGE CXX", "caffe2/CMakeLists.txt", ) + if self.spec.satisfies("@2.5:+rocm"): + filter_file( + "find_library(ROCM_ROCTX_LIB roctx64 HINTS ${ROCM_PATH}/lib)", + "find_library(ROCM_ROCTX_LIB roctx64 HINTS ${ROCM_PATH}/lib)\n" + "if(DEFINED ENV{ROCTRACER_INCLUDE_DIR})\n" + " set(ROCTRACER_INCLUDE_DIR $ENV{ROCTRACER_INCLUDE_DIR} CACHE PATH " + '"Roctracer include directory" FORCE)\n' + "endif()", + "cmake/public/LoadHIP.cmake", + string=True, + ) if self.spec.satisfies("@2.1:2.7+rocm"): filter_file( "${ROCM_INCLUDE_DIRS}/rocm-core/rocm_version.h", @@ -724,6 +762,19 @@ def enable_or_disable(variant, keyword="USE", var=None): enable_or_disable("rocm") if "+rocm" in self.spec: + # So libtorch_hip.so and dependents find ROCm/runtime libs at runtime and + # during binary cache relocation (avoids "=> not found" for e.g. + # libamdhip64.so.6, libhsa-runtime64.so.1). + for lib_dir in [ + self.spec["hip"].prefix.lib, + self.spec["hsa-rocr-dev"].prefix.lib, + self.spec["rocm-smi-lib"].prefix.lib, + ]: + env.append_flags("LDFLAGS", "-Wl,-rpath," + lib_dir) + # Link to rocm-smi-lib which provides rsmi_* symbols used by libtorch_hip.so + env.append_flags( + "LDFLAGS", "-L{} -lrocm_smi64".format(self.spec["rocm-smi-lib"].prefix.lib) + ) env.set("PYTORCH_ROCM_ARCH", ";".join(self.spec.variants["amdgpu_target"].value)) env.set("HSA_PATH", self.spec["hsa-rocr-dev"].prefix) env.set("ROCBLAS_PATH", self.spec["rocblas"].prefix) @@ -744,6 +795,7 @@ def enable_or_disable(variant, keyword="USE", var=None): if self.spec.satisfies("@2.5:"): env.set("TORCHINDUCTOR_CK_DIR", self.spec["composable-kernel"].prefix) env.set("AOTRITON_INSTALLED_PREFIX", self.spec["aotriton"].prefix) + env.prepend_path("CPATH", self.spec["aotriton"].prefix.include) if self.spec.satisfies("^hip@5.2.0:"): env.set("CMAKE_MODULE_PATH", self.spec["hip"].prefix.lib.cmake.hip) @@ -867,6 +919,22 @@ def enable_or_disable(variant, keyword="USE", var=None): def setup_run_environment(self, env: EnvironmentModifications) -> None: self.torch_cuda_arch_list(env) + if "+rocm" in self.spec: + env.prepend_path("LD_LIBRARY_PATH", self.spec["hip"].prefix.lib) + + def setup_dependent_build_environment(self, env, dependent_spec): + if "+rocm" in self.spec: + env.prepend_path("LD_LIBRARY_PATH", self.spec["hip"].prefix.lib) + # PyTorch headers (e.g. c10/util/complex.h) include ; + # dependents need rocthrust include so HIP device builds can find it. + env.set("THRUST_PATH", self.spec["rocthrust"].prefix) + env.prepend_path("CPATH", self.spec["rocthrust"].prefix.include) + + def setup_dependent_run_environment(self, env, dependent_spec): + """So dependents (e.g. py-torch-nvidia-apex, py-torchaudio) can find + libamdhip64.so when importing torch or running code that uses ROCm.""" + if "+rocm" in self.spec: + env.prepend_path("LD_LIBRARY_PATH", self.spec["hip"].prefix.lib) @run_before("install") def build_amd(self): diff --git a/repos/spack_repo/builtin/packages/py_torchvision/package.py b/repos/spack_repo/builtin/packages/py_torchvision/package.py index 3b46f39a243..ee66a1335a6 100644 --- a/repos/spack_repo/builtin/packages/py_torchvision/package.py +++ b/repos/spack_repo/builtin/packages/py_torchvision/package.py @@ -61,6 +61,9 @@ class PyTorchvision(PythonPackage): version("0.9.0", sha256="9351ed92aded632f8c7f59dfadac13c191a834babe682f5785ea47e6fcf6b472") version("0.8.2", sha256="9a866c3c8feb23b3221ce261e6153fc65a98ce9ceaa71ccad017016945c178bf") + # Fix duplicate symbol error when building with ROCm + patch("torchvision-0.26.0-rocm-vision-duplicate-symbol.patch", when="@0.26.0 ^py-torch+rocm") + desc = "Enable support for native encoding/decoding of {} formats in torchvision.io" variant("png", default=True, description=desc.format("PNG")) variant("jpeg", default=True, description=desc.format("JPEG")) @@ -207,6 +210,28 @@ def setup_build_environment(self, env: EnvironmentModifications) -> None: include.extend(query.headers.directories) library.extend(query.libs.directories) + # When building with ROCm, add all ROCm library include paths for HIP compilation + # PyTorch headers transitively include many ROCm headers that extensions need + if "^py-torch+rocm" in self.spec: + rocm_deps = [ + "rocthrust", + "rocprim", + "hipsparse", + "hipblas", + "hipblas-common", + "hipblaslt", + "hipfft", + "hiprand", + "hipsolver", + "rocblas", + "rocsparse", + "rocsolver", + "rocfft", + ] + for dep in rocm_deps: + if dep in self.spec: + include.append(self.spec[dep].prefix.include) + # CONTRIBUTING.md says to use TORCHVISION_INCLUDE and TORCHVISION_LIBRARY, but # these do not work for older releases. Build uses a mix of Spack's compiler wrapper # and the actual compiler, so this is needed to get parts of the build working. @@ -215,3 +240,10 @@ def setup_build_environment(self, env: EnvironmentModifications) -> None: env.set("TORCHVISION_LIBRARY", ":".join(library)) env.set("CPATH", ":".join(include)) env.set("LIBRARY_PATH", ":".join(library)) + + # For ROCm builds, also prepend ROCm includes to ensure hipcc can find them + if "^py-torch+rocm" in self.spec: + for dep in rocm_deps: + if dep in self.spec: + env.prepend_path("CPATH", self.spec[dep].prefix.include) + env.prepend_path("CPLUS_INCLUDE_PATH", self.spec[dep].prefix.include) diff --git a/repos/spack_repo/builtin/packages/py_torchvision/torchvision-0.26.0-rocm-vision-duplicate-symbol.patch b/repos/spack_repo/builtin/packages/py_torchvision/torchvision-0.26.0-rocm-vision-duplicate-symbol.patch new file mode 100644 index 00000000000..21cd51f0168 --- /dev/null +++ b/repos/spack_repo/builtin/packages/py_torchvision/torchvision-0.26.0-rocm-vision-duplicate-symbol.patch @@ -0,0 +1,23 @@ +diff --git a/setup.py b/setup.py +index 7bf83fc..a9c8e5d 100644 +--- a/setup.py ++++ b/setup.py +@@ -153,6 +153,10 @@ def make_C_extension(): + + list(CSRS_DIR.glob("ops/quantized/cpu/*.cpp")) + ) + mps_sources = list(CSRS_DIR.glob("ops/mps/*.mm")) ++ ++ # Exclude vision_hip.cpp from main sources when building with ROCm to avoid duplicate symbols ++ if IS_ROCM: ++ sources = [s for s in sources if "vision_hip.cpp" not in str(s)] + + if IS_ROCM: + from torch.utils.hipify import hipify_python +@@ -161,6 +165,7 @@ def make_C_extension(): + project_directory=str(ROOT_DIR), + output_directory=str(ROOT_DIR), + includes="torchvision/csrc/ops/cuda/*", ++ ignores=("torchvision/csrc/vision.cpp",), + show_detailed=True, + is_pytorch_extension=True, + ) diff --git a/stacks/ml-linux-x86_64/spack.yaml b/stacks/ml-linux-x86_64/spack.yaml index 50a030f9d9f..530b720931b 100644 --- a/stacks/ml-linux-x86_64/spack.yaml +++ b/stacks/ml-linux-x86_64/spack.yaml @@ -37,7 +37,6 @@ spack: - py-timm - py-torch - py-torch-geometric - - py-torch-nvidia-apex - py-torchaudio - py-torchgeo - py-torchmetrics @@ -84,6 +83,7 @@ spack: - $scikit_learn_specs - $tensorboard_specs - $tensorflow_specs + - py-torch-nvidia-apex # xgboost requires older cuda, which doesn't support newer GCC # - $xgboost_specs override: @@ -113,15 +113,11 @@ spack: - $jax_specs - matrix: - [$keras_specs] - exclude: - - py-keras backend=torch - # Does not yet support Spack-installed ROCm - # - $torch_specs + - $torch_specs - $scikit_learn_specs - $tensorboard_specs - $tensorflow_specs - # Does not yet support Spack-installed ROCm - # - $xgboost_specs + - $xgboost_specs override: packages: all: @@ -139,9 +135,13 @@ spack: - target=x86_64_v3 - ~cuda - +rocm + - ~distributed - amdgpu_target=gfx90a - - ~flash_attention + - +flash_attention - "%c,cxx=gcc" + hwloc: + require: + - ~rocm ci: pipeline-gen: