adding checks for windows and jetson devices

Author: apbose

py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py

@@ -15,11 +15,11 @@
     tensorrt_fused_nccl_all_gather_op,
     tensorrt_fused_nccl_reduce_scatter_op,
 )
-from torch_tensorrt.dynamo.utils import load_tensorrt_llm
+from torch_tensorrt.dynamo.utils import load_tensorrt_llm_for_nccl
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
-if load_tensorrt_llm():
+if load_tensorrt_llm_for_nccl():
 
     @dynamo_tensorrt_converter(tensorrt_fused_nccl_all_gather_op)
     def fused_nccl_gather(

py/torch_tensorrt/dynamo/utils.py

@@ -841,6 +841,29 @@ def is_tegra_platform() -> bool:
     return False
 
 
+def is_platform_supported_for_trtllm(platform: str) -> bool:
+    """
+    Checks if the current platform supports TensorRT-LLM plugins for NCCL backend
+    Returns:
+        bool: True if the platform supports TensorRT-LLM plugins for NCCL backend, False otherwise.
+    Note:
+        TensorRT-LLM plugins for NCCL backend are not supported on:
+        - Windows platforms
+        - Jetson devices (aarch64 architecture)
+    """
+    if "windows" in platform:
+        logger.info(
+            "TensorRT-LLM plugins for NCCL backend are not supported on Windows"
+        )
+        return False
+    if "aarch64" in platform:
+        logger.info(
+            "TensorRT-LLM plugins for NCCL backend are not supported on Jetson devices (aarch64)"
+        )
+        return False
+    return True
+
+
 @contextmanager
 def download_plugin_lib_path(platform: str) -> Iterator[str]:
     """
@@ -891,6 +914,7 @@ def download_plugin_lib_path(platform: str) -> Iterator[str]:
     if "linux" in platform:
         lib_filename = "libnvinfer_plugin_tensorrt_llm.so"
     else:
+        # This condition is never met though
         lib_filename = "libnvinfer_plugin_tensorrt_llm.dll"
 
     with tempfile.TemporaryDirectory() as tmpdir:
@@ -923,7 +947,7 @@ def download_plugin_lib_path(platform: str) -> Iterator[str]:
         yield plugin_lib_path
 
 
-def load_and_initialize_trtllm_plugin(plugin_lib_path: str, platform: str) -> bool:
+def load_and_initialize_trtllm_plugin(plugin_lib_path: str) -> bool:
     """
     Loads and initializes the TensorRT-LLM plugin from the given shared library path.
 
@@ -933,9 +957,6 @@ def load_and_initialize_trtllm_plugin(plugin_lib_path: str, platform: str) -> bo
     Returns:
         bool: True if successful, False otherwise.
     """
-    if "windows" in platform:
-        logger.info("NCCL backend is not supported on Windows")
-        return False
     try:
         handle = ctypes.CDLL(plugin_lib_path)
         logger.info(f"Successfully loaded plugin library: {plugin_lib_path}")
@@ -981,7 +1002,7 @@ def load_and_initialize_trtllm_plugin(plugin_lib_path: str, platform: str) -> bo
     return False
 
 
-def load_tensorrt_llm() -> bool:
+def load_tensorrt_llm_for_nccl() -> bool:
     """
     Attempts to load the TensorRT-LLM plugin and initialize it.
     Either the env variable TRTLLM_PLUGINS_PATH can specify the path
@@ -990,11 +1011,15 @@ def load_tensorrt_llm() -> bool:
     Returns:
         bool: True if the plugin was successfully loaded and initialized, False otherwise.
     """
-    plugin_lib_path = os.environ.get("TRTLLM_PLUGINS_PATH")
+    # Check platform compatibility first
     platform = Platform.current_platform()
     platform = str(platform).lower()
+    if not is_platform_supported_for_trtllm(platform):
+        return False
+    plugin_lib_path = os.environ.get("TRTLLM_PLUGINS_PATH")
+
     if plugin_lib_path:
-        return load_and_initialize_trtllm_plugin(plugin_lib_path, platform)
+        return load_and_initialize_trtllm_plugin(plugin_lib_path)
     else:
         # this option can be used by user if TRTLLM_PLUGINS_PATH is not set by user
         use_trtllm_plugin = os.environ.get("USE_TRTLLM_PLUGINS", "0").lower() in (
@@ -1010,5 +1035,5 @@ def load_tensorrt_llm() -> bool:
             return False
 
         with download_plugin_lib_path(platform) as plugin_lib_path:
-            return load_and_initialize_trtllm_plugin(plugin_lib_path, platform)
+            return load_and_initialize_trtllm_plugin(plugin_lib_path)
     return False

tests/py/dynamo/distributed/test_nccl_ops.py

@@ -1,4 +1,5 @@
 import os
+import unittest
 
 import torch
 import torch.distributed as dist
@@ -19,12 +20,13 @@
 platform_str = str(Platform.current_platform()).lower()
 
 
-@unittest.skipIf(
-    "win" in platform_str, "Skipped on Windows: NCCL backend is not supported."
-)
 class TestGatherNcclOpsConverter(DispatchTestCase):
+    @unittest.skipIf(
+        "win" or "aarch64" in platform_str,
+        "Skipped on Windows and Jetson: NCCL backend is not supported.",
+    )
     @parameterized.expand([8])
-    def test_nccl_ops(self, linear_layer_dim):
+    def test_nccl_ops_gather(self, linear_layer_dim):
         class DistributedGatherModel(nn.Module):
             def __init__(self, input_dim):
                 super().__init__()
@@ -48,6 +50,10 @@ def forward(self, x):
             enable_passes=True,
         )
 
+    @unittest.skipIf(
+        "win" or "aarch64" in platform_str,
+        "Skipped on Windows and Jetson: NCCL backend is not supported.",
+    )
     @parameterized.expand([8])
     def test_nccl_ops_scatter(self, linear_layer_dim):