[Question]: MInference Pre filling is slower than the vllm original version

[junior-zsy]

Describe the issue

code :

# Copyright (c) 2024 Microsoft
# Licensed under The MIT License [see LICENSE for details]

from vllm import LLM, SamplingParams

from minference import MInference
import time

def read_content_from_file(file_path, num_chars=5000):
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            content = file.read(num_chars)
        return content
    except FileNotFoundError:
        logging.error(f"File {file_path} not found.")
        return ""
    except Exception as e:
        logging.error(f"An error occurred while reading the file: {e}")
        return ""

content = read_content_from_file("./question.txt", 12000) + ",请总结上面的故事梗概。"

prompts = []
for _ in range(50):
    prompts.extend([content])

sampling_params = SamplingParams(
    temperature=0.8,
    top_p=0.95,
    max_tokens=1,
)
model_name = "/xxx/model/Qwen2-7B-Instruct"
llm = LLM(
    model_name,
    max_num_seqs=1,
    enforce_eager=True,
    tensor_parallel_size=1,
    max_model_len=128000,
)



start_time = time.time()
outputs = llm.generate(prompts, sampling_params)
end_time = time.time()

elapsed_time = end_time - start_time
print(f"vllm Generating text took {elapsed_time:.2f} seconds.")



# Patch MInference Module
minference_patch = MInference("vllm", model_name)
llm = minference_patch(llm)

start_time = time.time()
outputs = llm.generate(prompts, sampling_params)
end_time = time.time()

elapsed_time = end_time - start_time
print(f"minference Generating text took {elapsed_time:.2f} seconds.")

results of execution：

Processed prompts: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:41<00:00, 1.22it/s]
vllm Generating text took 41.57 seconds.
Patched model for minference with vLLM..
Processed prompts: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [01:34<00:00, 1.90s/it]
minference Generating text took 95.37 seconds.

why minference slower than vllm

[iofu728]

Hi @junior-zsy, thank you for your feedback.

Although some information is missing, based on your code, I guess you ran a batch size of 50 with prompt tokens length less than 12k.

MInference involves some additional approximations which introduce overhead, making it slower than dense attention for shorter context sizes. The latency performance breaks even at around 30k tokens, as shown in our end-to-end benchmark: End-to-End Benchmark. The latency can be further optimized by adjusting the sparsity rate of sparse attention, though more aggressive sparse attention in smaller context windows hasn't been extensively tested.

We tested 128k with Qwen2-7B using hf, and its speedup ratio matches the results with LLaMA-3-8B in triton==2.1.0.

➜  MInference git:(main) ✗ python3.9 experiments/benchmarks/benchmark_e2e.py --attn_type minference --context_window 100000 --model_name Qwen/Qwen2-7B-Instruct
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
Loading checkpoint shards: 100%|████████████████████████████████████| 4/4 [00:02<00:00,  1.44it/s]
Patched model for minference..
100000 13.43383584022522
➜  MInference git:(main) ✗ python3.9 experiments/benchmarks/benchmark_e2e.py --attn_type minference_with_dense --context_window 100000 --model_name Qwen/Qwen2-7B-Instruct
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
Loading checkpoint shards: 100%|████████████████████████████████████| 4/4 [00:02<00:00,  1.45it/s]
Patched model for minference..
100000 19.306786155700684

[SimJeg]

Hello,

I was able to successfully run the benchmark for the model gradientai/Llama-3-8B-Instruct-Gradient-1048k using --context_window 100_000on a A100 80GB. I obtained the following results:

• minference_with_dense : 24.0sec
• minference : 15.8sec

I noticed in L133 that the KV cache is not stored use_cache=False, so I wanted to move beyond the time to first token benchmark.

I tried to run a transformers pipeline as explained in the README. I used the following tokens similarly to benchmark_e2e.py:

from transformers import AutoTokenizer

model_name = "gradientai/Llama-3-8B-Instruct-Gradient-1048k"
prompt_size = 100_000
tokenizer = AutoTokenizer.from_pretrained(model_name)
prompt = open("./prompt_hardest.txt").read()
tokens = tokenizer(prompt)["input_ids"] 
tokens = tokens * (prompt_size // len(tokens) + 1)
prompt = tokenizer.decode(tokens[:prompt_size - 1])

When I tried to run the pipeline I get OOM

from transformers import pipeline

pipe = pipeline("text-generation", model=model_name, torch_dtype="auto", device="cuda:0")
pipe(prompt, max_new_tokens=50, do_sample=False)

So I tried with vLLM and it worked

from time import time
from vllm import LLM, SamplingParams
from minference import MInference

def run_llm(llm, prompt):
    start = time()
    results = llm.generate([prompt], SamplingParams(temperature=0.0))
    print("Answer: ", results[0].outputs[0].text)
    print(f"Time: {time() - start:.2f}s") 
    

llm = LLM(model_name, max_num_seqs=1, enforce_eager=True, max_model_len=200_100, device="cuda:0")
run_llm(llm, prompt)
minference_patch = MInference("vllm", model_name)
llm = minference_patch(llm)
run_llm(llm, prompt)

Good news is: the generated text is the same. Bad news is: MInference is slower (21.5sec vs 23.2sec). Good news again is that using prompt_size=200_000, I obtained 75sec vs 38sec so MInference was 2x faster.

Do you know why the threshold for MInference to be useful is higher when doing generation ? Also I think the statement "a single A100 can support up to 1.8M context prompts" is a bit exaggerated because it does not include KV cache.

[iofu728]

Hi @SimJeg,

Thanks for your interest in MInference and the detailed results.

1. The results and speedup at 100K are expected and make sense. You might achieve even better performance by using triton==2.1.0.
2. Since MInference 1.0 focuses only on the pre-filling stage, we set use_cache=False. The storage time for KV cache can be fully covered by prefetching, so the actual TTFT and the current benchmark are almost the same. In fact, if you set use_cache=True and run the benchmark again, you will get the same latency. Additionally, enabling use_cache=True for context sizes greater than 200K can cause OOM with the HF implementation.
3. The original HF implementation experiences OOM at 30K. You need to use the minference_with_dense mode.
4. Due to PageAttention, vllm results in higher TTFT compared to the HF implementation in our tests. Overall, I don't think the increase in latency due to KV cache storage is reasonable; it is caused by the implementation, as this overhead can be effectively hidden by computational overhead.
5. In our README, we state, "Based on our tests, a single A100 can support up to 1.8M context prompts during the pre-filling stage using LLaMA-3-8B-4M with bf16." We emphasized "during the pre-filling stage." Even considering the KV cache, this is feasible, and you can use --kv_cache_cpu or other KV cache management measures to solve OOM issues. Additionally, our 1.8M context window estimate is conservative; the upper limit for a single A100 context window is actually higher.

We will optimize this part of the vllm implementation in the future. If you have any further questions, please feel free to let me know.

[SimJeg]

Hello @iofu728, thanks for your detailed answer, I didn't now about the --kv_cache_cpu option. I still don't really understand why we don't see improvements in vLLM generation at 100k while the benchmark for HF tells us it's supposed to be faster. I guess as you mentioned it's a matter of vllm implementation. Better TTFT is great, but it shouldn't degrade the overall speed of generation. As MInference 1.0 focuses on prefilling, shouldn't the time for generation unchanged ?

[iofu728]

Hi @SimJeg, thanks for your support.

I tried to reproduce the vLLM latency results and built an end-to-end benchmark in vLLM. However, in my tests, the 100k context size still shows over a 1.5x speedup, and the 128k context size shows around a 1.7x speedup. You can refer to the code below or check out end-to-end-benchmark-using-vllm.

vLLM's FlashAttention is relatively faster than HF, primarily because our kernel is based on Triton, which performs slightly worse than FlashAttention libraries, with a performance loss of about 10-20%. Additionally, we implemented head pipeline parallelism, which also incurs some performance loss. However, these aspects can be optimized.

Overall, based on our tests, I believe MInference have not changed "the time for generation." and can effectively improve TTFT in both vllm and HF frameworks.

# Copyright (c) 2024 Microsoft
# Licensed under The MIT License [see LICENSE for details]

import time

import torch
from transformers import AutoTokenizer
from vllm import LLM, SamplingParams

from minference import MInference

m = 100_000
model_name = "gradientai/Llama-3-8B-Instruct-262k"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
prompt_complex = open("./prompt_hardest.txt").read()
input_ids = tokenizer(prompt_complex)["input_ids"]
n = len(input_ids)
b = m // n + 1

new_input_ids = (input_ids * b)[:m]
prompt = tokenizer.decode(new_input_ids)
prompts = [prompt]

sampling_params = SamplingParams(
    temperature=0.8,
    top_p=0.95,
    max_tokens=1,
)

llm = LLM(
    model_name,
    max_num_seqs=1,
    enforce_eager=True,
    max_model_len=128000,
)

# Patch MInference Module
minference_patch = MInference("vllm", model_name)
llm = minference_patch(llm)

torch.cuda.synchronize()
start = time.time()
with torch.no_grad():
    outputs = llm.generate(prompts, sampling_params)
torch.cuda.synchronize()
s = time.time() - start

print(m, s)

[SimJeg]

Thanks for doing these benchmarks, they illustrate the efficiency of Minference. I will try to reproduce it next week using triton==2.1.0 (I was using triton==2.3.0)