fix gqa

utils.py

@@ -76,7 +76,9 @@ def calculate_memory_components(
     overhead = 72 + 32 * mbs
 
     # Activations
-    decoder_layer_mib = (seq_len * mbs * hidden_size/tp) * (2/1024/1024) * (4*intermediate_size/hidden_size + 10)
+    # decoder_layer_mib = (seq_len * mbs * hidden_size/tp) * (2/1024/1024) * (4*intermediate_size/hidden_size + 10)
+    is_mha = num_key_value_heads == num_attention_heads
+    decoder_layer_mib = (seq_len * mbs * hidden_size/tp) * (2/1024/1024) * (4*intermediate_size/hidden_size + 12 + 2*num_key_value_heads/num_attention_heads + (2 if is_mha else 0))
     
     if pp > 1:
         activs = min(pp, batch_accum) * num_hidden_layers_in_pp * decoder_layer_mib