Move to small wheel approach for CUDA SBSA wheel #160720

tinglvv · 2025-08-15T07:01:14Z

Use download.pytorch.org's dependencies like x86 build instead of bundling libs into the wheel

cc @seemethere @malfet @atalman @ptrblck @msaroufim @eqy @jerryzh168 @snadampal @milpuz01 @aditew01 @nikhil-arm @fadara01 @nWEIdia

pytorch-bot · 2025-08-15T07:01:17Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/160720

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

Note: Links to docs will display an error until the docs builds have been completed.

⏳ 9 Pending, 16 Unrelated Failures

As of commit 3f3d1b2 with merge base 8906266 ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

linux-binary-libtorch / libtorch-rocm6_3-shared-with-deps-release-test (gh) (similar failure)
Can't find 'action.yml', 'action.yaml' or 'Dockerfile' under '/home/runner/_work/pytorch/pytorch/.github/actions/setup-rocm'. Did you forget to run actions/checkout before running your local action?
linux-binary-libtorch / libtorch-rocm6_4-shared-with-deps-release-test (gh) (similar failure)
Can't find 'action.yml', 'action.yaml' or 'Dockerfile' under '/home/runner/_work/pytorch/pytorch/.github/actions/setup-rocm'. Did you forget to run actions/checkout before running your local action?
linux-binary-manywheel / manywheel-py3_10-rocm6_3-test (gh) (similar failure)
Can't find 'action.yml', 'action.yaml' or 'Dockerfile' under '/home/runner/_work/pytorch/pytorch/.github/actions/setup-rocm'. Did you forget to run actions/checkout before running your local action?
linux-binary-manywheel / manywheel-py3_10-rocm6_4-test (gh) (similar failure)
Can't find 'action.yml', 'action.yaml' or 'Dockerfile' under '/home/runner/_work/pytorch/pytorch/.github/actions/setup-rocm'. Did you forget to run actions/checkout before running your local action?
linux-binary-manywheel / manywheel-py3_11-rocm6_3-test (gh) (similar failure)
Can't find 'action.yml', 'action.yaml' or 'Dockerfile' under '/home/runner/_work/pytorch/pytorch/.github/actions/setup-rocm'. Did you forget to run actions/checkout before running your local action?
linux-binary-manywheel / manywheel-py3_11-rocm6_4-test (gh) (similar failure)
Can't find 'action.yml', 'action.yaml' or 'Dockerfile' under '/home/runner/_work/pytorch/pytorch/.github/actions/setup-rocm'. Did you forget to run actions/checkout before running your local action?
linux-binary-manywheel / manywheel-py3_12-rocm6_3-test (gh) (similar failure)
Can't find 'action.yml', 'action.yaml' or 'Dockerfile' under '/home/runner/_work/pytorch/pytorch/.github/actions/setup-rocm'. Did you forget to run actions/checkout before running your local action?
linux-binary-manywheel / manywheel-py3_12-rocm6_4-test (gh) (similar failure)
Can't find 'action.yml', 'action.yaml' or 'Dockerfile' under '/home/runner/_work/pytorch/pytorch/.github/actions/setup-rocm'. Did you forget to run actions/checkout before running your local action?
linux-binary-manywheel / manywheel-py3_13-rocm6_3-test (gh) (similar failure)
Can't find 'action.yml', 'action.yaml' or 'Dockerfile' under '/home/runner/_work/pytorch/pytorch/.github/actions/setup-rocm'. Did you forget to run actions/checkout before running your local action?
linux-binary-manywheel / manywheel-py3_13-rocm6_4-test (gh) (similar failure)
Can't find 'action.yml', 'action.yaml' or 'Dockerfile' under '/home/runner/_work/pytorch/pytorch/.github/actions/setup-rocm'. Did you forget to run actions/checkout before running your local action?
linux-binary-manywheel / manywheel-py3_13t-rocm6_3-test (gh) (similar failure)
Can't find 'action.yml', 'action.yaml' or 'Dockerfile' under '/home/runner/_work/pytorch/pytorch/.github/actions/setup-rocm'. Did you forget to run actions/checkout before running your local action?
linux-binary-manywheel / manywheel-py3_13t-rocm6_4-test (gh) (similar failure)
Can't find 'action.yml', 'action.yaml' or 'Dockerfile' under '/home/runner/_work/pytorch/pytorch/.github/actions/setup-rocm'. Did you forget to run actions/checkout before running your local action?
linux-binary-manywheel / manywheel-py3_14-rocm6_3-test (gh) (similar failure)
Can't find 'action.yml', 'action.yaml' or 'Dockerfile' under '/home/runner/_work/pytorch/pytorch/.github/actions/setup-rocm'. Did you forget to run actions/checkout before running your local action?
linux-binary-manywheel / manywheel-py3_14-rocm6_4-test (gh) (similar failure)
Can't find 'action.yml', 'action.yaml' or 'Dockerfile' under '/home/runner/_work/pytorch/pytorch/.github/actions/setup-rocm'. Did you forget to run actions/checkout before running your local action?
linux-binary-manywheel / manywheel-py3_14t-rocm6_3-test (gh) (similar failure)
Can't find 'action.yml', 'action.yaml' or 'Dockerfile' under '/home/runner/_work/pytorch/pytorch/.github/actions/setup-rocm'. Did you forget to run actions/checkout before running your local action?
linux-binary-manywheel / manywheel-py3_14t-rocm6_4-test (gh) (similar failure)
Can't find 'action.yml', 'action.yaml' or 'Dockerfile' under '/home/runner/_work/pytorch/pytorch/.github/actions/setup-rocm'. Did you forget to run actions/checkout before running your local action?

This comment was automatically generated by Dr. CI and updates every 15 minutes.

Aidyn-A · 2025-08-15T07:11:36Z

.ci/aarch64_linux/aarch64_wheel_ci_build.py

+            "/usr/local/lib/libnvpl_lapack_lp64_gomp.so.0",
+            "/usr/local/lib/libnvpl_blas_lp64_gomp.so.0",
+            "/usr/local/lib/libnvpl_lapack_core.so.0",
+            "/usr/local/lib/libnvpl_blas_core.so.0",


There are nvpl-blas nvpl-lapack in PyPi, so we can utilize those as well.

Thanks Aidyn! let me try pulling those two as well.

Might have to do it on a later stage, haven't tested nvpl-blas and nvpl-lapack with pypi before.

tinglvv · 2025-08-15T08:22:58Z

@pytorchbot rebase

pytorchmergebot · 2025-08-15T08:24:24Z

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

pytorchmergebot · 2025-08-15T08:24:26Z

Successfully rebased sbsa-small-wheel onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout sbsa-small-wheel && git pull --rebase)

.ci/aarch64_linux/aarch64_ci_build.sh

tinglvv · 2025-09-04T17:55:22Z

Build for sbsa build was successful. If just install through pip install without any extra instruction, would default install from NCCL on PyPI, which would encounter error and face security risks. Should use pip install torch-2.9.0.dev20250904+cu130-cp312-cp312-manylinux_2_28_aarch64.whl --index-url https://download.pytorch.org/whl/nightly/cu130

root@03cd90fca0f1:/tmp/artifacts# pip install torch-2.9.0.dev20250904+cu130-cp312-cp312-manylinux_2_28_aarch64.whl 
...
INFO: pip is looking at multiple versions of torch to determine which version is compatible with other requirements. This could take a while.
ERROR: Could not find a version that satisfies the requirement nvidia-nccl-cu13==2.27.7; platform_system == "Linux" and platform_machine == "aarch64" (from torch) (from versions: 0.0.0a0)

[notice] A new release of pip is available: 25.0.1 -> 25.2
[notice] To update, run: pip install --upgrade pip
ERROR: No matching distribution found for nvidia-nccl-cu13==2.27.7; platform_system == "Linux" and platform_machine == "aarch64"

tinglvv · 2025-09-08T09:10:02Z

After successful installation, seeing this error

Traceback (most recent call last):
  File "/usr/local/lib/python3.12/site-packages/torch/__init__.py", line 323, in _load_global_deps
    ctypes.CDLL(global_deps_lib_path, mode=ctypes.RTLD_GLOBAL)
  File "/usr/local/lib/python3.12/ctypes/__init__.py", line 379, in __init__
    self._handle = _dlopen(self._name, mode)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^
OSError: libcudart.so.13: cannot open shared object file: No such file or directory

1st check: RPATH not set

readelf -d /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch.so 

Dynamic section at offset 0x3fc78 contains 31 entries:
  Tag        Type                         Name/Value
 0x0000000000000001 (NEEDED)             Shared library: [libtorch_cpu.so]
 0x0000000000000001 (NEEDED)             Shared library: [libtorch_cuda.so]
 0x0000000000000001 (NEEDED)             Shared library: [libgcc_s.so.1]
 0x0000000000000001 (NEEDED)             Shared library: [libpthread.so.0]
 0x0000000000000001 (NEEDED)             Shared library: [libc.so.6]
 0x0000000000000001 (NEEDED)             Shared library: [ld-linux-aarch64.so.1]
 0x000000000000000e (SONAME)             Library soname: [libtorch.so]
 0x000000000000001d (RUNPATH)            Library runpath: [$ORIGIN:/usr/local/cuda/lib64]

2nd check:
Through LD_DEBUG, found that the libtorch_global_deps.so need to be patched with RPATH fix, instead of just libtorch.so and libtorch_python.so:

root@b4c39583b159:/usr/local/lib/python3.12/site-packages/torch/lib# LD_DEBUG=libs python3 -c "import torch" 2>&1 | grep -A5 -B5 cudart
       409:     calling init: /usr/local/lib/python3.12/lib-dynload/_datetime.cpython-312-aarch64-linux-gnu.so
       409:
       409:
       409:     calling init: /usr/local/lib/python3.12/lib-dynload/array.cpython-312-aarch64-linux-gnu.so
       409:
       409:     find library=libcudart.so.13 [0]; searching
       409:      search path=/usr/local/lib/python3.12/site-packages/torch/lib:/usr/local/cuda/lib64               (RUNPATH from file /usr/local/lib/python3.12/site-packages/torch/lib/libtorch_global_deps.so)
       409:       trying file=/usr/local/lib/python3.12/site-packages/torch/lib/libcudart.so.13
       409:       trying file=/usr/local/cuda/lib64/libcudart.so.13
       409:      search cache=/etc/ld.so.cache
       409:      search path=/lib/aarch64-linux-gnu:/usr/lib/aarch64-linux-gnu:/lib:/usr/lib               (system search path)
       409:       trying file=/lib/aarch64-linux-gnu/libcudart.so.13
       409:       trying file=/usr/lib/aarch64-linux-gnu/libcudart.so.13
       409:       trying file=/lib/libcudart.so.13
       409:       trying file=/usr/lib/libcudart.so.13
       409:
Traceback (most recent call last):
  File "/usr/local/lib/python3.12/site-packages/torch/__init__.py", line 323, in _load_global_deps
    ctypes.CDLL(global_deps_lib_path, mode=ctypes.RTLD_GLOBAL)
  File "/usr/local/lib/python3.12/ctypes/__init__.py", line 379, in __init__
    self._handle = _dlopen(self._name, mode)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^
OSError: libcudart.so.13: cannot open shared object file: No such file or directory

tinglvv · 2025-09-08T09:19:32Z

After fixing RPATH with , new error

>>> import torch
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/lib/python3.12/site-packages/torch/__init__.py", line 415, in <module>
    from torch._C import *  # noqa: F403
    ^^^^^^^^^^^^^^^^^^^^^^
ImportError: libtorch.so: cannot open shared object file: No such file or directory

Through LD_DEBUG, can see that libtorch_python.so is loading libtorch.so, so need to fix RPATH for libtorch_python.so

root@b4c39583b159:/usr/local/lib/python3.12/site-packages/torch/lib# LD_DEBUG=libs python3 -c "import torch" 2>&1 | grep -A3 -B3 libtorch.so
       424:      search path=/usr/local/lib/python3.12/site-packages/torch/lib             (RPATH from file /usr/local/lib/python3.12/site-packages/torch/lib/libtorch_global_deps.so)
       424:       trying file=/usr/local/lib/python3.12/site-packages/torch/lib/libtorch_python.so
       424:
       424:     find library=libtorch.so [0]; searching

tinglvv · 2025-09-08T18:11:58Z

Need to fix the RPATH for all libs loading pypi deps in

root@57ad84516286:/usr/local/lib/python3.12/site-packages/torch/lib# ls
libarm_compute.so        libnvpl_blas_core.so.0         libtorch.so
libarm_compute_graph.so  libnvpl_blas_lp64_gomp.so.0    libtorch_cpu.so
libc10.so                libnvpl_lapack_core.so.0       libtorch_cuda.so
libc10_cuda.so           libnvpl_lapack_lp64_gomp.so.0  libtorch_cuda_linalg.so
libcaffe2_nvrtc.so       libshm                         libtorch_global_deps.so
libgfortran.so.5         libshm.so                      libtorch_nvshmem.so
libgomp.so.1             libshm_windows                 libtorch_python.so

tinglvv · 2025-09-08T18:24:14Z

Moving from zip to wheel pack will be in Andrey's PR #159481

tinglvv · 2025-09-09T00:11:31Z

Verified 13.0 wheel passes.

Follow up:

generalize the logic for cu13 wheels (to be compatible with future 13.1, etc).
metadata needs change for all cuda dependencies:
Requires-Dist: nvidia-cuda-nvrtc==13.0.48; platform_system == "Linux" and platform_machine == "aarch64" ->
Requires-Dist: nvidia-cuda-nvrtc==13.0.48; platform_system == "Linux"

atalman · 2025-09-09T00:16:44Z

@pytorchmergebot merge -f "lint and build looks good"

pytorchmergebot · 2025-09-09T00:18:15Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

huydhn · 2025-09-11T01:30:22Z

I'm also seeing an import error trying to load libarm_compute.so when building domain library, i.e. https://github.com/pytorch/vision/actions/runs/17612462583/job/50037380520#step:15:61

import torch
  File "/__w/_temp/conda_environment_17612462583/lib/python3.10/site-packages/torch/__init__.py", line 415, in <module>
    from torch._C import *  # noqa: F403
ImportError: libarm_compute.so: cannot open shared object file: No such file or directory

It's probably coming from this PR. Any thoughts? @Aidyn-A @tinglvv

Aidyn-A · 2025-09-11T01:35:08Z

I'm also seeing an import error trying to load libarm_compute.so when building domain library, i.e. https://github.com/pytorch/vision/actions/runs/17612462583/job/50037380520#step:15:61
import torch

  File "/__w/_temp/conda_environment_17612462583/lib/python3.10/site-packages/torch/__init__.py", line 415, in <module>

    from torch._C import *  # noqa: F403

ImportError: libarm_compute.so: cannot open shared object file: No such file or directory
It's probably coming from this PR. Any thoughts? @Aidyn-A @tinglvv

Yes, we are aware of that. @atalman has fixed it in #162566

pytorch#160673 Use download.pytorch.org's dependencies like x86 build instead of bundling libs into the wheel Pull Request resolved: pytorch#160720 Approved by: https://github.com/atalman

tinglvv added the ciflow/binaries Trigger all binary build and upload jobs on the PR label Aug 15, 2025

pytorchbot added the open source label Aug 15, 2025

Aidyn-A reviewed Aug 15, 2025

View reviewed changes

tinglvv changed the title ~~Move to small wheel approach for CUDA SBSA wheel~~ [WIP] Move to small wheel approach for CUDA SBSA wheel Aug 15, 2025

pytorchmergebot force-pushed the sbsa-small-wheel branch from a374266 to fa75e50 Compare August 15, 2025 08:24

Aidyn-A reviewed Aug 15, 2025

View reviewed changes

.ci/aarch64_linux/aarch64_ci_build.sh Outdated Show resolved Hide resolved

ptrblck added this to PyTorch + CUDA Aug 22, 2025

ptrblck moved this to In Progress in PyTorch + CUDA Aug 22, 2025

tinglvv added 4 commits September 3, 2025 23:27

pull from pypi index

f8b1958

Set EXTRA_INSTALL_REQUIREMENTS to aarch64

3f63a31

typo

9990b8e

Update CUDA 13 related changes

d570a4e

tinglvv force-pushed the sbsa-small-wheel branch from 98a3ee2 to d570a4e Compare September 4, 2025 06:34

fix lint

0c98ec6

tinglvv marked this pull request as ready for review September 4, 2025 17:52

tinglvv requested a review from a team as a code owner September 4, 2025 17:52

tinglvv added 2 commits September 7, 2025 22:32

Fix RPATH for libtorch.so

84966c9

fix lint

e5f16be

Fix the search path with RPATH

74ef78a

Do not disable NVSHMEM

3299310

Patch the rest of the libs

3f3d1b2

atalman changed the title ~~[WIP] Move to small wheel approach for CUDA SBSA wheel~~ Move to small wheel approach for CUDA SBSA wheel Sep 8, 2025

atalman approved these changes Sep 8, 2025

View reviewed changes

pytorchmergebot added the merging label Sep 9, 2025

pytorchmergebot added the Merged label Sep 9, 2025

pytorchmergebot closed this in 897c4e7 Sep 9, 2025

github-project-automation bot moved this from In Progress to Done in PyTorch + CUDA Sep 9, 2025

pytorchmergebot removed the merging label Sep 9, 2025

Aidyn-A mentioned this pull request Sep 10, 2025

[WIP] Small Wheel on SBSA: fix broken paths #162572

Closed

tinglvv mentioned this pull request Sep 23, 2025

[aarch64] Reduce CUDA SBSA Wheel Size #160673

Closed

atalman removed this from PyTorch + CUDA Sep 26, 2025

Move to small wheel approach for CUDA SBSA wheel #160720

Move to small wheel approach for CUDA SBSA wheel #160720

Uh oh!

Conversation

tinglvv commented Aug 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Aug 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/160720

⏳ 9 Pending, 16 Unrelated Failures

Uh oh!

Aidyn-A Aug 15, 2025

Choose a reason for hiding this comment

Uh oh!

tinglvv Aug 15, 2025

Choose a reason for hiding this comment

Uh oh!

tinglvv Sep 4, 2025

Choose a reason for hiding this comment

Uh oh!

tinglvv commented Aug 15, 2025

Uh oh!

pytorchmergebot commented Aug 15, 2025

Uh oh!

pytorchmergebot commented Aug 15, 2025

Uh oh!

Uh oh!

tinglvv commented Sep 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tinglvv commented Sep 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tinglvv commented Sep 8, 2025

Uh oh!

tinglvv commented Sep 8, 2025

Uh oh!

tinglvv commented Sep 8, 2025

Uh oh!

tinglvv commented Sep 9, 2025 • edited by atalman Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

atalman commented Sep 9, 2025

Uh oh!

pytorchmergebot commented Sep 9, 2025

Merge started

Uh oh!

huydhn commented Sep 11, 2025

Uh oh!

Aidyn-A commented Sep 11, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

tinglvv commented Aug 15, 2025 •

edited

Loading

pytorch-bot bot commented Aug 15, 2025 •

edited

Loading

tinglvv commented Sep 4, 2025 •

edited

Loading

tinglvv commented Sep 8, 2025 •

edited

Loading

tinglvv commented Sep 9, 2025 •

edited by atalman

Loading