minimind/docs/attention_experiment.md

Attention Architecture Experiments

This project supports four attention variants for controlled small-LLM experiments:

• mha: Multi-Head Attention, one K/V head per query head.
• gqa: Grouped-Query Attention, the default MiniMind-style setting.
• mqa: Multi-Query Attention, one shared K/V head.
• mla: Multi-head Latent Attention with low-rank KV cache and decoupled RoPE.

Training Commands

Run all commands from trainer/.

torchrun --nproc_per_node=4 train_pretrain.py --attention_type gqa
torchrun --nproc_per_node=4 train_pretrain.py --attention_type mqa --save_weight pretrain_mqa
torchrun --nproc_per_node=4 train_pretrain.py --attention_type mla --save_weight pretrain_mla

torchrun --nproc_per_node=4 train_full_sft.py --attention_type gqa --from_weight pretrain
torchrun --nproc_per_node=4 train_full_sft.py --attention_type mqa --from_weight pretrain_mqa --save_weight full_sft_mqa
torchrun --nproc_per_node=4 train_full_sft.py --attention_type mla --from_weight pretrain_mla --save_weight full_sft_mla

The legacy MLA flag still works:

torchrun --nproc_per_node=4 train_pretrain.py --use_mla 1

RL Stages

The same --attention_type argument is available in DPO, PPO, GRPO, LoRA, distillation, Agent RL, and DeepSpeed pretraining.

torchrun --nproc_per_node=4 train_dpo.py --attention_type gqa --from_weight full_sft
torchrun --nproc_per_node=4 train_ppo.py --attention_type gqa --from_weight full_sft
torchrun --nproc_per_node=4 train_grpo.py --attention_type gqa --from_weight full_sft

Benchmark

python scripts/benchmark_gqa_vs_mla.py --device cuda:0 --hidden_size 768

With checkpoints:

python scripts/benchmark_gqa_vs_mla.py \
  --device cuda:0 \
  --hidden_size 768 \
  --gqa_ckpt out/pretrain_768.pth \
  --mla_ckpt out/pretrain_mla_768_mla.pth \
  --data_path dataset/pretrain_t2t_mini.jsonl \
  --tokenizer_path model

Key metrics to record for a resume or report:

• Train setup: GPU type/count, batch size, sequence length, tokens trained.
• Architecture: total params, attention params/layer, KV-cache floats/token/layer.
• Efficiency: prefill latency, decode tokens/s, training tokens/s, peak GPU memory.
• Quality: pretrain loss, SFT PPL, DPO/PPO/GRPO reward and KL curves.

For hidden_size=768, the default GQA KV cache is 2 * 4 * 96 = 768 floats/token/layer. MLA with kv_lora_rank=128 and rope_dim=48 stores 128 + 48 = 176 floats/token/layer, a theoretical 4.4x KV-cache compression.