# Why tensor parallelism fails to split evenly across GPUs, and what TP size and hardware match actually requires

> Source: https://netray.co/troubleshoot/tensor-parallelism-uneven-gpu-split
> Publisher: Netray (netray.co), AI and integration engineering for ERP and manufacturing systems.
> Section: On-Prem AI Troubleshooting

## Short answer

Tensor parallelism fails when the chosen degree does not evenly divide the model's attention heads, and often its key/value heads and hidden size, so the framework cannot split the projection weights across ranks. It also fails in practice, without an assertion, when the GPUs assigned to the TP group differ in VRAM or compute, since the even weight split then fits some ranks and not others.

## Error Title

Tensor parallelism fails because the model does not split evenly across GPUs

## Error Signature

AssertionError: Number of attention heads (32) must be divisible by tensor parallel size (6)

## Alternate Signatures

- ValueError: total number of attention heads (32) is not divisible by tensor_parallel_size (6)
- RuntimeError: shape mismatch when splitting weight tensor across tensor parallel ranks

## Category

distributed

## Stack

- vllm
- pytorch
- cuda

## Affects

vLLM, TensorRT-LLM, DeepSpeed, and any transformers-based tensor-parallel setup, on any GPU model, most visible when TP size is not a power of two or GPUs are mixed models

## Quick Fix

### Match TP size to the model's head count and hardware first

**Steps:**

Print the model's actual head counts with python -c "from transformers import AutoConfig; c=AutoConfig.from_pretrained('MODEL'); print(c.num_attention_heads, getattr(c,'num_key_value_heads',None))" before picking a TP degree.

Choose a tensor-parallel size that evenly divides num_attention_heads, and if the model uses grouped-query attention, confirm it also divides num_key_value_heads.

Confirm every GPU in the tensor-parallel group is the identical model and VRAM size with nvidia-smi --query-gpu=name,memory.total --format=csv.

If you must use an odd number of GPUs, prefer combining tensor parallelism with pipeline or data parallelism (for example TP=4 with PP=2) instead of forcing an indivisible TP size.

Re-launch with the corrected --tensor-parallel-size (vLLM) or equivalent flag and confirm the server starts and reports even memory usage across all ranks.

## Symptoms

- Server or training job fails immediately at startup with an assertion about head counts or tensor parallel size
- Startup succeeds but one GPU in the group reports far higher memory usage than the others
- Works fine with TP=2 or TP=4 but fails or OOMs the moment you try TP=3, 5, 6, or 7
- Mixing two different GPU models in the same node causes one rank to OOM while the others have free VRAM

## Causes

### The chosen tensor-parallel size does not evenly divide the model's number of attention heads or hidden size

**Likelihood:** most-common

**Explanation:** Tensor parallelism works by splitting each attention head's projection matrices across ranks so each GPU computes a subset of heads. If the head count is not a multiple of the TP degree, at least one rank would need a fractional number of heads, which the sharding logic cannot express, so it raises an assertion instead of silently doing something wrong.

**Fix:** Pick a TP size from the divisors of the model's attention head count (for a 32-head model that is 1, 2, 4, 8, 16, or 32), not an arbitrary GPU count like 6, and verify with the model config before launching.

**Commands:**

- python -c "from transformers import AutoConfig; c=AutoConfig.from_pretrained('meta-llama/Llama-3.1-70B'); print(c.num_attention_heads)"

### GPUs in the tensor-parallel group are different models or have different VRAM capacities

**Likelihood:** common

**Explanation:** Tensor parallelism splits weights and activations equally by design, assuming every rank has the same compute and memory budget. Mixing a 24GB and a 48GB card in the same TP group means both ranks get the same shard size, but the 24GB card runs out of headroom for KV cache and activations long before the 48GB card does, causing an OOM on only one rank.

**Fix:** Only place identical GPU models with identical VRAM in the same tensor-parallel group; if your cluster is heterogeneous, group nodes by GPU type and route jobs to matching hardware, or reduce TP size to fit the smallest GPU's budget.

### The model uses grouped-query or multi-query attention, so TP size must also divide the smaller key/value head count

**Likelihood:** common

**Explanation:** Many current models reduce key/value heads relative to query heads to shrink KV cache size. Tensor parallelism must still assign whole KV heads to each rank, so a TP degree that divides the query head count but not the often much smaller KV head count still fails or silently mishandles KV heads.

**Fix:** Check num_key_value_heads specifically, not just num_attention_heads, and choose a TP size that divides the smaller of the two; for an 8 KV-head model, valid TP sizes are limited to divisors of 8.

### Vocabulary size or hidden dimension is not divisible by the TP size

**Likelihood:** occasional

**Explanation:** Frameworks typically pad the vocabulary to a TP-friendly size internally, but custom model code, older checkpoints, or hand-rolled parallelism implementations may not pad correctly, producing a shape mismatch specifically in the embedding or output projection layer rather than the attention layer.

**Fix:** Check hidden_size and vocab_size against your TP degree, and if using custom parallelism code, add explicit padding for the vocabulary dimension to the next multiple of the TP size.

### Custom sharding code assumes a power-of-two TP size and hardcodes bit-shift logic instead of general integer division

**Likelihood:** rare

**Explanation:** Some hand-written or older tensor-parallel implementations use bit operations (assuming TP is always 2, 4, 8, or 16) for speed, which silently produces incorrect results or crashes when someone later tries TP=3 or TP=6 on a non-power-of-two GPU count.

**Fix:** Review any custom sharding code for hardcoded power-of-two assumptions and replace with general modulo or division checks, or restrict deployment to power-of-two TP sizes until the code is fixed.

## Diagnostics

### Print the model's head and KV-head counts

**Command:** python -c "from transformers import AutoConfig; c=AutoConfig.from_pretrained('MODEL_NAME'); print('heads:',c.num_attention_heads,'kv_heads:',getattr(c,'num_key_value_heads',c.num_attention_heads))"

**Interpretation:** Compare both numbers against your intended TP size; your TP size must divide both, and if it divides only the head count, you will hit the KV-head sharding problem instead.

### Confirm all GPUs in the group are identical

**Command:** nvidia-smi --query-gpu=index,name,memory.total --format=csv

**Interpretation:** Every row should show the same GPU name and the same memory.total; any row that differs means that rank has a different memory budget and will OOM or underuse VRAM relative to the rest of the group.

### Check the actual launch flags being used

**Command:** ps aux | grep tensor-parallel-size

**Interpretation:** Confirms the TP size the process actually started with, which is useful when a wrapper script or orchestrator overrides what you think you configured.

## Prevention

- Standardize on a small set of TP sizes (2, 4, 8) that always divide your target models' head counts, and pick models with convenient head counts where possible.
- Tag GPUs in your inventory by model and VRAM so scheduling can guarantee a tensor-parallel group is always homogeneous.
- Add a pre-flight check to your deployment pipeline that reads the model config and rejects an incompatible TP size before attempting to launch.
- When procuring GPUs, buy in matched sets sized for your typical tensor-parallel group size, rather than adding one or two odd GPUs later.

## When To Escalate

If your GPU inventory is inherently heterogeneous, mixed generations or VRAM sizes acquired over time, and that mismatch is what is forcing awkward TP sizes, this becomes a hardware and cluster design question, not a config fix, since no amount of software tuning makes uneven hardware split evenly.

## Faq

### What TP sizes are valid for a 32-head model?

**Answer:** Any divisor of 32: 1, 2, 4, 8, 16, or 32. A TP size like 3, 5, 6, or 7 will fail immediately because it cannot split 32 heads into equal integer groups per GPU. If your GPU count does not match a divisor, combine tensor parallelism with pipeline or data parallelism instead.

### Can I use tensor parallelism across two different GPU models?

**Answer:** Technically the framework may let you launch, but it is not recommended. Tensor parallelism assumes identical per-rank compute and memory budgets, so a weaker or smaller-VRAM GPU in the group becomes the bottleneck or the one that OOMs, even though the weight split itself is mathematically even.

### Does grouped-query attention change what TP sizes are valid?

**Answer:** Yes. Your TP size must divide the number of key/value heads, which is often smaller than the number of query heads in models using grouped-query or multi-query attention. Check both numbers in the model config before choosing a TP degree.

## Search terms

tensor parallel size error, attention heads divisible by tensor parallel size, vllm tensor parallel size, uneven gpu split tensor parallelism, tensor parallelism gpu mismatch, num_key_value_heads tensor parallel

---
Talk to an engineer: https://netray.co/contact?intent=erp_ai | Machine catalog: https://netray.co/agent/catalog.json