diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml index 4200be6c6aab..3da8dab3c624 100644 --- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml +++ b/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 diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/collectives.py b/tensorrt_llm/_torch/auto_deploy/transform/library/collectives.py index 85dc6c48bed3..3db5973d7053 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/collectives.py +++ b/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, diff --git a/tests/unittest/_torch/auto_deploy/unit/multigpu/test_ad_allreduce_strategies.py b/tests/unittest/_torch/auto_deploy/unit/multigpu/test_ad_allreduce_strategies.py index 9d4e444e4d4a..ecac8280928e 100644 --- a/tests/unittest/_torch/auto_deploy/unit/multigpu/test_ad_allreduce_strategies.py +++ b/tests/unittest/_torch/auto_deploy/unit/multigpu/test_ad_allreduce_strategies.py @@ -187,8 +187,15 @@ 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: @@ -196,6 +203,8 @@ def test_allreduce_strategies(llm_root, shared_dataset, allreduce_strategy): # 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,6 +239,29 @@ 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( { @@ -236,17 +269,7 @@ def test_allreduce_strategies(llm_root, shared_dataset, allreduce_strategy): # "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}" )