Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
Table 7: Effect of number of attention heads on scaling on 8.3 billion of parameters with 8-way model parallelism.
['Attention heads', 'Hidden size per head', 'Scaling Efficiency']
This section studies the effect of attention heads on model parallel scaling. To this end, we consider the 8.3 billion parameter configuration with 8-way model parallelism and vary the number of heads from 16 to 32. As the number of attention heads increases, some of the GEMMS inside the self-attention layer become smaller and also the number of elements in the self attention softmax increases. This results in a slight decrease in scaling efficiency. Future research should be wary of this hyperparameter to design large transformer models that balance model speed and model accuracy.