bitsandbytes

Author: Borda

.azure/gpu-tests-fabric.yml

@@ -123,6 +123,17 @@ jobs:
           python requirements/pytorch/check-avail-extras.py
         displayName: "Env details"
 
+      - bash: |
+          # get pytorch version
+          PYTORCH_VERSION=$(python -c "import torch; print(torch.__version__.split('+')[0])")
+          # FixMe: uninstall bitsandbytes for pytorch 2.6 as it is not compatible with `triton.ops`
+          if [[ "${PYTORCH_VERSION}" == "2.6.0" ]]; then
+            pip uninstall -y bitsandbytes
+          else
+            python -c "import bitsandbytes"
+          fi
+        displayName: "Handle bitsandbytes"
+
       - bash: python -m pytest lightning_fabric
         workingDirectory: src
         # without succeeded this could run even if the job has already failed

.azure/gpu-tests-pytorch.yml

@@ -137,6 +137,17 @@ jobs:
           python requirements/pytorch/check-avail-extras.py
         displayName: "Env details"
 
+      - bash: |
+          # get pytorch version
+          PYTORCH_VERSION=$(python -c "import torch; print(torch.__version__.split('+')[0])")
+          # FixMe: uninstall bitsandbytes for pytorch 2.6 as it is not compatible with `triton.ops`
+          if [[ "${PYTORCH_VERSION}" == "2.6.0" ]]; then
+            pip uninstall -y bitsandbytes
+          else
+            python -c "import bitsandbytes"
+          fi
+        displayName: "Handle bitsandbytes"
+
       - bash: python -m pytest pytorch_lightning
         workingDirectory: src
         # without succeeded this could run even if the job has already failed

requirements/pytorch/check-avail-extras.py

@@ -4,7 +4,3 @@
     import matplotlib  # noqa: F401
     import omegaconf  # noqa: F401
     import rich  # noqa: F401
-    import torch
-
-    if torch.cuda.is_available():
-        import bitsandbytes  # noqa: F401