adding SBSA to the warning log and installing MPI libs for the distributed tests

Author: apbose

.github/scripts/install-mpi-linux-x86.sh

@@ -12,12 +12,6 @@ if [[ $(uname -m) == "aarch64" ]]; then
     install_cuda_aarch64
 fi
 
-if [[ "$(uname -s)" == "Linux" && "$(uname -m)" == "x86_64" ]]; then
-    # install MPI for Linux x86_64
-    source .github/scripts/install-mpi-linux-x86.sh
-    install_mpi_linux_x86
-fi
-
 # Install all the dependencies required for Torch-TensorRT
 pip install --pre -r ${PWD}/tests/py/requirements.txt
 # dependencies in the tests/py/requirements.txt might install a different version of torch or torchvision

.github/scripts/install-torch-tensorrt.sh

@@ -363,6 +363,7 @@ jobs:
         export USE_HOST_DEPS=1
         export CI_BUILD=1
         export USE_TRTLLM_PLUGINS=1
+        dnf install -y mpich mpich-devel openmpi openmpi-devel
         pushd .
         cd tests/py
         cd dynamo

.github/workflows/build-test-linux-x86_64.yml

@@ -856,7 +856,7 @@ def is_platform_supported_for_trtllm(platform: str) -> bool:
         return False
     if "aarch64" in platform:
         logger.info(
-            "TensorRT-LLM plugins for NCCL backend are not supported on Jetson devices (aarch64)"
+            "TensorRT-LLM plugins for NCCL backend are not supported on Jetson devices or SBSA (aarch64)"
         )
         return False
     return True

py/torch_tensorrt/dynamo/utils.py

@@ -44,10 +44,10 @@ def forward(self, x):
 platform_str = str(Platform.current_platform()).lower()
 
 
-class TestGatherNcclOpsConverter(DispatchTestCase):
+class TestNcclOpsConverter(DispatchTestCase):
     @unittest.skipIf(
         "win" in platform_str or "aarch64" in platform_str,
-        "Skipped on Windows and Jetson: NCCL backend is not supported.",
+        "Skipped on Windows, Jetson, SBSA: NCCL backend is not supported.",
     )
     @classmethod
     def setUpClass(cls):