[TRTLLM-10060][feat] Enable attention dp for Nemotron Super v3.

Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>

Author: nv-guomingz

tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py

@@ -11,8 +11,9 @@ class NemotronHHfWeightMapper(HfWeightMapper):
 
     def preprocess_weights(self, weights: dict) -> dict:
         config = self.config.pretrained_config
-        tp_size = self.config.mapping.tp_size
+        tp_size = 1 if self.config.mapping.enable_attention_dp else self.config.mapping.tp_size
         tp_rank = self.config.mapping.tp_rank
+
         d_inner = config.mamba_head_dim * config.mamba_num_heads
 
         def _split_mamba2_mixer_in_proj(w: torch.Tensor) -> torch.Tensor:

tensorrt_llm/_torch/models/modeling_nemotron_h.py

@@ -32,7 +32,7 @@
 from ..modules.decoder_layer import DecoderLayer
 from ..modules.embedding import Embedding
 from ..modules.fused_moe import MoEWeightLoadingMode, create_moe
-from ..modules.linear import Linear
+from ..modules.linear import Linear, TensorParallelMode
 from ..modules.mamba.mamba2_mixer import Mamba2Mixer
 from ..modules.mlp import MLP
 from ..modules.multi_stream_utils import maybe_execute_in_parallel
@@ -85,8 +85,10 @@ def __init__(
         self,
         model_config: ModelConfig[NemotronHConfig],
         layer_idx: int,
+        reduce_output: bool = False,
     ):
         config = model_config.pretrained_config
+
         super().__init__(
             hidden_size=config.hidden_size,
             num_attention_heads=config.num_attention_heads,
@@ -97,6 +99,7 @@ def __init__(
             layer_idx=layer_idx,
             dtype=config.torch_dtype,
             config=model_config,
+            reduce_output=reduce_output,
         )
 
     def forward(
@@ -154,6 +157,7 @@ def __init__(
             shared_expert_intermediate_size = (
                 config.moe_shared_expert_intermediate_size *
                 config.n_shared_experts)
+
             self.shared_experts = MLP(
                 hidden_size=config.hidden_size,
                 intermediate_size=shared_expert_intermediate_size,
@@ -193,11 +197,14 @@ def __init__(
             activation_type=self.activation_type,
         )
 
-        # AllReduce for combining shared and routed expert outputs in multi-GPU settings.
-        self.allreduce = AllReduce(
-            mapping=model_config.mapping,
-            strategy=model_config.allreduce_strategy,
-        )
+        if not model_config.mapping.enable_attention_dp:
+            # AllReduce for combining shared and routed expert outputs in multi-GPU settings.
+            self.allreduce = AllReduce(
+                mapping=model_config.mapping,
+                strategy=model_config.allreduce_strategy,
+            )
+        else:
+            self.allreduce = None
 
         # Setup latent projection layers.
         if self.use_latent_moe:
@@ -314,7 +321,11 @@ def __init__(
         elif layer_type == "-":
             self.mixer = MLPLayer(model_config, layer_idx)
         elif layer_type == "*":
-            self.mixer = TransformerLayer(model_config, layer_idx)
+            self.mixer = TransformerLayer(
+                model_config,
+                layer_idx,
+                reduce_output=not model_config.mapping.enable_attention_dp
+                and model_config.mapping.tp_size > 1)
         elif layer_type == "E":
             self.mixer = NemotronHMOE(model_config,
                                       layer_idx=layer_idx,
@@ -357,12 +368,24 @@ def __init__(self, model_config: ModelConfig[NemotronHConfig]):
             aux_stream_list[2],
         }
 
-        # calculate embeddings
-        self.embed_tokens = Embedding(
-            config.vocab_size,
-            config.hidden_size,
-            dtype=config.torch_dtype,
-        )
+        if model_config.mapping.enable_attention_dp:
+            # When attention_dp is enabled, we cannot do all_reduce since
+            # the problem size of different ranks are different.
+            # So, we don't do parallelism here.
+            self.embed_tokens = Embedding(
+                config.vocab_size,
+                config.hidden_size,
+                dtype=config.torch_dtype,
+            )
+        else:
+            self.embed_tokens = Embedding(
+                config.vocab_size,
+                config.hidden_size,
+                dtype=config.torch_dtype,
+                mapping=model_config.mapping,
+                tensor_parallel_mode=TensorParallelMode.COLUMN,
+                gather_output=True,
+            )
 
         # create layers
         layers = []

tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py

@@ -57,8 +57,23 @@ def __init__(
 
         config = config or ModelConfig()
         self.mapping = config.mapping
-        tp_rank = config.mapping.tp_rank
-        tp_size = config.mapping.tp_size
+
+        if config.mapping.enable_attention_dp:
+            from tensorrt_llm.mapping import Mapping
+            self.mapping = Mapping(
+                world_size=config.mapping.pp_size,
+                tp_size=1,
+                pp_size=config.mapping.pp_size,
+                rank=config.mapping.rank,
+                gpus_per_node=config.mapping.gpus_per_node,
+                enable_attention_dp=True,
+            )
+            tp_size = 1
+            tp_rank = 0
+        else:
+            self.mapping = config.mapping
+            tp_size = config.mapping.tp_size
+            tp_rank = config.mapping.tp_rank
 
         d_inner = head_dim * nheads
         d_in_proj = 2 * d_inner + 2 * n_groups * d_state + nheads

tensorrt_llm/_torch/modules/mlp.py

@@ -4,6 +4,8 @@
 import torch
 from torch import nn
 
+from tensorrt_llm.mapping import Mapping
+
 from ..model_config import ModelConfig
 from ..peft.lora.layer import LoraLayer, LoraModuleType
 from .linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig
@@ -21,24 +23,33 @@ def __init__(self,
                  config: Optional[ModelConfig] = None,
                  layer_idx: Optional[int] = None,
                  reduce_output: bool = True):
-
+        if config.mapping.enable_attention_dp:
+            mapping = Mapping(
+                world_size=config.mapping.pp_size,
+                tp_size=1,
+                pp_size=config.mapping.pp_size,
+                rank=config.mapping.rank,
+                gpus_per_node=config.mapping.gpus_per_node,
+                enable_attention_dp=True,
+            )
+        else:
+            mapping = config.mapping
         super().__init__()
         self.layer_idx = layer_idx
         self.hidden_size = hidden_size
         self.intermediate_size = intermediate_size
         self.activation = activation
 
         config = config or ModelConfig()
-        self.up_lora = LoraLayer(
-            [LoraModuleType.MLP_H_TO_4H],
-            [self.intermediate_size // config.mapping.tp_size])
+        self.up_lora = LoraLayer([LoraModuleType.MLP_H_TO_4H],
+                                 [self.intermediate_size // mapping.tp_size])
 
         self.up_proj = Linear(
             self.hidden_size,
             self.intermediate_size,
             bias=bias,
             dtype=dtype,
-            mapping=config.mapping,
+            mapping=mapping,
             tensor_parallel_mode=TensorParallelMode.COLUMN,
             weights_loading_config=WeightsLoadingConfig(
                 weight_mode=WeightMode.VANILLA),
@@ -55,7 +66,7 @@ def __init__(self,
             self.hidden_size,
             bias=bias,
             dtype=dtype,
-            mapping=config.mapping,
+            mapping=mapping,
             tensor_parallel_mode=TensorParallelMode.ROW,
             quant_config=config.get_quant_config(),
             skip_create_weights_in_init=config.skip_create_weights_in_init,

tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py

@@ -45,7 +45,7 @@ def __init__(
         self.mamba_ssm_cache_dtype = ssm_cache_dtype
 
         # get tp size
-        tp_size = mapping.tp_size
+        tp_size = mapping.tp_size if not mapping.enable_attention_dp else 1
 
         # derive mamba parameters for conv and ssm states
         d_inner = head_dim * num_heads

tests/integration/test_lists/waives.txt

@@ -444,7 +444,6 @@ unittest/_torch/auto_deploy/unit/singlegpu/test_ad_build_small_single.py::test_b
 triton_server/test_triton.py::test_opt[opt] SKIP (https://nvbugs/5739981)
 unittest/llmapi/test_llm_pytorch.py::test_tinyllama_logits_processor[False] SKIP (https://nvbugs/5771838)
 unittest/llmapi/test_llm_pytorch.py::test_tinyllama_logits_processor[True] SKIP (https://nvbugs/5771838)
-accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[dp4-trtllm-auto] SKIP (https://nvbugs/5772417)
 accuracy/test_cli_flow.py::TestPhi2::test_auto_dtype SKIP (https://nvbugs/5744293)
 accuracy/test_cli_flow.py::TestPhi2::test_tp2 SKIP (https://nvbugs/5744293)
 accuracy/test_cli_flow.py::TestPhi3Mini4kInstruct::test_auto_dtype SKIP (https://nvbugs/5744293)