Fix t5 encoder tests

Author: kevin314

fastvideo/models/encoders/t5.py

@@ -181,7 +181,8 @@ def __init__(self,
 
         self.qkv_proj = QKVParallelLinear(
             self.d_model,
-            self.d_model // self.total_num_heads,
+            #self.d_model // self.total_num_heads,
+            self.key_value_proj_dim,
             self.total_num_heads,
             self.total_num_kv_heads,
             bias=False,
@@ -199,7 +200,8 @@ def __init__(self,
                                        padding_size=self.relative_attention_num_buckets,
                                        quant_config=quant_config)
         self.o = RowParallelLinear(
-            self.d_model,
+            #self.d_model,
+            self.total_num_heads * self.key_value_proj_dim,
             self.d_model,
             bias=False,
             quant_config=quant_config,
@@ -298,10 +300,12 @@ def forward(
     ) -> torch.Tensor:
         bs, seq_len, _ = hidden_states.shape
         num_seqs = bs
-        n, c = self.n_heads, self.d_model // self.total_num_heads
+        #n, c = self.n_heads, self.d_model // self.total_num_heads
+        n, c = self.n_heads, self.key_value_proj_dim
         qkv, _ = self.qkv_proj(hidden_states)
         # Projection of 'own' hidden state (self-attention). No GQA here.
-        q, k, v = qkv.split(self.inner_dim, dim=-1)
+        #q, k, v = qkv.split(self.inner_dim, dim=-1)
+        q, k, v = qkv.split(self.qkv_proj.output_sizes, dim=-1)
         q = q.reshape(bs, seq_len, n, c)
         k = k.reshape(bs, seq_len, n, c)
         v = v.reshape(bs, seq_len, n, c)

fastvideo/tests/encoders/test_t5_encoder.py

@@ -169,6 +169,15 @@ def test_t5_large_encoder():
     # Check number of parameters
     logger.info("Model1 has %s parameters", len(params1))
     logger.info("Model2 has %s parameters", len(params2))
+    
+    # # Print parameter names for comparison
+    # logger.info("Model1 parameters:")
+    # for name in sorted(params1.keys()):
+    #     logger.info("  %s: %s", name, params1[name].shape)
+    
+    # logger.info("Model2 parameters:")
+    # for name in sorted(params2.keys()):
+    #     logger.info("  %s: %s", name, params2[name].shape)
 
     weight_diffs = []
     # check if embed_tokens are the same