From ec2443860ac6cd26eddfdd458a76f84cfd1b46fa Mon Sep 17 00:00:00 2001 From: Yuanyuan Chen Date: Fri, 23 Oct 2026 05:11:04 +0000 Subject: [PATCH] Load ROCm from the SDK wheels the torch wheel was built against ROCm wheels built against an unpacked TheRock SDK (PYTORCH_ROCM_EXTERNAL_SDK) declared no ROCm dependency and shipped no torch/_rocm_init.py, so `import torch` resolved librocm-openblas and the HIP libraries only through LD_LIBRARY_PATH or whatever ROCm the host's ld.so.conf exposed. On a Breadboard guest that is the accelerator pack's ROCm 10.0, not the 10.1 SDK the wheel was built and patched against; without LD_LIBRARY_PATH the import fails on librocm-openblas.so.0. The build now takes the SDK wheels' version from PYTORCH_ROCM_SDK_VERSION. tools/metadata/dependencies.py pins rocm[libraries] to it, and torch/CMakeLists.txt writes the _rocm_init.py that the existing install rule and torch/__init__.py already expect, so `import torch` preloads the SDK libraries from those wheels through rocm_sdk.initialize_process and checks the version. The preload list is TheRock's build_prod_wheels.py list minus hipdnn, which this SDK does not build. PYTORCH_ROCM_EXTERNAL_SDK is gone; the alternative of keeping it and fixing the guest's ld.so.conf would leave torch tied to whatever ROCm the host provides. Test Plan: ``` cmake -P test.cmake # the new block, with and without PYTORCH_ROCM_SDK_VERSION python -c 'import dependencies; ...' # rocm[libraries]== env -u LD_LIBRARY_PATH python -c 'import torch; ...' ``` With the generated _rocm_init.py added to the published gdb71170c wheel on a gfx1201 sandbox, `import torch` succeeds without LD_LIBRARY_PATH and /proc/self/maps shows libamdhip64, libhsa-runtime64, librocblas, libMIOpen and librocm-openblas loaded from the rocm_sdk wheels. A full wheel build with the change was not run. Authored with an AI assistant (Claude Code). Co-Authored-By: Claude Opus 5.5 --- tools/metadata/dependencies.py | 18 ++++++++++-------- torch/CMakeLists.txt | 18 ++++++++++++++++++ 2 files changed, 28 insertions(+), 8 deletions(-) diff --git a/tools/metadata/dependencies.py b/tools/metadata/dependencies.py index 0b1be2acf01ea..5696acfde992a 100644 --- a/tools/metadata/dependencies.py +++ b/tools/metadata/dependencies.py @@ -61,14 +61,16 @@ def dynamic_metadata( extra = os.environ.get("PYTORCH_EXTRA_INSTALL_REQUIREMENTS") if extra: deps.extend(r.strip() for r in extra.split("|") if r.strip()) - elif _is_truthy(os.environ.get("USE_ROCM")) and not _is_truthy( - os.environ.get("PYTORCH_ROCM_EXTERNAL_SDK") - ): - try: - rocm_version = version("rocm") - except PackageNotFoundError: - pass - else: + elif _is_truthy(os.environ.get("USE_ROCM")): + # Pin the ROCm SDK wheels the build used: PYTORCH_ROCM_SDK_VERSION when + # building against an unpacked SDK, else the installed rocm package. + rocm_version = os.environ.get("PYTORCH_ROCM_SDK_VERSION") + if not rocm_version: + try: + rocm_version = version("rocm") + except PackageNotFoundError: + rocm_version = None + if rocm_version: deps.append(f"rocm[libraries]=={rocm_version}") # cupti-python has no free-threaded wheels; PEP 508 cannot express this ABI constraint. diff --git a/torch/CMakeLists.txt b/torch/CMakeLists.txt index e2479ab4f98e0..953d7ce7e3ccd 100644 --- a/torch/CMakeLists.txt +++ b/torch/CMakeLists.txt @@ -466,6 +466,24 @@ endif() # _rocm_init.py is gitignored (injected by ROCm build scripts to bootstrap # torch/__init__.py). See https://github.com/ROCm/TheRock/blob/main/docs/packaging/python_packaging.md +# When building against the ROCm SDK wheels, PYTORCH_ROCM_SDK_VERSION names +# their version: generate the file so `import torch` preloads the SDK's +# libraries from those wheels and checks that the installed SDK matches. +if(USE_ROCM AND NOT "$ENV{PYTORCH_ROCM_SDK_VERSION}" STREQUAL "") + set(_rocm_preloads + amd_comgr amd_smi amdhip64 hiprtc hipblas hipfft hiprand hipsparse + hipsparselt hipsolver hipblaslt miopen rocm_sysdeps_liblzma rocm-openblas + rocprofiler-sdk rocprofiler-sdk-roctx roctracer64 roctx64 rccl hipfile) + list(TRANSFORM _rocm_preloads PREPEND "'") + list(TRANSFORM _rocm_preloads APPEND "'") + list(JOIN _rocm_preloads ", " _rocm_preloads) + file(WRITE "${TORCH_SRC_DIR}/_rocm_init.py" + "def initialize():\n" + " import rocm_sdk\n" + " rocm_sdk.initialize_process(\n" + " preload_shortnames=[${_rocm_preloads}],\n" + " check_version='$ENV{PYTORCH_ROCM_SDK_VERSION}')\n") +endif() # scikit-build-core won't include it via Python package collection. Install it # explicitly via cmake instead. if(SKBUILD AND USE_ROCM AND EXISTS "${TORCH_SRC_DIR}/_rocm_init.py")