test using CONV1D

Author: ChaoPang

src/cehrbert/models/hf_models/hf_cehrbert.py

@@ -8,6 +8,7 @@
 from transformers.activations import gelu_new
 from transformers.models.bert import modeling_bert
 from transformers.models.bert.modeling_bert import BertEncoder, BertOnlyMLMHead, BertPooler
+from transformers.pytorch_utils import Conv1D
 from transformers.utils import is_flash_attn_2_available, logging
 
 if is_flash_attn_2_available():
@@ -32,12 +33,9 @@ def __init__(self, config, position_embedding_type=None):
         self.num_attention_heads = config.num_attention_heads
         self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
         self.all_head_size = self.num_attention_heads * self.attention_head_size
-
-        self.query = nn.Linear(config.hidden_size, self.all_head_size)
-        self.key = nn.Linear(config.hidden_size, self.all_head_size)
-        self.value = nn.Linear(config.hidden_size, self.all_head_size)
-
-        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+        self.split_size = config.hidden_size
+        self.c_attn = Conv1D(3 * config.hidden_size, config.hidden_size)
+        self.dropout_rate = config.attention_probs_dropout_prob
 
     def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
         new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
@@ -58,8 +56,9 @@ def forward(
         dtype = hidden_states.dtype
 
         batch_size = hidden_states.size(0)
+        query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2)
         # (batch, n_heads, seq_length, head_dim)
-        query_layer = self.transpose_for_scores(self.query(hidden_states))
+        query_layer = self.transpose_for_scores(query)
         # If this is instantiated as a cross-attention module, the keys
         # and values come from an encoder; the attention mask needs to be
         # such that the encoder's padding tokens are not attended to.
@@ -71,21 +70,17 @@ def forward(
             value_layer = past_key_value[1]
             attention_mask = encoder_attention_mask
         elif is_cross_attention:
-            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
-            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
+            key_layer = self.transpose_for_scores(key)
+            value_layer = self.transpose_for_scores(value)
             attention_mask = encoder_attention_mask
         elif past_key_value is not None:
-            key_layer = self.transpose_for_scores(self.key(hidden_states))
-            value_layer = self.transpose_for_scores(self.value(hidden_states))
+            key_layer = self.transpose_for_scores(key)
+            value_layer = self.transpose_for_scores(value)
             key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
             value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
         else:
-            key_layer = self.transpose_for_scores(self.key(hidden_states))
-            value_layer = self.transpose_for_scores(self.value(hidden_states))
-
-        query_layer = query_layer.to(torch.bfloat16)
-        key_layer = key_layer.to(torch.bfloat16)
-        value_layer = value_layer.to(torch.bfloat16)
+            key_layer = self.transpose_for_scores(key)
+            value_layer = self.transpose_for_scores(value)
 
         # Flash Attention forward pass
         # Use the built-in scaled_dot_product_attention with Flash Attention
@@ -95,7 +90,7 @@ def forward(
                 key_layer,
                 value_layer,
                 attn_mask=attention_mask,
-                dropout_p=self.dropout.p,
+                dropout_p=self.dropout_rate,
                 is_causal=False,
                 scale=None,  # Default is 1/sqrt(head_dim)
             )
@@ -270,7 +265,7 @@ class CehrBertPreTrainedModel(PreTrainedModel):
 
     def _init_weights(self, module):
         """Initialize the weights."""
-        if isinstance(module, nn.Linear):
+        if isinstance(module, (nn.Linear, Conv1D)):
             # Slightly different from the TF version which uses truncated_normal for initialization
             # cf https://github.com/pytorch/pytorch/pull/5617
             module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)