68 lines
2.5 KiB
Markdown
68 lines
2.5 KiB
Markdown
# 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/`.
|
|
|
|
```bash
|
|
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
|
|
```
|
|
|
|
```bash
|
|
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:
|
|
|
|
```bash
|
|
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.
|
|
|
|
```bash
|
|
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
|
|
|
|
```bash
|
|
python scripts/benchmark_gqa_vs_mla.py --device cuda:0 --hidden_size 768
|
|
```
|
|
|
|
With checkpoints:
|
|
|
|
```bash
|
|
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.
|