Llama2-7b模型文本生成耗时过长，求优化提速方案

问题描述 #

我已在本地加载HuggingFace的Llama2-7b模型，模型加载仅需约2分钟，但针对简单prompt（如“What is an LLM?”）生成响应耗时超2小时。以下是我的实现代码：

class LlamaInference():
    """基于Llama 70b参数LLM的文本生成类，用于生成数据集。"""
    def __init__(self, output_file_path):
        """ 初始化LlamaInference类及所有变量
        参数说明:
        output_file_path -- 数据集生成后输出文件的存储路径
        """
        self.output_file_path = output_file_path
        self.pipeline = None
        self.output_list = []
        self.output_list_condensed = []

    def __set_up(self):
        """ 使用指定模型及版本配置HuggingFace pipeline
        """
        #model = "meta-llama/Llama-2-70b-chat-hf"
        #revision = "e6152b720bd3cd67afc66e36d06893a0e1f84b48"
        model = "meta-llama/Llama-2-7b-chat-hf"
        revision = "08751db2aca9bf2f7f80d2e516117a53d7450235"

        self.tokenizer = AutoTokenizer.from_pretrained(model, padding_side="left")

        self.pipeline = transformers.pipeline(
            "text-generation",
            model=model,
            tokenizer=self.tokenizer,
            torch_dtype=torch.float16,
            device_map="auto",
            revision=revision,
            do_sample=True,
            return_full_text=False
        )
        self.pipeline.tokenizer.pad_token_id = self.tokenizer.eos_token_id

    def gen_text(self, prompt, **kwargs):
        """ 根据prompt生成文本

        参数说明:
        prompt -- 向LLM提出的问题
        **kwargs -- 传递给self.pipeline的参数
        """
        try:
            if self.pipeline is None:
                self.__set_up()
            if "batch_size" not in kwargs:
                kwargs["batch_size"] = 1
            
            if "max_new_tokens" not in kwargs:
                kwargs["max_new_tokens"] = 2048

            kwargs.update(
                {
                    "pad_token_id": self.pipeline.tokenizer.eos_token_id, 
                    "eos_token_id": self.pipeline.tokenizer.eos_token_id,
                }
            )
            display('started pipeline')
            token_outputs = self.tokenizer(prompt)
            display(token_outputs)
            outputs = self.pipeline(prompt, **kwargs)
            display('end pipeline')
            return outputs
        except Exception as error:
            display(f'__gen_text error: {error}')

inference = LlamaInference('test.json')
result = inference.gen_text("Hello", max_new_tokens=2048, batch_size=1, temperature=0.5)

提速优化建议 #

1. 确保硬件加速生效 #

• 验证GPU调用：添加print(torch.cuda.is_available())确认GPU可用，检查device_map="auto"是否将模型分配到GPU。如果模型运行在CPU上，7B模型生成速度必然极慢，必须优先解决GPU调用问题。
• 启用Flash Attention：加载模型时添加attn_implementation="flash_attention_2"参数，需先安装flash-attn库，这能大幅降低注意力机制的计算耗时。

2. 调整生成参数 #

• 减少max_new_tokens：简单prompt无需生成2048个token，可改为256或512，直接缩短生成流程的计算量。
• 优化采样策略：如果不需要生成多样化内容，将do_sample=False改用贪婪解码；若保留采样，可适当调高temperature或限制top_k/top_p值，减少不必要的候选token计算。

3. 模型加载优化 #

• 使用量化模型：通过bitsandbytes库加载8-bit或4-bit量化版本，在pipeline中添加load_in_8bit=True或load_in_4bit=True，既减少显存占用，又能显著提升生成速度。
• 确认浮点精度：确保torch_dtype=torch.float16生效，避免模型在CPU与GPU间频繁数据传输。

4. 代码细节优化 #

• 移除调试代码：删除display('started pipeline')、display(token_outputs)等调试用输出，减少IO耗时。
• 简化tokenizer配置：在初始化tokenizer时直接设置self.tokenizer.pad_token = self.tokenizer.eos_token，避免重复配置pad_token_id。

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