diff --git a/megatron/training/training.py b/megatron/training/training.py
index 34e65337703..cf1cdda2685 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -573,7 +573,7 @@ def transformer_flops():
                 ## for in the dsv4_hybrid branch below.
                 q_term = args.q_lora_rank * (
                     args.hidden_size
-                    + args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim)
+                    + args.num_attention_heads * args.v_head_dim
                     + 1  # q norm
                 )
                 kv_term = args.hidden_size * args.v_head_dim + args.v_head_dim  # kv proj + kv norm