llama2

GeLU

For more detail, check GeLU.

RMSNorm

In recent settings, RMSNorm is typically placed before a layer, a.k.a pre-norm. One possible explanation is the increase of stability for the training process.

In the original RMSNorm paper, the authors claimed that recentering step in layer norm yields insignificant boost performance, while removing it speeds up the computation noticeably.

GQA

GQA(Grouped Query Attention), is a balance between multi-head attention(MHA) from vanilla transformers and MQA(Multi Query Attention), leveraging the advantages from both sides.

Grouped query attention is somewhat an interpolation of MHA and MQA with the goal of balancing the pros and cons. Tokens are first grouped together according to positions and share key and value vectors within groups.

The following table is a direct comparison between dimensions of query, key, value and output for MHA and MQA. $B$ is batch size, $T$ is sequence length, $d$ is embedding dimension, $\text{n\_heads}$ is number of heads, $d_k=d/\text{n\_heads}$ is query/key/value vector dimension, $\text{n\_kv\_heads}$ is the number of key/value heads and $G=\text{n\_heads/n\_kv\_heads}$ is the number of groups.

	MHA	MQA	GQA
$X$	$(B,T,d)$	$(B,T,d)$	$(B,T,d)$
$Q$	$(B,\text{n\_heads},T,d_k)$	$(B,\text{n\_heads},T,d_k)$	$(B,G,\text{n\_kv\_heads},T,d_k)$
$K$	$(B,\text{n\_heads},T,d_k)$	$(B,\text{1},T,d_k)$	$(B,\text{n\_kv\_heads},T,d_k)$
$V$	$(B,\text{n\_heads},T,d_k)$	$(B,\text{1},T,d_k)$	$(B,\text{n\_kv\_heads},T,d_k)$
$O$	$(B,T,d)$	$(B,T,d)$	$(B,\text{n\_kv\_heads},T,d_k)$
$W_O$	$(d,d)$	$(d,d)$	$\text{n\_kv\_heads},d$
Y	$(B,T,d)$	$(B,T,d)$	$(B,T,d)$