MLKV: Multi-Layer Key-Value Sharing

Experiments on EleutherAI's Pythia models

Setup

git clone https://github.com/zaydzuhri/pythia-mlkv.git
cd pythia-mlkv
pip install -r requirements.txt

Convert Pythia models to MQA/GQA/MLKV models

git lfs install
git clone https://huggingface.co/EleutherAI/pythia-160m-deduped
rm -rf pythia-160m-deduped/.git
python3 convert_to_mlkv.py --weights_path pythia-160m-deduped --num-key-value-layers 6 --num-key-value-heads 1

Here are all the 8+1 configs needed for all experiments:

Name	Num. of layers	Num. of attention heads	Num. of layers with KV heads (num-key-value-layers)	Num. of KV heads in a layer (num-key-value-heads)	Total num. of KV heads	Num. of parameters
MHA-144	12	12	12	12	144	160M
GQA-48	12	12	12	4	48	160M
MLKV-48	12	12	4	12	48	160M
MQA-12	12	12	12	1	12	160M
MLKV-12	12	12	4	3	12	160M
MLKV-6	12	12	6	1	6	160M
MLKV-4	12	12	4	1	4	160M
MLKV-2	12	12	2	1	2	160M
MLKV-1	12	12	1	1	1	160M

Uptraining

The dataset has been prepared to Huggingface, so you can directly uptrain:

CUDA_VISIBLE_DEVICES=0,1 python3 uptrain.py --output-dir pythia-160m-mlkv-6-b12-g2-v1 --model pythia-160m-deduped_mlkv_6_1 --batch-size 12 --gradient-accumulate-every 1 --learning-rate 6e-4 --warmup-ratio 0.2  --wandb pythia-160m-mlkv-6-b12-g2-v1