Installation overwrites PyTorch version on Jetson AGX Orin

[Andymulb]

🐛 Bug

Hi, I just tried to install xformers on an NVIDIA Jetson AGX Orin with Ubuntu 22.04, Jetpack 6.1 and CUDA 12.6. It requires the specific PyTorch version 2.5.0a0+872d972e41.nv24.08. I'm building xformers from source as described in the README. However, after the build process finished, it uninstalls PyTorch 2.5.0a0+872d972e41.nv24.08 and installs PyTorch 2.5.1.

python -m xformers.info
WARNING[XFORMERS]: xFormers can't load C++/CUDA extensions. xFormers was built for:
    PyTorch 2.5.0a0+872d972e41.nv24.08 with CUDA 1206 (you have 2.5.1)
    Python  3.10.16 (you have 3.10.16)
  Please reinstall xformers (see https://github.com/facebookresearch/xformers#installing-xformers)
  Memory-efficient attention, SwiGLU, sparse and more won't be available.
  Set XFORMERS_MORE_DETAILS=1 for more details
xFormers 0.0.30+6440945.d20250113
memory_efficient_attention.ckF:                    unavailable
memory_efficient_attention.ckB:                    unavailable
memory_efficient_attention.ck_decoderF:            unavailable
memory_efficient_attention.ck_splitKF:             unavailable
memory_efficient_attention.cutlassF-pt:            available
memory_efficient_attention.cutlassB-pt:            available
[email protected]:         available
[email protected]:         available
[email protected]:             unavailable
[email protected]:             unavailable
memory_efficient_attention.triton_splitKF:         available
indexing.scaled_index_addF:                        unavailable
indexing.scaled_index_addB:                        unavailable
indexing.index_select:                             unavailable
sequence_parallel_fused.write_values:              unavailable
sequence_parallel_fused.wait_values:               unavailable
sequence_parallel_fused.cuda_memset_32b_async:     unavailable
sp24.sparse24_sparsify_both_ways:                  unavailable
sp24.sparse24_apply:                               unavailable
sp24.sparse24_apply_dense_output:                  unavailable
sp24._sparse24_gemm:                               unavailable
[email protected]:                 available
[email protected]:                        available
swiglu.dual_gemm_silu:                             unavailable
swiglu.gemm_fused_operand_sum:                     unavailable
swiglu.fused.p.cpp:                                not built
is_triton_available:                               False
pytorch.version:                                   2.5.1
pytorch.cuda:                                      not available
dcgm_profiler:                                     unavailable
build.info:                                        available
build.cuda_version:                                1206
build.hip_version:                                 None
build.python_version:                              3.10.16
build.torch_version:                               2.5.0a0+872d972e41.nv24.08
build.env.TORCH_CUDA_ARCH_LIST:                    8.7
build.env.PYTORCH_ROCM_ARCH:                       None
build.env.XFORMERS_BUILD_TYPE:                     None
build.env.XFORMERS_ENABLE_DEBUG_ASSERTIONS:        None
build.env.NVCC_FLAGS:                              None
build.env.XFORMERS_PACKAGE_FROM:                   None
source.privacy:                                    open source

Environment

Collecting environment information...
PyTorch version: 2.5.0a0+872d972e41.nv24.08
Is debug build: False
CUDA used to build PyTorch: 12.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.5 LTS (aarch64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.31.4
Libc version: glibc-2.35

Python version: 3.10.16 (main, Dec 11 2024, 16:18:56) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.148-tegra-aarch64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.6.85
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: Orin (nvgpu)
Nvidia driver version: 540.4.0
cuDNN version: Probably one of the following:
/usr/lib/aarch64-linux-gnu/libcudnn.so.9.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_adv.so.9.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_cnn.so.9.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_engines_precompiled.so.9.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_graph.so.9.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_heuristic.so.9.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_ops.so.9.6.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       aarch64
CPU op-mode(s):                     32-bit, 64-bit
Byte Order:                         Little Endian
CPU(s):                             12
On-line CPU(s) list:                0-11
Vendor ID:                          ARM
Model name:                         Cortex-A78AE
Model:                              1
Thread(s) per core:                 1
Core(s) per cluster:                4
Socket(s):                          -
Cluster(s):                         3
Stepping:                           r0p1
CPU max MHz:                        2201,6001
CPU min MHz:                        115,2000
BogoMIPS:                           62.50
Flags:                              fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm lrcpc dcpop asimddp uscat ilrcpc flagm paca pacg
L1d cache:                          768 KiB (12 instances)
L1i cache:                          768 KiB (12 instances)
L2 cache:                           3 MiB (12 instances)
L3 cache:                           6 MiB (3 instances)
NUMA node(s):                       1
NUMA node0 CPU(s):                  0-11
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:           Mitigation; __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; CSV2, but not BHB
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.5.0a0+872d972e41.nv24.8
[pip3] torchaudio==0.13.1
[pip3] torchvision==0.14.1
[pip3] torchviz==0.0.2
[pip3] triton==3.2.0+git7cc6799d
[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] torch                     2.5.0a0+872d972e41.nv24.8          pypi_0    pypi
[conda] torchvision               0.14.1                   pypi_0    pypi
[conda] triton                    3.2.0+git7cc6799d          pypi_0    pypi

• PyTorch Version: 2.5.0a0+872d972e41.nv24.08
• OS: Ubuntu 22.04
• How you installed PyTorch:

pip install numpy==1.26.4
export TORCH_INSTALL=https://developer.download.nvidia.com/compute/redist/jp/v61/pytorch/torch-2.5.0a0+872d972e41.nv24.08.17622132-cp310-cp310-linux_aarch64.whl
pip install $TORCH_INSTALL torchvision torchaudio

• Build command you used: pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers
• Python version: 3.10.16
• CUDA/cuDNN version: 12.6

nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2024 NVIDIA Corporation
Built on Tue_Oct_29_23:53:06_PDT_2024
Cuda compilation tools, release 12.6, V12.6.85
Build cuda_12.6.r12.6/compiler.35059454_0

• GPU models and configuration: Nvidia Jetson AGX Orin 64GB

[johnnynunez]

pip install pytorch xformers --index-url=https://pypi.jetson-ai-lab.dev/jp6/cu126