[#9847][fix] AutoDeploy: Added option to skip first layer of allreduce+rmsnorm fusion to avoid a hang

tensorrt_llm/_torch/auto_deploy/config/default.yaml

@@ -125,6 +125,7 @@ transforms:
     backend: trtllm
   fuse_allreduce_residual_rmsnorm:
     stage: post_load_fusion
+    skip_first_match: false
   # TODO (lucaslie): add backend selection as part of configurable inference optimizers
   fuse_rmsnorm:
     stage: post_load_fusion

tensorrt_llm/_torch/auto_deploy/transform/library/collectives.py

@@ -5,16 +5,22 @@
 The torch backend (demollm mode) does not benefit from fusion.
 """
 
-from functools import partial
-from typing import Tuple
+from typing import Tuple, Type
 
 import torch
+from pydantic import Field
 from torch.fx import GraphModule
 
 from ...models.factory import ModelFactory
 from ...shim.interface import CachedSequenceInterface
 from ...utils.pattern_matcher import ADPatternMatcherPass, register_ad_pattern
-from ..interface import BaseTransform, SharedConfig, TransformInfo, TransformRegistry
+from ..interface import (
+    BaseTransform,
+    SharedConfig,
+    TransformConfig,
+    TransformInfo,
+    TransformRegistry,
+)
 
 # TODO: This is an overly simplified model that works well for vanilla Llama models.
 # However, we eventually want to consider more sophisticated patterns such as
@@ -87,6 +93,19 @@ def _allreduce_residual_rmsnorm_replacement(
 # ============================================================================
 
 
+class FuseAllreduceResidualRMSNormConfig(TransformConfig):
+    """Configuration for the allreduce + residual + RMSNorm fusion transform."""
+
+    skip_first_match: bool = Field(
+        default=False,
+        description=(
+            "Skip fusing the first allreduce+rmsnorm match. This is a workaround for a hang "
+            "that occurs with torch.compile + NCCL when the first collective triggers lazy "
+            "initialization while ranks are desynchronized during JIT compilation."
+        ),
+    )
+
+
 @TransformRegistry.register("fuse_allreduce_residual_rmsnorm")
 class FuseAllreduceResidualRMSNorm(BaseTransform):
     """Fuse (allreduce + residual add + RMSNorm) into one fused op with tuple output.
@@ -96,6 +115,10 @@ class FuseAllreduceResidualRMSNorm(BaseTransform):
     this fusion and uses unfused operations.
     """
 
+    @classmethod
+    def get_config_class(cls) -> Type[TransformConfig]:
+        return FuseAllreduceResidualRMSNormConfig
+
     def _apply(
         self,
         gm: GraphModule,
@@ -132,11 +155,35 @@ def _apply(
             add_order="x_first", strategy=strategy
         )
 
+        # Create replacement functions, optionally skipping the first match
+        skip_first = self.config.skip_first_match
+        match_counter = [0]  # Use list to allow mutation in nested function
+
+        def _make_conditional_repl(pattern_fn):
+            def _conditional_repl(
+                x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
+            ):
+                match_counter[0] += 1
+                if skip_first and match_counter[0] == 1:
+                    # Skip first match - return unfused pattern
+                    return pattern_fn(x, residual, weight, eps)
+                # Fuse this match
+                return _allreduce_residual_rmsnorm_replacement(x, residual, weight, eps, strategy)
+
+            return _conditional_repl
+
+        _conditional_repl_residual_first = _make_conditional_repl(
+            _allreduce_residual_rmsnorm_pattern_trtllm
+        )
+        _conditional_repl_x_first = _make_conditional_repl(
+            _allreduce_residual_rmsnorm_pattern2_trtllm
+        )
+
         # Register TRT-LLM backend patterns only (no torch backend fusion)
         # Pattern 1: residual + allreduce(x)
         register_ad_pattern(
             search_fn=_allreduce_residual_rmsnorm_pattern_trtllm,
-            replace_fn=partial(_allreduce_residual_rmsnorm_replacement, strategy=strategy),
+            replace_fn=_conditional_repl_residual_first,
             patterns=patterns,
             dummy_args=dummy_args,
             op_ignore_types=op_ignore_types,
@@ -146,7 +193,7 @@ def _apply(
         # Pattern 2: allreduce(x) + residual
         register_ad_pattern(
             search_fn=_allreduce_residual_rmsnorm_pattern2_trtllm,
-            replace_fn=partial(_allreduce_residual_rmsnorm_replacement, strategy=strategy),
+            replace_fn=_conditional_repl_x_first,
             patterns=patterns,
             dummy_args=dummy_args,
             op_ignore_types=op_ignore_types,

tests/unittest/_torch/auto_deploy/unit/multigpu/test_ad_allreduce_strategies.py

@@ -187,15 +187,24 @@ def _prepare_dataset(root_dir: str, temp_dir: str, model_path_or_name: str, num_
         "SYMM_MEM",
     ],
 )
-def test_allreduce_strategies(llm_root, shared_dataset, allreduce_strategy):  # noqa: F811
-    """Test different allreduce strategies with multi-GPU configuration making sure that there are no crashes or hangs.
+@pytest.mark.parametrize(
+    "compile_backend",
+    [
+        "torch-cudagraph",
+        "torch-opt",
+    ],
+)
+def test_allreduce_strategies(llm_root, shared_dataset, allreduce_strategy, compile_backend):  # noqa: F811
+    """Test different allreduce strategies and compile backends with multi-GPU configuration.
 
     Configuration:
         The allreduce_strategy is set in the transforms config:
         ```yaml
         transforms:
           detect_sharding:
             allreduce_strategy: "ONESHOT"  # or AUTO, NCCL, TWOSHOT, etc.
+          compile_model:
+            backend: "torch-cudagraph"  # or torch-opt
         ```
 
     Test configuration:
@@ -208,6 +217,7 @@ def test_allreduce_strategies(llm_root, shared_dataset, allreduce_strategy):  #
         llm_root: Root directory fixture
         shared_dataset: Shared dataset fixture (prepared once for all test runs)
         allreduce_strategy: Strategy to test (AUTO, ONESHOT, TWOSHOT, MIN_LATENCY, NCCL)
+        compile_backend: Compile backend to test (torch-cudagraph, torch-opt)
     """
     # Fixed timeout for all strategies (5 minutes should be enough)
     TEST_TIMEOUT_SECONDS = 300
@@ -229,24 +239,37 @@ def test_allreduce_strategies(llm_root, shared_dataset, allreduce_strategy):  #
 
         # Create configuration with specified allreduce strategy in transforms
         extra_llm_api_options_path = f"{temp_dir}/extra_llm_api_options.yaml"
+
+        # Build transforms config
+        transforms_config = {
+            "detect_sharding": {
+                "stage": "sharding",
+                "allreduce_strategy": allreduce_strategy,
+            },
+            "compile_model": {
+                "stage": "compile",
+                "backend": compile_backend,
+                "cuda_graph_batch_sizes": [1, 2, 4, 8, 16, 32, 64, 128, 256],
+            },
+        }
+
+        # For torch-opt backend, enable skip_first_match workaround to prevent hang
+        # when torch.compile + NCCL triggers lazy initialization during JIT compilation
+        # see https://github.com/NVIDIA/TensorRT-LLM/issues/9847
+        if compile_backend == "torch-opt":
+            transforms_config["fuse_allreduce_residual_rmsnorm"] = {
+                "enabled": True,
+                "skip_first_match": True,
+            }
+
         with open(extra_llm_api_options_path, "w") as f:
             yaml.dump(
                 {
                     **config["args"],
                     "max_batch_size": max_batch_size,
                     "max_num_tokens": max_num_tokens,
                     "max_seq_len": 256,
-                    "transforms": {
-                        "detect_sharding": {
-                            "stage": "sharding",
-                            "allreduce_strategy": allreduce_strategy,
-                        },
-                        "compile_model": {
-                            "stage": "compile",
-                            "backend": "torch-cudagraph",
-                            "cuda_graph_batch_sizes": [1, 2, 4, 8, 16, 32, 64, 128, 256],
-                        },
-                    },
+                    "transforms": transforms_config,
                 },
                 f,
             )
@@ -288,7 +311,8 @@ def test_allreduce_strategies(llm_root, shared_dataset, allreduce_strategy):  #
                 assert result.exit_code == 0, f"Benchmark failed with output: {result.output}"
         except TimeoutError as e:
             pytest.fail(
-                f"Test timed out after {TEST_TIMEOUT_SECONDS}s for strategy {allreduce_strategy}. "
+                f"Test timed out after {TEST_TIMEOUT_SECONDS}s for strategy {allreduce_strategy} "
+                f"with backend {compile_backend}. "
                 f"This might indicate a hang (e.g., TWOSHOT without C++ fix). Error: {e}"
             )