[ROCm] Allowing for more flexibility in matching ROCm-specific PyTorch wheels to the installed ROCm version

[nikhil-tensorwave]

Feature request

When compiling bitsandbytes from source on ROCm, it should allow for having a ROCm-specific PyTorch wheel within 1 minor version; i.e. if ROCm 6.3.x is installed, then any of pytorch+rocm 6.2.x, 6.3.x, or 6.4.x should work.

Motivation

PyTorch wheels lag behind ROCm releases, especially for stable versions. Additionally, there are very significant performance improvements made to nightly builds of PyTorch, which also tend to lag behind ROCm releases. ROCm 6.3 has been released with additional improvements, but bitsandbytes no longer works due to the mismatch of installed ROCm and the PyTorch wheel ROCm version. Allowing for more flexibility would help with continued work with bitsandbytes on ROCm 6.3.

Your contribution

We can submit a PR making this change and additionally test compilation and performance on our MI300X servers. Please help direct us to where in the codebase this check for matching ROCm versions happens.

[matthewdouglas]

cc: @pnunna93

[pnunna93]

Please check this function - https://github.com/bitsandbytes-foundation/bitsandbytes/blob/multi-backend-refactor/bitsandbytes/cextension.py#L72

It looks for the compiled bitsandbytes .so file that matches with rocm version from torch.