ci: fix GPU unit and integration test failures on multi-GPU runners

Author: sarunsingla11722

.github/workflows/run_tests_against_package.yml

@@ -91,6 +91,8 @@ jobs:
             || (inputs.device_type == 'cpu' && 'tpu' || inputs.device_type)
           }}
         ALLOW_MULTIPLE_LIBTPU_LOAD: ${{ inputs.device_type == 'cpu' && 'true' || '' }} # bypass /tmp/libtpu_lockfile check for cpu tests, which don't actually use accelerators (to allow concurrency)
+        NCCL_P2P_DISABLE: ${{ inputs.device_type == 'cuda12' && '1' || '' }}
+        NCCL_IB_DISABLE: ${{ inputs.device_type == 'cuda12' && '1' || '' }}
       options: ${{ inputs.container_resource_option }}
     steps:
       - name: Checkout MaxText
@@ -172,6 +174,11 @@ jobs:
                 done
               fi
             fi
+            # Restrict GPU unit tests to a single GPU to avoid multi-device NCCL failures
+            if [[ "${INPUTS_PYTEST_MARKER}" == *"not integration_test"* ]]; then
+              export CUDA_VISIBLE_DEVICES=0
+              echo "Restricting GPU unit tests to a single GPU: CUDA_VISIBLE_DEVICES=0"
+            fi
           fi
           if [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ]; then
             $PYTHON_EXE -m pip install --quiet pytest-split pytest-xdist