[Feature]: Tensor Parallelism with non divisble amount of attention heads

[NadavShmayo]

🚀 The feature, motivation and pitch

I am trying to run a 70B model on a node with 3XA100-80Gi.
2XA100-80Gi does not contain enough VRAM to run the model, and when I try to run vLLM with tensor parallelism of 3, it raises an error saying that the number of attention heads is not divisble by 3.

I looked into changing the tensor parallelism feature so that it supports an uneven division of the tensors between GPUs.
But I might be missing something here as there are a lot of validations in the codebase to avoid this scenario.
Is it possible to implement tensor parallelism this way?

Alternatives

No response

Additional context

No response

[robertgshaw2-redhat]

Supporting this case will add complications, particularly for:

• managing KV cache memory on each shard
• making sure attention happens locally on each shard
• loading quantized weights

I think in theory this could be done, but it would require quite a bit of refactoring to many parts of the vllm system to support it

[theycallmeloki]

Just coming here from having installed 7xGPUs on a Threadripper and I'm unable to run llama 3 70b on 7x24 GB cards so I have a valid use for this and would like to help with what I can as well

Can perhaps guide me vaguely on problems I should be aware of in attempting to make this change happen I was thinking to start from reading the tests and going from there to get an idea of how to make this change happen

I would also like to have a way to use llama 3 8b or something like that locally while I am able to piece together this for vllm so if there's any alternative that might work like llama.cpp or any other project that might intermittently help me have an idea of how it's implemented without this particular restriction

[robertgshaw2-redhat]

Just coming here from having installed 7xGPUs on a Threadripper and I'm unable to run llama 3 70b on 7x24 GB cards so I have a valid use for this and would like to help with what I can as well

Can perhaps guide me vaguely on problems I should be aware of in attempting to make this change happen I was thinking to start from reading the tests and going from there to get an idea of how to make this change happen

I would also like to have a way to use llama 3 8b or something like that locally while I am able to piece together this for vllm so if there's any alternative that might work like llama.cpp or any other project that might intermittently help me have an idea of how it's implemented without this particular restriction

The biggest problem that you will have is that we have a central KV cache management system which manages the mapping from physical to logical blocks. There is a large implicit assumption that each shard has the same mapping since there is one set of metadata passed to each shard during the forward pass

It will require deep understanding and refactoring of the entire system to unwind this assumption

[NadavShmayo]

Just coming here from having installed 7xGPUs on a Threadripper and I'm unable to run llama 3 70b on 7x24 GB cards so I have a valid use for this and would like to help with what I can as well
Can perhaps guide me vaguely on problems I should be aware of in attempting to make this change happen I was thinking to start from reading the tests and going from there to get an idea of how to make this change happen
I would also like to have a way to use llama 3 8b or something like that locally while I am able to piece together this for vllm so if there's any alternative that might work like llama.cpp or any other project that might intermittently help me have an idea of how it's implemented without this particular restriction

The biggest problem that you will have is that we have a central KV cache management system which manages the mapping from physical to logical blocks. There is a large implicit assumption that each shard has the same mapping since there is one set of metadata passed to each shard during the forward pass

It will require deep understanding and refactoring of the entire system to unwind this assumption

I'd be happy to try and implement this, but I'm less familiar with this part of the vLLM codebase.
Can you please refer me to the primary files I'd have to change for this, or give a brief explanation of the central KV cache management system you mentioned?
(I am familiar with PagedAttention and block management, but less in the distributed inference context).

[robertgshaw2-redhat]

Just coming here from having installed 7xGPUs on a Threadripper and I'm unable to run llama 3 70b on 7x24 GB cards so I have a valid use for this and would like to help with what I can as well
Can perhaps guide me vaguely on problems I should be aware of in attempting to make this change happen I was thinking to start from reading the tests and going from there to get an idea of how to make this change happen
I would also like to have a way to use llama 3 8b or something like that locally while I am able to piece together this for vllm so if there's any alternative that might work like llama.cpp or any other project that might intermittently help me have an idea of how it's implemented without this particular restriction

The biggest problem that you will have is that we have a central KV cache management system which manages the mapping from physical to logical blocks. There is a large implicit assumption that each shard has the same mapping since there is one set of metadata passed to each shard during the forward pass
It will require deep understanding and refactoring of the entire system to unwind this assumption

I'd be happy to try and implement this, but I'm less familiar with this part of the vLLM codebase. Can you please refer me to the primary files I'd have to change for this, or give a brief explanation of the central KV cache management system you mentioned? (I am familiar with PagedAttention and block management, but less in the distributed inference context).

You're going to have to touch many subsystems.

• The Scheduler, BlockManager, and Attention to get a sense of how the KVCache is managed.
• Worker + ModelExecutor handle the tensor parallelism

[NadavShmayo]

Just coming here from having installed 7xGPUs on a Threadripper and I'm unable to run llama 3 70b on 7x24 GB cards so I have a valid use for this and would like to help with what I can as well
Can perhaps guide me vaguely on problems I should be aware of in attempting to make this change happen I was thinking to start from reading the tests and going from there to get an idea of how to make this change happen
I would also like to have a way to use llama 3 8b or something like that locally while I am able to piece together this for vllm so if there's any alternative that might work like llama.cpp or any other project that might intermittently help me have an idea of how it's implemented without this particular restriction

The biggest problem that you will have is that we have a central KV cache management system which manages the mapping from physical to logical blocks. There is a large implicit assumption that each shard has the same mapping since there is one set of metadata passed to each shard during the forward pass
It will require deep understanding and refactoring of the entire system to unwind this assumption

I'd be happy to try and implement this, but I'm less familiar with this part of the vLLM codebase. Can you please refer me to the primary files I'd have to change for this, or give a brief explanation of the central KV cache management system you mentioned? (I am familiar with PagedAttention and block management, but less in the distributed inference context).

You're going to have to touch many subsystems.

• The Scheduler, BlockManager, and Attention to get a sense of how the KVCache is managed.
• Worker + ModelExecutor handle the tensor parallelism

Hey, I implement this feature (:
Would really appreciate your feedback on this #5367

[ArmykOliva]

Is there any way to use tensor parallelism with uneven amount of GPUs?

[github-actions]

This issue has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this issue should remain open. Thank you!

[github-actions]

This issue has been automatically closed due to inactivity. Please feel free to reopen if you feel it is still relevant. Thank you!

[groxaxo]

this has worked for me, i've got 3 3090 . python -m vllm.entrypoints.openai.api_server
--model ./stelterlab_openhands-lm-32b-v0.1-AWQ
--tensor-parallel-size 1
--pipeline-parallel-size 3
--quantization awq_marlin
--dtype float16
--max-model-len 53000
--gpu-memory-utilization 0.95
--max-num-batched-tokens 4096
--max-num-seqs 32
--enable-prefix-caching
--swap-space 0
--host 0.0.0.0
--port 8000