Can we support multi GPUs when doing inference?

[sparkFeiYang]

Now I tried run this 27b Gemma3 model on 2*40G VRAM A100 GPUs. If I run this script directly, it will shows out of memory, because bf16 model is 55G and one A100 can't afford this. I use DeepSpeed to modify the script, and finally I succeed run this model on 2 A100 GPUs. But the time cost is too much for each run, especially when I set large output_len like 1500, and this will cost about 2 hours for 1 prompt run with 27b bf16.

[Balakrishna-Chennamsetti]

Hi @sparkFeiYang ,

That is a common challenge with running very large models like Gemma 27B, even with powerful hardware like A100s, especially when generating long sequences.

You've already confirmed the necessity and effectiveness of multi-GPU inference by using DeepSpeed to fit the model (BF16≈55GB) onto your two 40GB A100 GPUs. The next step is speeding up that multi-GPU inference, which is indeed supported and is the core challenge you're facing.

The issue of slow inference, particularly for a large output_len like 1500 tokens, points to potential bottlenecks in how DeepSpeed is managing the multi-GPU setup.

Here are the key areas to investigate and different multi-GPU inference solutions you can use to significantly reduce your time cost:

1. Tensor Parallelism (Model Parallelism): DeepSpeed uses this (which is what allowed the model to fit on 2 GPUs) by sharding the model's layers across the GPUs. For your two-GPU setup, ensuring it's configured for optimal tensor parallelism (mp_size=2) is crucial.

Action: Ensure you're initializing the DeepSpeed engine for inference with the maximum tensor parallelism degree that your setup allows, which is mp_size=2. You can also look for the option replace_with_kernel_inject=True to use DeepSpeed's optimized CUDA kernels.

2. vLLM: This framework is known for its PagedAttention algorithm, which optimizes KV-cache usage. It has native support for Tensor Parallelism, which is what you need for Gemma 27B on 2 A100s.

Action: Try running your model with vLLM, setting the tensor parallelism size to 2:

vllm serve --model google/gemma-3-27b-it --tensor-parallel-size 2

3. Consider Quantization:
   Your model is 55GB in BF16. Even though you got it running, the memory is very tight. This might be indirectly impacting speed.

Quantization (e.g., GPTQ or AWQ to 4-bit): Quantizing the model's weights can drastically reduce the memory footprint.

Gemma 27B in 4-bit only requires about 21GB of VRAM (plus activation/KV cache memory).

Benefit: This frees up significant VRAM, which can be used for a larger KV-Cache, thus supporting your large output_len (1500 tokens) much more efficiently, or potentially allowing you to run multiple prompts concurrently (batching).

Trade-off: Quantization can sometimes slightly decrease model quality and might introduce a slight overhead due to dequantization, but modern quantized kernels (especially in vLLM) are highly optimized.

Thanks.