Hi. I was researching a way to make tensor parallel possible in wide variety of hardware configurations. The problem it follow the pattern 2 4 8 , which makes 3 5 6 7 … Unusable and then the other choice is to go for pipeline parallel. The problem with pipeline parallel is that is is not the best case and not good for training also.
My Idea is to have the attention heads reduced from 64 to 60 for large models and the head size reduced to 3840. Those two numbers unlock many advantages compared to other head prune options.
3840 is a mathematically excellent choice. It is a highly composite number that divides perfectly by 2, 3, 4, and 5.
Using 3840 as your hidden size gives you immense architectural flexibility for configuring your attention heads while keeping your GPU running at peak efficiency.
Head Configurations Allowed by 3840
If you keep the industry-standard head dimension (head_dim = 64), a hidden size of 3840 gives you exactly 60 attention heads (3840 \div 64 = 60).
Because 60 is divisible by 2, 3, 4, and 5, you can group your attention heads perfectly in multiple ways: [1]
- Divisible by 2: 30 pairs of heads
- Divisible by 3: 20 triplets of heads
- Divisible by 4: 15 groups of 4 heads
- Divisible by 5: 12 groups of 5 heads
Why 3840 is Great for Grouped-Query Attention (GQA)
This extreme divisibility makes 3840 perfect for modern GQA optimization. If you have 60 Query (Q) heads, you can compress your Key/Value (KV) cache into clean integer groupings without any fractional remainders:
- 4 KV Groups: 15 Q heads per KV group (Divisible by 4)
- 5 KV Groups: 12 Q heads per KV group (Divisible by 5)
- 6 KV Groups: 10 Q heads per KV group (Divisible by 6)
GPU Performance: Why 3840 Beats 4032
Unlike 4032, 3840 is a multiple of 128 (3840 \div 128 = 30) and 256 (3840 \div 256 = 15).
This completely eliminates the memory throughput penalties discussed earlier:
- Perfect Memory Alignment: 3840 aligns perfectly with 128-byte GPU memory cache lines. There is zero wasted bandwidth when transferring weights from HBM to SRAM.
- FlashAttention Compatibility: Most highly optimized inference engines (like vLLM and FlashAttention-3) split matrix operations into blocks of 128 or 256 tokens/channels. 3840 natively fits into these fast-path kernels.
Any thoughts on this. I think Nvidia modelopt can handle this job, but I don’t know if it is possible on even number of nodes, because calibration and other runs needs to be done first which is not possible.