Low GPU usage during inference (~25%) on 4 A100 80 GB
Author: Kanishk-Kumar
Pull request: no
Labels: name: Core: Pipeline

name: Distributed Training / Models

name: bug

### System Info

I do not have access to the system at the moment.

### Who can help?

@ArthurZucker @Rocketknight1

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Hi, to reproduce, please run the following code:

```
import transformers
import torch
from secret import HF_TOKEN

model_id = "meta-llama/Meta-Llama-3.1-70B-Instruct"
tokenizer = transformers.AutoTokenizer.from_pretrained(model_id)
streamer = transformers.TextStreamer(tokenizer, skip_prompt=True)
pipeline = transformers.pipeline(
    "text-generation",
    model=model_id,
    model_kwargs={"torch_dtype": torch.bfloat16},
    device_map="auto",
    token = HF_TOKEN,
    streamer=streamer,
)

messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Write an essay on kites."},
]
outputs = pipeline(
    messages,
    max_new_tokens=5000,
    temperature=0.01,
    pad_token_id=tokenizer.eos_token_id
)
```

You'll see up to 100% GPU usage when model is loading, but after, each GPU will only have ~25% usage when model starts writing the output. I thought this is due to data getting across GPUs and bandwidth being the bottleneck, but then I ran the same code parallelly on two separate JuypterLab notebooks and GPU usage was ~50% during inference.

I do not have access to the system at the moment, but I'll get the access again in next few days and provide the `transformers-cli env` output and any other details if needed. Any help will be appreciated, thank you.

(Similar usage pattern was seen in 8xA100 GPU when running a larger model [meta-llama/Meta-Llama-3.1-405B-Instruct-FP8](https://huggingface.co/meta-llama/Meta-Llama-3.1-405B-Instruct-FP8).)

### Expected behavior

GPU usage should be >80%.

Comments:
I’m not an expert, so please correct me if I'm wrong. However, I believe the issue arises because you're only performing inference on one message, which doesn't fully utilize the GPU. To maximize GPU utility, I suggest using a larger dataset and setting `batch_size=128` (or experimenting with different batch sizes) in your `pipeline`.

Also, if your primary goal is inference, consider loading the model on each GPU. Depending on the precision you're using—and potentially even better with quantization—you might achieve better utilization by distributing the model across the GPUs.

>  I believe the issue arises because you're only performing inference on one message, which doesn't fully utilize the GPU.

But it should utilize, shouldn't it? Writing essay on kites (see code) takes a while.😅

> To maximize GPU utility, I suggest using a larger dataset and setting batch_size=128 (or experimenting with different batch sizes) in your pipeline.

What dataset? As you can see the code, I'm only using for inference. Not sure what `batch_size` does in case of inference, I thought it'll affect speed at the cost of accuracy and I don't want accuracy to be affected but I'll anyway try this, thanks!

> Also, if your primary goal is inference, consider loading the model on each GPU. Depending on the precision you're using—and potentially even better with quantization—you might achieve better utilization by distributing the model across the GPUs.

I'm using large models with high context length that require VRAM from multiple GPUs, I don't know what you mean by "distributed" but the model layers are getting offloaded across all the GPUs correctly. But during inference, each gets used only around ~25% (`watch -d -n 1 nvidia-smi`). I want full, either F16 or 32 bit precision (as you can also see in my code).

Hi @Kanishk-Kumar, I think the problem is that `device_map="auto"` is splitting the model across GPUs such that only one GPU is active at a time. @muellerzr is there a good solution in accelerate to get parallel sharding instead that works for generation?


> But it should utilize, shouldn't it? Writing essay on kites (see code) takes a while.😅

Longer token sequences, like generating 1,000 tokens, use more RAM. according to [this Blog](https://huggingface.co/blog/llama31), for the LLaMA 3.1 70B model, token generation memory usage looks something like this:

- 1k tokens → 0.313 GB
- 16k tokens → 4.88 GB
- 128k tokens → 39.06 GB

However, the GPU utilization stays relatively constant, which is why the GPU utilization doesn’t increase much, even with longer token sequences. So while generating a long essay may take time, it's still using a similar number of compute resources per token.

> 
> What dataset? As you can see the code, I'm only using for inference. Not sure what `batch_size` does in case of inference, I thought it'll affect speed at the cost of accuracy and I don't want accuracy to be affected but I'll anyway try this, thanks!
> 
I assumed you might be doing more than one inference, in which case batching would help. When running a single inference, like in your case, GPUs like the A100 are designed for parallel processing, and maximum utilization is usually achieved with large batches of data that can fully utilize all the GPU cores. If you’re only doing one inference at a time, the GPU won’t be fully utilized.

That said, if you’re looking to increase efficiency, you might want to explore optimizations with `ONNX `or `TensorRT`, which can speed up inference. However, depending on your use case, it's probably not worth the extra effort.

Regarding `Accelerate`, I found this quote [in the docs](https://huggingface.co/docs/accelerate/en/usage_guides/distributed_inference) that might relate to your situation:

> 2. Loading parts of a model onto each GPU and processing a single input at one time

 but I couldn’t find a specific demonstration for it in the docs.


Hi @Rocketknight1 , I saw constant ~25% usage on each GPU after model started streaming generated tokens, so I don't think only one GPU was active at a time.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.