Why tensor parallelism fails to split evenly across GPUs, and what TP size and hardware match actually requires
AssertionError: Number of attention heads (32) must be divisible by tensor parallel size (6)
Also appears as
- 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
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.
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
Match TP size to the model's head count and hardware first
- 1Print 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.
- 2Choose 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.
- 3Confirm every GPU in the tensor-parallel group is the identical model and VRAM size with nvidia-smi --query-gpu=name,memory.total --format=csv.
- 4If 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.
- 5Re-launch with the corrected --tensor-parallel-size (vLLM) or equivalent flag and confirm the server starts and reports even memory usage across all ranks.
How to confirm this is your problem
- 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
Root causes and fixes
The chosen tensor-parallel size does not evenly divide the model's number of attention heads or hidden size
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.
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
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
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
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
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.
Diagnostic commands
Print the model's head and KV-head counts
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))"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
nvidia-smi --query-gpu=index,name,memory.total --format=csv
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
ps aux | grep tensor-parallel-size
Confirms the TP size the process actually started with, which is useful when a wrapper script or orchestrator overrides what you think you configured.
Stopping it from happening again
- 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 this becomes an architecture problem
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.
Frequently asked questions
What TP sizes are valid for a 32-head model?
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?
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?
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.
Size it properly next time
Free calculators that prevent this class of failure before you provision hardware.
Multi-GPU Tensor Parallelism Calculator
Model how tensor-parallel throughput actually scales across multiple GPUs, accounting for interconnect overhead that keeps scaling sub-linear.
Free ToolGPU Sizing Calculator for LLM Inference
Work out how many GPUs you need to serve a given open-weight model to your user base, based on memory footprint and token throughput.
Free ToolNVIDIA GPU Selector for LLM Workloads
Score your workload across model size, concurrency, latency, budget, and facility power to get a recommended GPU tier from RTX-class to multi-node B200 clusters.
Related problems
NCCL error during multi-GPU training or inference
An NCCL error during multi-GPU training or inference is almost always a symptom of a rank that crashed, a version mismatch across processes, or bad GPU topology, not a bug in NCCL itself. Enable NCCL_DEBUG=INFO first and read the per-rank logs before touching timeouts or retry logic.
vLLM tensor-parallel-size must divide the number of attention heads
vLLM shards attention heads and KV heads evenly across GPUs for tensor parallelism, so --tensor-parallel-size must be a divisor of the model's head count and must equal the number of GPUs vLLM can actually see. Pick a TP size from the model's valid divisors, commonly 1, 2, 4, or 8, and make sure CUDA_VISIBLE_DEVICES exposes exactly that many GPUs.
Pipeline parallelism scales poorly, throughput does not improve with more stages
Pipeline parallelism scales poorly when the number of microbatches per training step is too small relative to the number of pipeline stages, because the unavoidable fill and drain bubble at the start and end of each step is proportional to the stage count minus one divided by the number of microbatches. With too few microbatches, GPUs spend a large fraction of every step idle waiting for the pipeline to fill or drain, and adding more stages without also adding more microbatches makes this worse, not better.
GPU peer-to-peer (P2P) access not working between GPUs on the same node
GPU peer-to-peer (P2P) access fails when the PCIe topology, IOMMU or ACS settings, or the GPU model itself does not support a direct memory path between two devices, forcing all transfers through the CPU and host memory instead of GPU-to-GPU. Setting NCCL_P2P_DISABLE=1 is a useful diagnostic to confirm P2P is the problem, but it only removes the crash by falling back to a slower path; it does not restore the P2P bandwidth you actually need for good multi-GPU performance.
GuideMulti-GPU LLM Serving: Tensor vs Pipeline Parallelism
Multi-GPU LLM serving explained: tensor parallelism vs pipeline parallelism, NCCL interconnect requirements, and when to split a model across GPUs.
GuidevLLM Production Deployment: A Practitioner's Guide
Deploy vLLM in production: continuous batching, PagedAttention, config flags that matter, and the metrics to watch before you trust it with real traffic.
GuideNVIDIA H100 vs H200 vs B200 for Enterprise AI in 2026
Compare NVIDIA H100, H200, and B200 GPUs on specs, price, availability, and performance per dollar for enterprise LLM inference and training in 2026.
Still stuck, or tired of fighting your own infrastructure?
Netray deploys and operates on-prem AI for regulated manufacturers and defense suppliers. We have debugged this stack in production, on air-gapped networks, at scale.