[CI] Trick Bagua into installing appropriate wheel in GPU tests (#14380)

Bagua trick needs to be replicated on everywhere applicable

Author: otaj

.azure/gpu-tests.yml

@@ -73,7 +73,8 @@ jobs:
         python -c "fname = 'requirements/pytorch/strategies.txt' ; lines = [line for line in open(fname).readlines() if 'horovod' not in line] ; open(fname, 'w').writelines(lines)"
         python -c "fname = 'requirements/pytorch/strategies.txt' ; lines = [line for line in open(fname).readlines() if 'bagua' not in line] ; open(fname, 'w').writelines(lines)"
         CUDA_VERSION_MM=$(python -c "import torch ; print(''.join(map(str, torch.version.cuda.split('.')[:2])))")
-        pip install "bagua-cuda$CUDA_VERSION_MM>=0.9.0"
+        CUDA_VERSION_BAGUA=$(python -c "print([ver for ver in [115,113,111,102] if $CUDA_VERSION_MM >= ver][0])")
+        pip install "bagua-cuda$CUDA_VERSION_BAGUA>=0.9.0"
         pip install -e .[strategies]
         pip install -U deepspeed  # TODO: remove when docker images are upgraded
         pip install --requirement requirements/pytorch/devel.txt

dockers/base-conda/Dockerfile

@@ -141,8 +141,9 @@ RUN \
 RUN \
     # install Bagua
     CUDA_VERSION_MM=$(python -c "print(''.join('$CUDA_VERSION'.split('.')[:2]))") && \
-    pip install "bagua-cuda$CUDA_VERSION_MM==0.9.0" && \
-    python -c "import bagua_core; bagua_core.install_deps()" && \
+    CUDA_VERSION_BAGUA=$(python -c "print([ver for ver in [115,113,111,102] if $CUDA_VERSION_MM >= ver][0])") && \
+    pip install "bagua-cuda$CUDA_VERSION_BAGUA==0.9.0" && \
+    if [[ "$CUDA_VERSION_MM" = "$CUDA_VERSION_BAGUA" ]]; then python -c "import bagua_core; bagua_core.install_deps()"; fi && \
     python -c "import bagua; print(bagua.__version__)"
 
 RUN \