TL/CUDA: Add new data types into NVLS MPI CI

This commit adds new data types and enables reduce scatter
mpi test for nvls CI.

Signed-off-by: Juee14Desai <jueehimalbha@nvidia.com>

Author: Juee14Desai

.ci/pipeline/test_nvls_matrix.yaml

@@ -80,13 +80,22 @@ steps:
     onfail: |
       sudo -E -u svcnbu-swx-hpcx ${WORKSPACE}/.ci/scripts/stop_slurm_allocation.sh
 
-  - name: Run UCC NVLS MPI tests
+  - name: Run UCC NVLS MPI tests (allreduce)
     containerSelector: "{name: 'build_helper'}"
     timeout: "${TEST_TIMEOUT_MINUTES}"
     run: |
       set -x
       export DOCKER_IMAGE_NAME="${registry_host}#torch-ucc/${UCC_URI_SUFFIX}:${DOCKER_IMAGE_TAG}"
       export SLURM_JOB_ID=$(cat ${WORKSPACE}/job_id.txt)
       sudo -E -u svcnbu-swx-hpcx ${WORKSPACE}/.ci/scripts/run_nvls_slurm.sh '/opt/nvidia/src/ucc/.ci/scripts/run_tests_ucc_nvls_mpi.sh' ${NVLS_MPI_PPN:-4}
+
+  - name: Run UCC NVLS MPI tests (reduce_scatter)
+    containerSelector: "{name: 'build_helper'}"
+    timeout: "${TEST_TIMEOUT_MINUTES}"
+    run: |
+      set -x
+      export DOCKER_IMAGE_NAME="${registry_host}#torch-ucc/${UCC_URI_SUFFIX}:${DOCKER_IMAGE_TAG}"
+      export SLURM_JOB_ID=$(cat ${WORKSPACE}/job_id.txt)
+      sudo -E -u svcnbu-swx-hpcx ${WORKSPACE}/.ci/scripts/run_nvls_slurm.sh '/opt/nvidia/src/ucc/.ci/scripts/run_tests_ucc_nvls_mpi_reduce_scatter.sh' ${NVLS_MPI_PPN:-4}
     always: |
       sudo -E -u svcnbu-swx-hpcx ${WORKSPACE}/.ci/scripts/stop_slurm_allocation.sh

.ci/scripts/run_tests_ucc_nvls_mpi.sh

@@ -15,11 +15,8 @@ export UCC_TL_CUDA_NVLS_SM_COUNT=4
 
 EXE="/opt/nvidia/src/ucc/build/test/mpi/ucc_test_mpi"
 EXE+=" --set_device 2 --mtypes cuda"
+DTYPES="float32,int32,uint32,int64,uint64"
 
 echo "INFO: NVLS MPI tests (allreduce) ..."
-UCC_TL_CUDA_TUNE="allreduce:cuda:@0" $EXE -c allreduce -d float32 -o sum -m 1024:33554432
+UCC_TL_CUDA_TUNE="allreduce:cuda:@0" $EXE -c allreduce -d ${DTYPES} -o sum -m 1024:33554432
 echo "INFO: NVLS MPI tests (allreduce) ... DONE"
-
-# echo "INFO: NVLS MPI tests (reduce_scatter) ..."
-# UCC_TL_CUDA_TUNE="reduce_scatter:cuda:@3" $EXE -c reduce_scatter -d float32 -o sum -m 1024:33554432
-# echo "INFO: NVLS MPI tests (reduce_scatter) ... DONE"

.ci/scripts/run_tests_ucc_nvls_mpi_reduce_scatter.sh

@@ -0,0 +1,24 @@
+#!/bin/bash -xe
+# NVLS reduce_scatter only. Run as a separate srun step (separate MPI job).
+# -m 1024:33554432:4 keeps per-rank counts NVLS-aligned.
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+source "${SCRIPT_DIR}/env.sh"
+
+if [ "${SLURM_LOCALID:-0}" = "0" ]; then
+    "${SCRIPT_DIR}/check_nvls_fabric.sh"
+fi
+
+export OMPI_MCA_coll=^hcoll
+export OMPI_MCA_coll_ucc_enable=0
+export UCC_TLS=cuda,ucp
+export UCC_LOG_LEVEL=info
+export UCC_TL_CUDA_NVLS_SM_COUNT=4
+
+EXE="/opt/nvidia/src/ucc/build/test/mpi/ucc_test_mpi"
+EXE+=" --set_device 2 --mtypes cuda"
+DTYPES="float32,int32,uint32,int64,uint64"
+
+echo "INFO: NVLS MPI tests (reduce_scatter) ..."
+UCC_TL_CUDA_TUNE="reduce_scatter:cuda:@3" $EXE -c reduce_scatter -d ${DTYPES} -o sum -m 1024:33554432:4
+echo "INFO: NVLS MPI tests (reduce_scatter) ... DONE"