[None][perf] Improve PyTorch encoder-decoder support and performance

tensorrt_llm/_torch/models/modeling_bart.py

@@ -18,20 +18,21 @@
 
 Key differences from T5:
     - LayerNorm instead of RMSNorm.
-    - Post-norm (residual → add → LayerNorm) instead of pre-norm.
+    - BART uses post-norm; mBART uses pre-norm and final stack norms.
     - Learned absolute positional embeddings (not relative bias).
-    - GELU activation (not ReLU / gated).
+    - The checkpoint selects the MLP activation (typically GELU for BART and
+      ReLU for mBART).
     - Bias in attention and MLP projections.
-    - Embedding scale = sqrt(d_model).
+    - mBART scales token embeddings by sqrt(d_model).
 """
 
 import math
 from typing import Dict, Optional
 
 import torch
-import torch.nn.functional as F
 from torch import nn
 from transformers import BartConfig
+from transformers.activations import ACT2FN
 
 from ..attention_backend import AttentionMetadata
 from ..attention_backend.interface import PredefinedAttentionMask
@@ -86,6 +87,18 @@ def _bart_head_dim(config: BartConfig) -> int:
     return config.d_model // config.encoder_attention_heads
 
 
+def _bart_normalize_before(config: BartConfig) -> bool:
+    return getattr(config, "model_type", None) == "mbart" or bool(
+        getattr(config, "normalize_before", False)
+    )
+
+
+def _bart_add_final_layer_norm(config: BartConfig) -> bool:
+    return getattr(config, "model_type", None) == "mbart" or bool(
+        getattr(config, "add_final_layer_norm", False)
+    )
+
+
 def _packed_position_ids(
     position_ids: Optional[torch.IntTensor],
     hidden_states: torch.Tensor,
@@ -167,7 +180,7 @@ def __init__(
 
 
 class BartEncoderLayer(nn.Module):
-    """BART encoder layer: self-attention → add+LN → MLP → add+LN (post-norm)."""
+    """BART/mBART encoder layer with configurable pre- or post-norm."""
 
     def __init__(
         self,
@@ -179,6 +192,7 @@ def __init__(
         hidden_size = config.d_model
         ffn_dim = _bart_encoder_ffn_dim(config)
         num_heads = _bart_encoder_num_heads(config)
+        self.normalize_before = _bart_normalize_before(config)
 
         self.self_attn = BartSelfAttention(model_config, num_heads=num_heads, layer_idx=layer_idx)
 
@@ -187,13 +201,14 @@ def __init__(
             eps=1e-5,
             dtype=config.torch_dtype,
             has_bias=True,
+            residual_in_fp32=False,
         )
 
         self.mlp = MLP(
             hidden_size=hidden_size,
             intermediate_size=ffn_dim,
             bias=True,
-            activation=F.gelu,
+            activation=ACT2FN[config.activation_function],
             dtype=config.torch_dtype,
             config=model_config,
             layer_idx=layer_idx,
@@ -204,6 +219,7 @@ def __init__(
             eps=1e-5,
             dtype=config.torch_dtype,
             has_bias=True,
+            residual_in_fp32=False,
         )
 
     def forward(
@@ -214,19 +230,27 @@ def forward(
         **kwargs,
     ) -> torch.Tensor:
         residual = hidden_states
+        if self.normalize_before:
+            hidden_states = self.self_attn_layer_norm(hidden_states)
         hidden_states = self.self_attn(
             position_ids=position_ids,
             hidden_states=hidden_states,
             attn_metadata=attn_metadata,
             attention_mask=PredefinedAttentionMask.FULL,
         )
-        hidden_states = residual + hidden_states
-        hidden_states = self.self_attn_layer_norm(hidden_states)
+        if self.normalize_before:
+            hidden_states = residual + hidden_states
+        else:
+            hidden_states, _ = self.self_attn_layer_norm(hidden_states, residual)
 
         residual = hidden_states
+        if self.normalize_before:
+            hidden_states = self.final_layer_norm(hidden_states)
         hidden_states = self.mlp(hidden_states)
-        hidden_states = residual + hidden_states
-        hidden_states = self.final_layer_norm(hidden_states)
+        if self.normalize_before:
+            hidden_states = residual + hidden_states
+        else:
+            hidden_states, _ = self.final_layer_norm(hidden_states, residual)
 
         return hidden_states
 
@@ -237,7 +261,7 @@ def forward(
 
 
 class BartDecoderLayer(nn.Module):
-    """BART decoder layer: self-attn → add+LN → cross-attn → add+LN → MLP → add+LN."""
+    """BART/mBART decoder layer with configurable pre- or post-norm."""
 
     def __init__(
         self,
@@ -249,6 +273,7 @@ def __init__(
         hidden_size = config.d_model
         ffn_dim = _bart_decoder_ffn_dim(config)
         num_heads = _bart_decoder_num_heads(config)
+        self.normalize_before = _bart_normalize_before(config)
 
         self.self_attn = BartSelfAttention(model_config, num_heads=num_heads, layer_idx=layer_idx)
 
@@ -257,6 +282,7 @@ def __init__(
             eps=1e-5,
             dtype=config.torch_dtype,
             has_bias=True,
+            residual_in_fp32=False,
         )
 
         self.cross_attn = BartCrossAttention(model_config, layer_idx=layer_idx)
@@ -266,13 +292,14 @@ def __init__(
             eps=1e-5,
             dtype=config.torch_dtype,
             has_bias=True,
+            residual_in_fp32=False,
         )
 
         self.mlp = MLP(
             hidden_size=hidden_size,
             intermediate_size=ffn_dim,
             bias=True,
-            activation=F.gelu,
+            activation=ACT2FN[config.activation_function],
             dtype=config.torch_dtype,
             config=model_config,
             layer_idx=layer_idx,
@@ -283,6 +310,7 @@ def __init__(
             eps=1e-5,
             dtype=config.torch_dtype,
             has_bias=True,
+            residual_in_fp32=False,
         )
 
     def forward(
@@ -295,34 +323,46 @@ def forward(
         skip_cross_kv_projection: bool = False,
         **kwargs,
     ) -> torch.Tensor:
-        # Self-attention (post-norm)
+        # Self-attention
         residual = hidden_states
+        if self.normalize_before:
+            hidden_states = self.self_attn_layer_norm(hidden_states)
         hidden_states = self.self_attn(
             position_ids=position_ids,
             hidden_states=hidden_states,
             attn_metadata=attn_metadata,
             attention_mask=PredefinedAttentionMask.CAUSAL,
         )
-        hidden_states = residual + hidden_states
-        hidden_states = self.self_attn_layer_norm(hidden_states)
+        if self.normalize_before:
+            hidden_states = residual + hidden_states
+        else:
+            hidden_states, _ = self.self_attn_layer_norm(hidden_states, residual)
 
-        # Cross-attention (post-norm)
+        # Cross-attention
         residual = hidden_states
+        if self.normalize_before:
+            hidden_states = self.cross_attn_layer_norm(hidden_states)
         hidden_states = self.cross_attn(
             hidden_states=hidden_states,
             encoder_hidden_states=encoder_hidden_states,
             attn_metadata=attn_metadata,
             cross_attn_metadata=cross_attn_metadata,
             skip_cross_kv_projection=skip_cross_kv_projection,
         )
-        hidden_states = residual + hidden_states
-        hidden_states = self.cross_attn_layer_norm(hidden_states)
+        if self.normalize_before:
+            hidden_states = residual + hidden_states
+        else:
+            hidden_states, _ = self.cross_attn_layer_norm(hidden_states, residual)
 
-        # MLP (post-norm)
+        # MLP
         residual = hidden_states
+        if self.normalize_before:
+            hidden_states = self.final_layer_norm(hidden_states)
         hidden_states = self.mlp(hidden_states)
-        hidden_states = residual + hidden_states
-        hidden_states = self.final_layer_norm(hidden_states)
+        if self.normalize_before:
+            hidden_states = residual + hidden_states
+        else:
+            hidden_states, _ = self.final_layer_norm(hidden_states, residual)
 
         return hidden_states
 
@@ -333,15 +373,15 @@ def forward(
 
 
 class BartEncoder(nn.Module):
-    """BART encoder: positional embedding + encoder layers."""
+    """BART/mBART encoder: positional embedding + encoder layers."""
 
     def __init__(self, model_config: ModelConfig[BartConfig]):
         super().__init__()
         config = model_config.pretrained_config
         num_layers = _bart_encoder_num_layers(config)
 
-        # HF BART uses offset=2 for the padding token, so the actual embedding
-        # table has max_position_embeddings + 2 entries.
+        # HF BART/mBART uses offset=2 for the padding token, so the actual
+        # embedding table has max_position_embeddings + 2 entries.
         self.embed_positions = Embedding(
             config.max_position_embeddings + 2,
             config.d_model,
@@ -353,6 +393,16 @@ def __init__(self, model_config: ModelConfig[BartConfig]):
             dtype=config.torch_dtype,
             has_bias=True,
         )
+        self.layer_norm = (
+            LayerNorm(
+                hidden_size=config.d_model,
+                eps=1e-5,
+                dtype=config.torch_dtype,
+                has_bias=True,
+            )
+            if _bart_add_final_layer_norm(config)
+            else None
+        )
         self.layers = nn.ModuleList(
             [BartEncoderLayer(model_config, layer_idx=i) for i in range(num_layers)]
         )
@@ -374,11 +424,13 @@ def forward(
                 attn_metadata=attn_metadata,
                 position_ids=position_ids,
             )
+        if self.layer_norm is not None:
+            hidden_states = self.layer_norm(hidden_states)
         return hidden_states
 
 
 class BartDecoder(nn.Module):
-    """BART decoder: positional embedding + decoder layers."""
+    """BART/mBART decoder: positional embedding + decoder layers."""
 
     def __init__(self, model_config: ModelConfig[BartConfig]):
         super().__init__()
@@ -396,6 +448,16 @@ def __init__(self, model_config: ModelConfig[BartConfig]):
             dtype=config.torch_dtype,
             has_bias=True,
         )
+        self.layer_norm = (
+            LayerNorm(
+                hidden_size=config.d_model,
+                eps=1e-5,
+                dtype=config.torch_dtype,
+                has_bias=True,
+            )
+            if _bart_add_final_layer_norm(config)
+            else None
+        )
         self.layers = nn.ModuleList(
             [BartDecoderLayer(model_config, layer_idx=i) for i in range(num_layers)]
         )
@@ -423,6 +485,8 @@ def forward(
                 cross_attn_metadata=cross_attn_metadata,
                 skip_cross_kv_projection=skip_cross_kv_projection,
             )
+        if self.layer_norm is not None:
+            hidden_states = self.layer_norm(hidden_states)
         return hidden_states
 
 
@@ -623,13 +687,15 @@ def _convert_hf_bart_weights(
         model.shared.weight
         model.encoder.embed_positions.weight
         model.encoder.layernorm_embedding.{weight,bias}
+        model.encoder.layer_norm.{weight,bias}  # mBART
         model.encoder.layers.{i}.self_attn.{q_proj,k_proj,v_proj,out_proj}.{weight,bias}
         model.encoder.layers.{i}.self_attn_layer_norm.{weight,bias}
         model.encoder.layers.{i}.fc1.{weight,bias}
         model.encoder.layers.{i}.fc2.{weight,bias}
         model.encoder.layers.{i}.final_layer_norm.{weight,bias}
         model.decoder.embed_positions.weight
         model.decoder.layernorm_embedding.{weight,bias}
+        model.decoder.layer_norm.{weight,bias}  # mBART
         model.decoder.layers.{i}.self_attn.{q_proj,k_proj,v_proj,out_proj}.{weight,bias}
         model.decoder.layers.{i}.self_attn_layer_norm.{weight,bias}
         model.decoder.layers.{i}.encoder_attn.{q_proj,k_proj,v_proj,out_proj}.{weight,bias}
@@ -675,6 +741,8 @@ def _wb(prefix: str) -> dict:
     # Encoder positional embedding
     out["model.encoder.embed_positions"] = [{"weight": _get(f"{p}encoder.embed_positions.weight")}]
     out["model.encoder.layernorm_embedding"] = [_wb(f"{p}encoder.layernorm_embedding")]
+    if _maybe(f"{p}encoder.layer_norm.weight") is not None:
+        out["model.encoder.layer_norm"] = [_wb(f"{p}encoder.layer_norm")]
 
     # Encoder layers
     for i in range(enc_layers):
@@ -700,6 +768,8 @@ def _wb(prefix: str) -> dict:
     # Decoder positional embedding
     out["model.decoder.embed_positions"] = [{"weight": _get(f"{p}decoder.embed_positions.weight")}]
     out["model.decoder.layernorm_embedding"] = [_wb(f"{p}decoder.layernorm_embedding")]
+    if _maybe(f"{p}decoder.layer_norm.weight") is not None:
+        out["model.decoder.layer_norm"] = [_wb(f"{p}decoder.layer_norm")]
 
     # Decoder layers
     for i in range(dec_layers):