model.generate single CPU core bottleneck

[dhcracchiolo]

System Info

I'm running inference on a GPU EC2 instance using CUDA. After doing a little profiling I noticed the model.generate method was the clear bottleneck. Upon closer inspection running htop showed that during this method call only a single cpu core is used and is maxed out to 100%. I've made sure sure all of my weights, biases and activations are all on the gpu. Running nvidia-smi shows the proper amount of VRAM usage. My question is, is there a way to speed this method up using multiple CPU cores? I haven't dug deep into the method to see exactly what is causing the issue yet. Just curious if there's something obvious I'm missing.

(Basic sample code)

from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig
import torch

tokenizer = LlamaTokenizer.from_pretrained("decapoda-research/llama-7b-hf")

model = LlamaForCausalLM.from_pretrained(
    "decapoda-research/llama-7b-hf",
    load_in_8bit=True,
    device_map="auto",
)

PROMPT = f"""### Below is an instruction that describes a task. Write a response that appropriately completes the request.
### Instruction:
### Input: "What is the difference between a llama and a vicuna?"
### Response:"""

inputs = tokenizer(
    PROMPT,
    return_tensors="pt",
)
input_ids = inputs["input_ids"].cuda()

generation_config = GenerationConfig(
    temperature=0.6,
    top_p=0.95,
    repetition_penalty=1.15,
)

  generation_output = model.generate(
      input_ids=input_ids,
      generation_config=generation_config,
      return_dict_in_generate=True,
      output_scores=True,
      max_new_tokens=128,
  )

for s in generation_output.sequences:
    print(tokenizer.decode(s))

Who can help?

No response

Information

• The official example scripts
• 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

Run this method with any model and profile the cpu cores to see it's only using a single core:
generation_output = model.generate(
input_ids=input_ids,
generation_config=generation_config,
return_dict_in_generate=True,
output_scores=True,
max_new_tokens=128,
)

Expected behavior

Multiprocessing to help balance the load.

[sgugger]

cc @gante

[gante]

Hey @dhcracchiolo 👋

Good timing on your question, we are precisely exploring what we can do to speed up LLaMA within python + PyTorch + .generate API, in this repo.

In a nutshell, the model is CPU bound, as is most PyTorch code for models of this size. If you install the repo above, run python run_llama.py --model huggingface/llama-7b --preallocate --profile, and check the profiler results on TensorBoard, you'll see that the CPU bottleneck is in orchestrating the PyTorch execution on GPU. This is why you see high-speed inference packages to go straight into CUDA, to get rid of this overhead.

Feel free to use the code from the repo above, it should get you ~50% speedup :) It doesn't have all the options that we have in .generate, though. If you have concrete suggestions on how to further speed it up, we're all ears 👂

[dhcracchiolo]

Hey @dhcracchiolo 👋

Good timing on your question, we are precisely exploring what we can do to speed up LLaMA within python + PyTorch + .generate API, in this repo.

In a nutshell, the model is CPU bound, as is most PyTorch code for models of this size. If you install the repo above, run python run_llama.py --model huggingface/llama-7b --preallocate --profile, and check the profiler results on TensorBoard, you'll see that the CPU bottleneck is in orchestrating the PyTorch execution on GPU. This is why you see high-speed inference packages to go straight into CUDA, to get rid of this overhead.

Feel free to use the code from the repo above, it should get you ~50% speedup :) It doesn't have all the options that we have in .generate, though. If you have concrete suggestions on how to further speed it up, we're all ears 👂

Thanks for the response! Excellent! I'll check the repo out and report back. 👍

[github-actions]

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 are likely to be ignored.

[calvintwr]

@gante Can i check if using other interfaces like pipeline circumvent this issue?

[gante]

@calvintwr ofc you can, no guarantees it would be merged though ;) Depends on the benefits-costs tradeoff, as any contribution does.

pipeline calls generate under the hood, so the bottleneck would be the same.

[calvintwr]

@gante I don’t mind exploring the solution. Can I ask if the code in question here?:

transformers/src/transformers/generation/utils.py

Line 1295 in 1791ef8

def generate(

If so, I skimmed, and is the reason because of extensive use of torch methods at many different places? Or is there a main culprit?

[gante]

The main culprit is the CPU bottleneck when issuing instructions (i.e. the time in Python issuing GPU instructions is not optimized in many segments of the model forward pass). Check this comment, it has instructions on how to dive deeper :) Note that most of the bottleneck is on the forward pass of the model itself.

This means that speedups for existing models can come mostly from two angles:
1 - Spending less time on CPU preparing GPU instructions
2 - Picking faster GPU operations that yield the same results

We rarely optimize our models, so there is room for some speed-up in our models (up to 10%, I'd say).

[calvintwr]

@gante I see that. 10% is not a lot…

What’s your view on offloading this portion to Cpp, and use SharedArray to transfer data between?

We can help, but just trying to get a sense of the extent of the task. But if this is achievable, this will bring transformer to the next level.

[gante]

@calvintwr offloading to other programming languages is not within the scope of transformers, at least not for now :) SharedArray could be an option!

In transformers, we love performance upgrades! However, our biggest focus is on ease of use, so we don't want to introduce performance gains at the cost of complex dependencies or complex code 🤗