Implement solution for faster release of PyTorch easyconfigs

[lexming]

The goal of this issue is to discuss a solution in EB to allow for faster development of the PyTorch ecosystem in each toolchain. The two main motivations are:

• there are more and more packages that depend on PyTorch, and the easyconfigs of PyTorch are a bottleneck to start working on those
• building PyTorch is also becoming more and more complex (just look at the number of patches in 2.6), which reduces the number of people that is willing or can tackle the development of easyconfigs for new versions

The initial idea is to have a 2-stage distribution of PyTorch easyconfigs:

• Stage-1: easyconfigs using pre-built binaries of PyTorch from wheels
  easy to update in EB: with usual PythonBundle easyblock installs
  no need for extensive testing: the pre-built binaries are already tested, so we can just check that they work in the host system avoiding testing the thousand features in PyTorch (as we do with builds from source)
  has to integrate with the dependencies in EB
• Stage-2: easyconfigs building PyTorch from source (as usual)
  the same type of installation we currently do where PyTorch is built from source

The key feature for Stage-2 is that it should be transparent to users. This means that there should be some mechanism that allows deploying the optimized PyTorch side by side to the wheel PyTorch and automagically swap the wheel module with the optimized one without reinstalling any of the software depending on that version of PyTorch.

Some preliminary info:

• The wheels from the PyTorch project already cover CPU, CUDA and ROCm; so there is no need to look elsewhere for our purpose.
• In Brussels we already have some experience in doing module swaps with PyTorch. We installed two different versions of PyTorch in the same toolchain without reinstalling any of the packages on top of it. This happened for 2023a, we started installing PyTorch v2.1.2, then the tens of packages needing PyTorch were gradually added on top of it; and months later we installed v2.3.0 and just instructed our users to manually swap the PyTorch modules. That worked very well and we have not received any issue with that approach. However, that was before EB 5 and without RPATH.

[lexming]

I tested the compatibility between upstream wheels from https://download.pytorch.org/whl and our own builds from source of PyTorch and torchvision, and the results show that we cannot easily swap them.

Tests consisted in running the example script from https://github.com/pytorch/examples/tree/main/super_resolution with PyTorch v2.6.0 and torchvision v0.21.0 without accelerators.

Building/Installing from scratch always results in a working stack.

• PyTorch wheel + torchvision wheel : GOOD
• PyTorch wheel + torchvision build from source on top of PyTorch wheel: GOOD
• PyTorch build from source + torchvision build from source on top of PyTorch from source: GOOD

Swapping modules between the previous configuration causes issues though:

• PyTorch build from source + torchvision wheel: BAD
• PyTorch build from source + torchvision build from source on top of PyTorch wheel: BAD
• PyTorch wheel + torchvision build from source on top of PyTorch from source: BAD

All swaps cause the same error in this example script. The following is the traceback for swapping out "PyTorch wheel" for "PyTorch build from source" while keeping the same "torchvision build from source on top of PyTorch wheel":

The following have been reloaded with a version change:
  1) PyTorch/2.6.0-foss-2024a-binary => PyTorch/2.6.0-foss-2024a

Traceback (most recent call last):
  File "/vscmnt/brussel_pixiu_data/_data_brussel/vo/000/bvo00005/vsc10122/tests/pytorch/examples/super_resolution/super_resolve.py", line 5, in <module>
    from torchvision.transforms import ToTensor
  File "/user/brussel/101/vsc10122/easybuild/install/zen4/software/torchvision/0.21.0-foss-2024a-monster/lib/python3.12/site-packages/torchvision/__init__.py", line 10, in <module>
    from torchvision import _meta_registrations, datasets, io, models, ops, transforms, utils  # usort:skip
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/user/brussel/101/vsc10122/easybuild/install/zen4/software/torchvision/0.21.0-foss-2024a-monster/lib/python3.12/site-packages/torchvision/_meta_registrations.py", line 163, in <module>
    @torch.library.register_fake("torchvision::nms")
     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/apps/brussel/RL9/zen4/software/PyTorch/2.6.0-foss-2024a/lib/python3.12/site-packages/torch/library.py", line 828, in register
    use_lib._register_fake(op_name, func, _stacklevel=stacklevel + 1)
  File "/apps/brussel/RL9/zen4/software/PyTorch/2.6.0-foss-2024a/lib/python3.12/site-packages/torch/library.py", line 198, in _register_fake
    handle = entry.fake_impl.register(func_to_register, source)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/apps/brussel/RL9/zen4/software/PyTorch/2.6.0-foss-2024a/lib/python3.12/site-packages/torch/_library/fake_impl.py", line 31, in register
    if torch._C._dispatch_has_kernel_for_dispatch_key(self.qualname, "Meta"):
       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: operator torchvision::nms does not exist

This means that we must swap all modules that link to libraries from PyTorch to avoid breakage. Only those modules that depend on PyTorch at the Python level (i.e. just do import torch) are safe to keep.

[lexming]

The following packages have shared libraries that link to libtorch.so:

• Packages in PyTorch-bundle
  • torchvision
  • torchaudio
  • torchtext
  • torio
  • torchrl
  • tensordict
• Other packages
  • MONAI
  • OpenMM-Torch
  • photontorch
  • PyTorch-Geometric

This makes quite hard to have a setup with swapping modules to work safely.

[akesandgren]

So where does torchvision::nms come from?
Is it something we fail to build into our PyTorch or ?

[lexming]

@akesandgren it's weird because this torchvision::nms is there in both cases. Torchvision works well in the stack made of wheels and in our stack build from source. This error only happens when swapping PyTorch from one into the other. And it does happen with both swaps, so:

• you take the stack made of wheels and swap in our PyTorch build from source; then error about torchvision::nms
• you take the stack build from source and swap in the PyTorch from wheels, then also error about torchvision::nms

[Flamefire]

At least for torchvision I have an explanation:
Using the EB-build PyTorch 2.7:

eb-shell> pip install 'torchvision' --prefix /tmp/easybuild-tmp/eb-hxhejtep/tmp03q3yd29 --no-deps 
...
Successfully installed torchvision-0.23.0
eb-shell> python -c 'import torchvision'
...
RuntimeError: operator torchvision::nms does not exist

eb-shell> pip install 'torchvision~=0.22.0' --prefix /tmp/easybuild-tmp/eb-hxhejtep/tmp03q3yd29 --no-deps 
...
Successfully installed torchvision-0.22.1
eb-shell> python -c 'import torchvision'
eb-shell> echo $?
0

So basically it was just a version mismatch: We cannot mix packages build against different PyTorch versions.

This is also reflected in the requirements:

eb-shell> pip check
torchvision 0.23.0 has requirement torch==2.8.0, but you have torch 2.7.1.

This page lists the required PyTorch versions for each torchvision version

I tested the other way round: PyTorch-bundle/2.1.2-foss-2023a-CUDA-12.1.1 with PyPI torch==2.1.2

$ python -c 'import torchvision'
/software/PyTorch-bundle/2.1.2-foss-2023a-CUDA-12.1.1/lib/python3.11/site-packages/torchvision/io/image.py:13: UserWarning: Failed to load image Python extension: '/software/PyTorch-bundle/2.1.2-foss-2023a-CUDA-12.1.1/lib/python3.11/site-packages/torchvision/image.so: undefined symbol: _ZN5torch3jit17parseSchemaOrNameERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE'If you don't plan on using image functionality from `torchvision.io`, you can ignore this warning. Otherwise, there might be something wrong with your environment. Did you have `libjpeg` or `libpng` installed before building `torchvision` from source?
  warn(

Checking the build config of PyTorch (python -c "import torch; print(torch.__config__.show())") I see that the offical builds are built with -D_GLIBCXX_USE_CXX11_ABI=0 -fabi-version=11, the define causes this issue, for us it is "1"

In the source I see this documented with

    # Please note this is required in order to ensure compatibility between gcc 9 and gcc 7
    # This could be removed when all Linux PyTorch binary builds are compiled by the same toolchain again
    append_cxx_flag_if_supported("-fabi-version=11" CMAKE_CXX_FLAGS)

I can't tell why the official builds at that time included that disable-define but at least nowadays this isn't an issue.

Mixing a non-CUDA torch with a PYPI torchvision fails with

ImportError: libcudnn.so.8: cannot open shared object file: No such file or directory

And finally we can't mix different CUDA versions:

Detected that PyTorch and torchvision were compiled with different CUDA major versions. .... Please reinstall the torchvision that matches your PyTorch install."

[lexming]

That's interesting. In my case the error about torchvision::nms is not due a version mismatch. I tested this with the exact same versions of PyTorch and torchvision

• case A: with wheels from PyPI

  $ python -c "import torch; import torchvision; print(torch.__version__); print(torchvision.__version__); "
  2.6.0+cpu
  0.21.0+cpu
  

• case B: build from source

  $ python -c "import torch; import torchvision; print(torch.__version__); print(torchvision.__version__); "
  2.6.0
  0.21.0
  

Swapping the PyTorch module in the previous cases breaks torchvision as described with errors about torchvision::nms. So there is no version mismatch, unless the +cpu tag is causing the mismatch. It's either this or the CXX11 ABI.

[Flamefire]

So there is no version mismatch, unless the +cpu tag is causing the mismatch. It's either this or the CXX11 ABI.

The latter will cause another issue (undefined symbols), so I assume it is the CPU versions of torch are the issue. Not sure how you'd get those from PyPI. IIRC there were only GPU packages and you'd need to use different source-index to get them. Found e.g. https://docs.astral.sh/uv/guides/integration/pytorch/#installing-pytorch related to that:

the +cpu builds are published on https://download.pytorch.org/whl/cpu, while the +cu121 builds are published on https://download.pytorch.org/whl/cu121.

PyTorch build from source + torchvision wheel: BAD
PyTorch build from source + torchvision build from source on top of PyTorch wheel: BAD
PyTorch wheel + torchvision build from source on top of PyTorch from source: BAD

Just so I understand the last 2 cases correctly: You installed PyTorch from PyPI into the PyTorch module, then build the torchvision module (from source) then rebuild the PyTorch module (from source).
And the last case is a regular EB installation of both and then installing PyPI torch into the PyTorch module.

[lexming]

The wheels I'm using are the official ones from pytorch.org. Those are documented in https://pytorch.org/get-started/previous-versions/ and the repository for the CPU-only variant is indeed located in https://download.pytorch.org/whl/cpu.

Just so I understand the last 2 cases correctly: You installed PyTorch from PyPI into the PyTorch module, then build the torchvision module (from source) then rebuild the PyTorch module (from source).
And the last case is a regular EB installation of both and then installing PyPI torch into the PyTorch module.

I build/installed two stacks with the exact same versions for PyTorch and torchvision that work:

• A: PyTorch and torchvision build from sources (using the easyconfigs already in here)
• B: PyTorch and torchvision installed from official wheels (using custom easyconfigs)

Both stacks works fine. However, swapping modules between A and B breaks.

• PyTorch build from source + torchvision wheel:
  this is swapping PyTorch module in B with PyTorch from A
• PyTorch wheel + torchvision build from source on top of PyTorch from source:
  this is swapping PyTorch module in A with PyTorch from B
• PyTorch build from source + torchvision build from source on top of PyTorch wheel:
  this is a special one, I also build torchvision from source on top of the official PyTorch wheel. Swapping that module also breaks.

[Flamefire]

News from PyTorch devs:

Yes, with the new stable ABI + our C++ convenience wrappers, extensions that use the limited subset of APIs can be built in a libtorch stable way since 2.9.

So we might be fine with 2.9+

[lexming]

The current C++11 ABI setting does indeed make it impossible to switch from the libtorch libraries in upstream wheels and the libraries compiled from source in EB. The symbols in those libraries have different names.

For instance:

• _ZN5torch3jit17parseSchemaOrNameERKSsb in upstream wheel
• _ZN5torch3jit17parseSchemaOrNameERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEEb in build from source