doc: add n_query_groups to attention notation table (#2092)

Author: raishish

litgpt/model.py

@@ -377,8 +377,26 @@ def forward(
         # - T          | time-step (sequence length)
         # - C          | model's embeddings size (n_embd)
         # - C*         | attentions's embeddings size
-        # - nh_(q,k,v) | number of heads for query, key and value
         # - hs         | head size
+        # - nh_(q,k,v) | number of heads for query, key and value
+        # - n_query_groups = nh_k = nh_v | number of query groups sharing key and value heads
+        # alternative notation: num_kv_groups = n_query_groups
+        # ┌───┐┌───┐┌───┐┌───┐     ┌───┐    ┌───┐             ┌───┐
+        # │ v ││ v ││ v ││ v │     │ v │    │ v │             │ v │
+        # └───┘└───┘└───┘└───┘     └───┘    └───┘             └───┘
+        #   │    │    │    │         │        │                 │
+        # ┌───┐┌───┐┌───┐┌───┐     ┌───┐    ┌───┐             ┌───┐
+        # │ k ││ k ││ k ││ k │     │ k │    │ k │             │ k │
+        # └───┘└───┘└───┘└───┘     └───┘    └───┘             └───┘
+        #   │    │    │    │      ┌──┴──┐  ┌──┴──┐      ┌────┬──┴─┬────┐
+        # ┌───┐┌───┐┌───┐┌───┐  ┌───┐┌───┐┌───┐┌───┐  ┌───┐┌───┐┌───┐┌───┐
+        # │ q ││ q ││ q ││ q │  │ q ││ q ││ q ││ q │  │ q ││ q ││ q ││ q │
+        # └───┘└───┘└───┘└───┘  └───┘└───┘└───┘└───┘  └───┘└───┘└───┘└───┘
+        # ◀──────────────────▶  ◀──────────────────▶  ◀──────────────────▶
+        #         MHA                    GQA                   MQA
+        #   n_query_groups=4       n_query_groups=2      n_query_groups=1
+        #
+        # credit https://arxiv.org/pdf/2305.13245.pdf
         head_size = self.config.head_size
         n_head = self.config.n_head
         n_query_groups = self.config.n_query_groups