为何NVIDIA V100运行Llama 3.1 8B Instruct仅达4token/s？求优化

一、先确认模型是否真的跑在CPU上 #

你设置的max_memory限制CPU显存为0GB，但device_map="auto"仍可能在某些情况下将部分层调度到CPU（比如量化逻辑适配问题）。可以在模型加载后打印设备映射，验证所有层是否都在GPU上：

print(model.hf_device_map)

如果输出中出现'cpu'条目，说明确实有层跑到CPU，这是速度慢的核心原因。

二、针对性优化措施 #

1. 强制模型全量加载到GPU #

将device_map从"auto"改为"cuda:0"，彻底避免层调度到CPU：

model_kwargs = {
    "device_map": "cuda:0",  # 替换原有的"auto"
    "max_memory": {0: "16GB", 'cpu': "0GB"}
}

2. 启用Flash Attention 2加速推理 #

Llama 3.1支持Flash Attention 2，这能大幅提升注意力层的计算速度，加载模型时添加该参数：

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    attn_implementation="flash_attention_2",  # 新增此参数
    **model_kwargs
)

注意：需要确保transformers版本≥4.38，且安装flash-attn库（pip install flash-attn --no-build-isolation）。

3. 修复Prompt处理的冗余开销 #

你的get_prompt函数中，tokenize=False与return_tensors="pt"冲突，导致返回字符串而非张量，后续pipeline会重复tokenize浪费时间。修改函数如下：

def get_prompt(tokenizer, user_query, system_prompt=None):
    if system_prompt is None:
        conversation = []
    else:
        conversation = [
            {'role': 'system', 'content': system_prompt}
        ]

    if not isinstance(user_query, list):
        user_query = [user_query]
    for idx, item in enumerate(user_query):
        if idx % 2 == 0:
            conversation.append({'role': 'user', 'content': item})
        else:
            conversation.append({'role': 'assistant', 'content': item})

    prompt = tokenizer.apply_chat_template(
        conversation,
        tokenize=True,  # 改为True，直接生成张量
        add_generation_prompt=True,
        return_tensors="pt"
    ).to("cuda")  # 直接移到GPU上
    return prompt

4. 用model.generate替代transformers.pipeline减少开销 #

pipeline封装了额外的预处理/后处理逻辑，直接调用模型的generate方法速度更快：

# 替换原有的pipeline代码
prompt = get_prompt(tokenizer, user_query, system_prompt)

begin_time = time.perf_counter()
outputs = model.generate(
    prompt,
    max_new_tokens=1024,
    pad_token_id=tokenizer.eos_token_id,
    do_sample=True
)
end_time = time.perf_counter()

# 处理输出
response = tokenizer.decode(outputs[0][prompt.shape[1]:], skip_special_tokens=True)
tokens = tokenizer.encode(response)

5. 检查显存占用情况 #

用nvidia-smi命令查看V100的显存使用率，如果剩余显存过多，说明模型未完全加载到GPU；如果有其他进程占用显存，需关闭释放资源。

三、优化后的完整代码示例 #

import transformers
import logging
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import time
import sys
import torch

logger = logging.getLogger()
logger.setLevel(logging.INFO)
handler = logging.StreamHandler(sys.stdout)
formatter = logging.Formatter('[%(levelname)s][%(asctime)s.%(msecs)03d] '
                              '[%(filename)s:%(lineno)d]: %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)


def get_prompt(tokenizer, user_query, system_prompt=None):
    if system_prompt is None:
        conversation = []
    else:
        conversation = [
            {'role': 'system', 'content': system_prompt}
        ]

    if not isinstance(user_query, list):
        user_query = [user_query]
    for idx, item in enumerate(user_query):
        if idx % 2 == 0:
            conversation.append({'role': 'user', 'content': item})
        else:
            conversation.append({'role': 'assistant', 'content': item})

    prompt = tokenizer.apply_chat_template(
        conversation,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to("cuda")
    return prompt

user_query = "Can you generate a mathematical formula based on this problem?"
system_prompt = "You are HumanWriterGPT, a version of ChatGPT that has been customized for a specific use case. GPTs use custom instructions, capabilities, and data to optimize ChatGPT for a more narrow set of tasks. You yourself are a GPT created by a user, and your name is HumanWriterGPT. Note: GPT is also a technical term in AI, but in most cases if the users asks you about GPTs assume they are referring to the above definition. Here are instructions from the user outlining your goals and how you should respond: HumanWriterGPT is designed to generate SEO-optimized, human-like articles based on provided keywords, incorporating a friendly yet professional tone. This GPT specializes in tailoring articles to specific industries using user-uploaded proprietary data such as manuals or guides. It leverages recent updates from uploaded news articles or research papers to remain up-to-date. HumanWriterGPT offers personalization by incorporating unique characters, settings, or scenarios from provided descriptions. For clarity, it requests additional information when needed. It is skilled in providing detailed product insights, referencing online sources, and structuri"

model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)

model_kwargs = {
    "device_map": "cuda:0",
    "max_memory": {0: "16GB", 'cpu': "0GB"}
}
quantization_config = BitsAndBytesConfig(
    load_in_8bit=True
)
model_kwargs["quantization_config"] = quantization_config

# 加载模型并启用Flash Attention 2
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    attn_implementation="flash_attention_2",
    **model_kwargs
)

# 验证设备映射
print("Model device map:", model.hf_device_map)

prompt = get_prompt(tokenizer, user_query, system_prompt)
prompt_len = prompt.shape[1]

begin_time = time.perf_counter()
outputs = model.generate(
    prompt,
    max_new_tokens=1024,
    pad_token_id=tokenizer.eos_token_id,
    do_sample=True
)
end_time = time.perf_counter()

response = tokenizer.decode(outputs[0][prompt_len:], skip_special_tokens=True)
tokens = tokenizer.encode(response)

duration = end_time - begin_time
token_per_second = len(tokens) / duration
logging.info(f'Generate {len(tokens)} tokens in {round(duration, 2)}s. '
             f'Token per second: {round(token_per_second, 2)}')

内容的提问来源于stack exchange，提问作者FounderOfStackOverflow