clever switching of graph applications

Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

Author: nv-lschneider

tensorrt_llm/_torch/compilation/pattern_applier.py

@@ -8,13 +8,15 @@
 
 import torch
 import torch.nn as nn
+from torch._inductor.pattern_matcher import PatternMatcherPass
 from torch.fx import GraphModule
 
 import tensorrt_llm
 from tensorrt_llm import logger
 from tensorrt_llm.mapping import Mapping
 
 from .backend import Backend
+from .patterns.ar_residual_norm import register_nccl_symmetric_patterns
 
 
 def apply_nccl_symmetric_patterns_to_model(model: nn.Module, mapping: Mapping) -> nn.Module:
@@ -51,6 +53,10 @@ def __init__(self, mapping: Mapping):
                 max_num_streams=1,
                 mapping=mapping,
             )
+            # Override custom_passes to register NCCL_SYMMETRIC patterns with match_auto_strategy=True
+            # This ensures we match AUTO strategy nodes (which is what's in the graph) and convert to NCCL_SYMMETRIC
+            self.custom_passes = [PatternMatcherPass()]
+            register_nccl_symmetric_patterns(self.custom_passes, mapping, match_auto_strategy=True)
 
         def __call__(self, gm: GraphModule, example_inputs: List[torch.Tensor]) -> callable:
             """

tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py

@@ -748,14 +748,29 @@ def target_finalize_pattern(
 
 
 def register_nccl_symmetric_patterns(custom_passes: List[PatternMatcherPass],
-                                     mapping: Mapping):
+                                     mapping: Mapping,
+                                     match_auto_strategy: bool = False):
+    """
+    Register NCCL_SYMMETRIC patterns.
+
+    Args:
+        custom_passes: List of pattern matcher passes to register patterns in
+        mapping: Mapping configuration for TP group
+        match_auto_strategy: If True, match AUTO strategy nodes (strategy=3) and convert to NCCL_SYMMETRIC.
+                           If False, match NCCL_SYMMETRIC strategy nodes (strategy=8) directly.
+                           Should be True when NCCL_SYMMETRIC is requested but graph nodes have AUTO.
+    """
 
     def register_convert_supported_ar_to_nccl_symmetric(
             custom_pass: PatternMatcherPass):
         # FX normalizes keyword arguments to positional arguments based on schema
         # Pattern should match positional arguments in schema order:
         # input, residual, norm_weight, scale, bias, workspace, group, strategy, op, eps, trigger_completion_at_end
-        strategy = int(AllReduceStrategy.NCCL_SYMMETRIC)
+        # When match_auto_strategy=True, match AUTO strategy nodes (strategy=3) since graph nodes
+        # typically have AUTO even when NCCL_SYMMETRIC is requested in config. The pattern will convert them to NCCL_SYMMETRIC.
+        # When match_auto_strategy=False, match NCCL_SYMMETRIC strategy nodes (strategy=8) directly.
+        strategy = int(AllReduceStrategy.AUTO) if match_auto_strategy else int(
+            AllReduceStrategy.NCCL_SYMMETRIC)
         input_node = KeywordArg('input')
         fusion = KeywordArg('op')  # Schema uses 'op', not 'fusion_op'
         trtllm_allreduce_default = CallFunction(
@@ -765,8 +780,15 @@ def register_convert_supported_ar_to_nccl_symmetric(
             fusion, KeywordArg('eps'), Ignored())
 
         if tensorrt_llm.mpi_rank() == 0:
-            logger.debug("[NCCL_SYMMETRIC] Pattern: Registering pattern with "
-                         f"strategy={strategy}, group={mapping.tp_group}")
+            if match_auto_strategy:
+                logger.debug(
+                    "[NCCL_SYMMETRIC] Pattern: Registering pattern to match AUTO strategy "
+                    f"(strategy={strategy}) and convert to NCCL_SYMMETRIC, group={mapping.tp_group}"
+                )
+            else:
+                logger.debug(
+                    "[NCCL_SYMMETRIC] Pattern: Registering pattern to match NCCL_SYMMETRIC strategy "
+                    f"(strategy={strategy}), group={mapping.tp_group}")
 
         def empty_convert_supported_ar_to_nccl_symmetric(
             input: torch.Tensor,
@@ -1085,5 +1107,8 @@ def register_ar_fusions(custom_passes: List[PatternMatcherPass],
     if enable_ub:
         register_ub_patterns(custom_passes, mapping)
 
-    # Always register NCCL_SYMMETRIC patterns (they only match when strategy is explicitly NCCL_SYMMETRIC)
-    register_nccl_symmetric_patterns(custom_passes, mapping)
+    # Register NCCL_SYMMETRIC patterns - by default match NCCL_SYMMETRIC strategy nodes directly
+    # Set match_auto_strategy=True only when NCCL_SYMMETRIC is requested but graph has AUTO nodes
+    register_nccl_symmetric_patterns(custom_passes,
+                                     mapping,
+                                     match_auto_strategy=False)