vllm generate推理与Huggingface generate推理对齐(长样本)

自学AI的鲨鱼儿

已于 2024-08-02 22:08:22 修改

阅读量8.9k

点赞数 34

分类专栏： # 训练文章标签：个人笔记深度学习 LLM

于 2024-06-24 15:30:23 首次发布

本文链接：https://blog.youkuaiyun.com/qq_16555103/article/details/139927787

版权

一、离线推理

1.1、开发环境

1.2、推理对齐

1.2.1、样本生成参数gen_kwargs对齐

1.2.2、文本生成eos token对齐

1.2.3、model 初始化数据类型对齐

1.2.4、vllm generate推理选择适当的batch size

1.5.2、Huggingface generate

二、在线部署

2.1、部署相关的基础知识

2.2、类openai vLLM 服务器部署

2.2.1 类openai vLLM 服务器参数解析

2.2.2 使用py脚本来实现 python -m vllm.entrypoints.openai.api_server 功能进行部署，只是可以同py脚本添加更多自定义的功能，代码如下

2.2.3 vLLM python -m vllm.entrypoints.openai.api_server 原本不支持自定义的序列惩罚，故修改vllm.entrypoints.openai.api_server 源码以支持推理序列抑制惩罚

2.2.4 client 请求

一、离线推理

本文中所有的结论皆是个人通过下列开发环境中评测出来的，仅作为参考

1.1、开发环境

个人数据长度：prompt普遍在9k以上

linux: ubuntu
GPU：单张A6000（48G）
python=3.10
torch==2.3.0
transformers==4.41.2
vllm==0.5.0.post1
vllm-flash-attn==2.5.9
flash-attn==2.5.9.post1

实现目标：
vllm generate函数批量输入prompt（普遍在9k以上）贪婪解码推理结果与 Huggingface generate函数 贪婪解码推理完全一致

1.2、推理对齐

若想vllm generate推理与Huggingface generate推理完全一致主要需要做以下几个方面修改

1.2.1、样本生成参数gen_kwargs对齐

首先要确保vllm生成参数与Huggingface生成参数对齐，但两者某些参数命名不完全相同，需要找出各个重要参数的映射关系

VLLM(generate 函数)	Huggingface(generate 函数)	说明
max_tokens	max_new_tokens	允许生成新token最大数值
top_p	top_p
top_k（默认为50）	top_k（默认为-1，全词库）
temperature	temperature	温度系数，Huggingface 不允许temperature=0，而vllm中temperature=0 表示贪婪解码
repetition_penalty	repetition_penalty	重复惩罚
/	do_sample	huggingface 中do_sample=True 表示做概率采样，do_sample=False 表示做贪婪解码，而vllm中没有此参数，使用temperature=0 就可以贪婪解码

根据上表，我们首先需要将两个框架的生成参数对齐，为了保证实验的稳定性，两者均采用贪婪解码，需要注意的是，即使使用贪婪解码，top-p的值也会影响生成的结果，故两者均设置参数如下

# 双方启用贪婪解码, 确保推理的一致性
gen_kwargs_vllm = {
    "max_tokens": 1150,
    "top_p": 0.9,
    "top_k": 50,
    "temperature": 0.0,
    "repetition_penalty": 1.0,
}
gen_kwargs_hug = {
    "max_new_tokens": 1150,
    "top_p": 0.9,
    "top_k": 50,
    "temperature": 0.35,
    "repetition_penalty": 1.0,
    "do_sample": False
}

1.2.2、文本生成eos token对齐

Huggingface generate 函数设置文本结束符号，这里面的llama模型选中的llama3-instruct，具体需要根据自己的模型进行更换

# 组织生成文本的超参
if model_type == 'llama':
    # 设置文本生成的结束 token ID
    gen_kwargs_hug['eos_token_id'] = [self.tokenizer.eos_token_id, self.tokenizer.convert_tokens_to_ids("<|end_of_text|>")]
    gen_kwargs_hug['pad_token_id'] = self.tokenizer.eos_token_id
elif model_type == 'qwen2':
    # 设置文本生成的结束 token ID
    gen_kwargs_hug['eos_token_id'] = [151645,151643]
    gen_kwargs_hug['pad_token_id'] = 151643
else:
    raise ValueError(f"Only support 'llama or qwen2' now, but got '{model_type}'")

vllm generate 函数设置文本结束符号，这里面的llama模型选中的llama3-instruct，具体需要根据自己的模型进行更换

# 组织生成文本的超参
if self.model_type == 'llama':
    # 设置文本生成的结束 token ID,  # gen_kwargs_vllm['pad_token_id'] = self.tokenizer.eos_token_id
    gen_kwargs_vllm['stop_token_ids'] = [self.tokenizer.eos_token_id, self.tokenizer.convert_tokens_to_ids("<|end_of_text|>")]
elif self.model_type == 'qwen2':
    # 设置文本生成的结束 token ID, # gen_kwargs_vllm['pad_token_id'] = 151643
    gen_kwargs_vllm['stop_token_ids'] = [151645,151643]
else:
    raise ValueError(f"Only support llama and qwen2, model_type {self.model_type} is not supported")

1.2.3、model 初始化数据类型对齐

Huggingface from_pretrained 模型的数据类型也要与VLLM 模型初始化时数据类型保持一致

Huggingface 加载模型

# Huggingface
model_ = AutoModelForCausalLM.from_pretrained(
            model_name_or_path,
            trust_remote_code=True,
            low_cpu_mem_usage=True,     
            torch_dtype = torch.float16,   # float16 加载, 后面的vllm 也要使用 float16 加载
).eval()

VLLM 加载模型

一些其他的参数（待确定）：

max_num_batched_tokens=4096, # 控制批处理中的最大token数

max_num_seqs=256, # 控制批处理中的最大序列数

class VLLMInference():
    def __init__(self,
                 model_name_or_path:str,
                 model_type:str,
                 # dtype 模型加载的数据类型, 'auto' 表示自动, torch.float32 表示 fp32, torch.float16 表示 fp16
                 # 与 Huggingface from_pretrained dtype 尽量保持一致, 为了Huggingface generate结果对齐
                 dtype: str,  
                 seed: int=0, # VLLM 默认为0, 默认即可
                 trust_remote_code: bool = True, 
                 tensor_parallel_size: int = 1,   # GPU 张量并行的卡数
                 gpu_memory_utilization: float = 0.9, # GPU 内存占用率
                 max_seq_len_to_capture: int = 9800, # 提升效率的cuda, 可以选择样本中中位数适中的文本长度
                **kwargs
                ):
        
        self.SYSTEM_PROMPT = SYSTEM_PROMPT
        self.model_name_or_path = model_name_or_path
        self.model_name_suffix = os.path.split(model_name_or_path)[-1]
        assert model_type in ['llama', 'qwen2'], f"model_type must be in ['llama', 'qwen2'], but got {model_type}"
        self.model_type = model_type
        
        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path,
                                                    trust_remote_code=True,
                                                    use_fast=True
                                                    # use_fast=False if model.config.model_type == 'llama' else True
                                                )
        # VLLM 模型初始化
        # vllm class 详情见: https://docs.vllm.ai/en/latest/dev/offline_inference/llm.html
        self.model = LLM(model=model_name_or_path,
                         trust_remote_code=trust_remote_code,
                         tensor_parallel_size=tensor_parallel_size,
                         dtype=dtype,  # 与 Huggingface from_pretrained dtype 尽量保持一致, 为了Huggingface generate结果对齐
                         gpu_memory_utilization=gpu_memory_utilization,
                         max_seq_len_to_capture=max_seq_len_to_capture,
                         **kwargs
                        )
        self.model.set_tokenizer(tokenizer)
        self.tokenizer = self.model.get_tokenizer()

1.2.4、vllm generate推理选择适当的batch size

经个人实验多次证明，虽然vllm generate函数可以一次性输入所有的数据，但如果一次性输入的数量过多，也会导致vllm推理的结果与Huggingface 推理结果存在差异

个人数据集中的prompt多数在9k左右，属于较长的样本，在48G的A6000多次实验batch size <= 15，可以贪婪解码得到与Huggingface generate 贪婪解码完全一致的结果

1.3、推理代码

VLLM 代码，需要注意的是，如果vllm使用微调PI扩展长度后的model推理，vllm 会自动根据model config.rope_scaling 的缩放值调整vllm 的max_seq_len 这个参数，所以不用担心vllm截断超过原始模型输入的长度，比如本人数据prompt长度均在9k，已经超出了llama3 8k的输入长度，但通过PI扩展长度微调后，config.rope_scaling=2，而vllm加载model后max_seq_len=16384 而不是 8192

class VLLMInference():
    def __init__(self,
                 model_name_or_path:str,
                 model_type:str,
                 # dtype 模型加载的数据类型, 'auto' 表示自动, torch.float32 表示 fp32, torch.float16 表示 fp16
                 # 与 Huggingface from_pretrained dtype 尽量保持一致, 为了Huggingface generate结果对齐
                 dtype: str,  
                 seed: int=0, # VLLM 默认为0, 默认即可
                 trust_remote_code: bool = True, 
                 tensor_parallel_size: int = 1,   # GPU 张量并行的卡数
                 gpu_memory_utilization: float = 0.9, # GPU 内存占用率
                 max_seq_len_to_capture: int = 9800, # 提升效率的cuda, 可以选择样本中中位数适中的文本长度
                **kwargs
                ):
        
        self.SYSTEM_PROMPT = SYSTEM_PROMPT
        self.model_name_or_path = model_name_or_path
        self.model_name_suffix = os.path.split(model_name_or_path)[-1]
        assert model_type in ['llama', 'qwen2'], f"model_type must be in ['llama', 'qwen2'], but got {model_type}"
        self.model_type = model_type
    
        # vllm class 详情见: https://docs.vllm.ai/en/latest/dev/offline_inference/llm.html
        # 需要注意的是，如果vllm使用微调PI扩展长度后的model推理，vllm 会自动根据model config.rope_scaling 的缩放值调整vllm 的max_seq_len 这个参数，
        # 所以不用担心vllm截断超过原始模型输入的长度，比如本人数据prompt长度均在9k，已经超出了llama3 8k的输入长度，
        # 但通过PI扩展长度微调后，config.rope_scaling=2，而vllm加载model后max_seq_len=16384 而不是 8192
        self.model = LLM(model=model_name_or_path,
                         trust_remote_code=trust_remote_code,
                         tensor_parallel_size=tensor_parallel_size,
                         dtype=dtype,
                         gpu_memory_utilization=gpu_memory_utilization,
                         max_seq_len_to_capture=max_seq_len_to_capture,
                         **kwargs
                        )
        
        # tokenizer = AutoTokenizer.from_pretrained(model_name_or_path,
        #                                             trust_remote_code=True,
        #                                             padding_side="left",   # 推理侧需要左pad
        #                                             use_fast=True
        #                                             # # llama不支持fast
        #                                             # use_fast=False if model.config.model_type == 'llama' else True
        #                                         )
        # self.model.set_tokenizer(tokenizer)
        self.tokenizer = self.model.get_tokenizer()
        logger.info(f"vllm tokenizer: \n{self.tokenizer}")
        

    def _generate(self, 
                 text:Union[str,List[str]],
                 gen_kwargs:dict,
                 batch_size:int=10,  # vllm 处理的 batch size, 如果不传入合适的batch size 会导致与Huggingface generate 不能对齐 (可能是传入大量的数据内存优化导致的)
                ):
        """
        解释一下 vLLM 在处理大量输入数据时的行为：
        输入处理：
            当你将所有数据作为一个列表传入 model.generate() 方法时，vLLM 确实会接收所有这些输入。但是，它不会简单地将整个输入列表的大小作为单一的 batch size 来处理。
        动态批处理：
            vLLM 使用动态批处理（dynamic batching）技术。这意味着它不会一次性处理所有输入，而是根据当前的计算资源和模型容量，动态地决定如何最efficiently地处理这些输入。
        内部批处理机制：
            vLLM 会内部管理一个请求队列。
            它会根据当前可用的 GPU 内存和计算资源，动态地决定每次处理多少输入。
            这个动态决定的批大小通常小于你提供的全部输入数量，特别是当输入数量很大时。
        最大批大小限制：
            虽然你可能输入了成千上万的提示，但 vLLM 有内部的最大批大小限制。
            这个限制是为了确保efficient的内存使用和计算效率。
            默认情况下，这个限制通常远小于你可能提供的全部数据量。
        连续处理：
            vLLM 会连续地处理输入队列，每次取出一定数量的输入进行处理。
            这个过程会持续进行，直到所有输入都被处理完毕。
        性能优化：
            这种方法允许 vLLM 在处理大量数据时保持高效率。
            它可以充分利用 GPU 资源，同时避免因为一次性加载过多数据而导致的内存问题。
        用户角度：
            从用户的角度来看，当你调用 model.generate(all_prompts, sampling_params) 时，vLLM 会处理所有输入，但内部是分批进行的。
            你会得到所有输入的输出，就好像它们是一次性处理的一样。
        控制批大小：
            如果你确实需要更精细地控制批处理大小，可以考虑使用 AsyncLLMEngine 或者在服务器模式下设置 max_batch_size 参数。
            但在大多数情况下，让 vLLM 自动管理批处理会得到更好的性能。
        """
        gen_kwargs = copy.deepcopy(gen_kwargs)
        # 组织生成文本的超参
        if self.model_type == 'llama':
            # 设置文本生成的结束 token ID,  # gen_kwargs['pad_token_id'] = self.tokenizer.eos_token_id
            gen_kwargs['stop_token_ids'] = [self.tokenizer.eos_token_id, self.tokenizer.convert_tokens_to_ids("<|end_of_text|>")]
        elif self.model_type == 'qwen2':
            # 设置文本生成的结束 token ID, # gen_kwargs['pad_token_id'] = 151643
            gen_kwargs['stop_token_ids'] = [151645,151643]
        else:
            raise ValueError(f"Only support llama and qwen2, model_type {self.model_type} is not supported")
        
        # SamplingParams 参数详情见: https://docs.vllm.ai/en/latest/dev/sampling_params.html#vllm.SamplingParams
        sampling_params = SamplingParams(**gen_kwargs)
        logger.warning(f"Now vllm running sampling_params \n{sampling_params}")

        # 组织文本
        if isinstance(text, str):
            text = [text]
        text = [i.strip() for i in text]

        text = [[
                {"role": "system", "content": self.SYSTEM_PROMPT},
                {"role": "user", "content": i}
            ] for i in text]
        
        text = [self.tokenizer.apply_chat_template(i, 
                                                    add_generation_prompt=True, 
                                                    tokenize=False   # False 表示返回 str,不转化为id
                                                ) for i in text]
        # 这里选择转化为id传给prompt_token_ids, 因为vllm默认的tokenizer encode 可能会加入前缀特殊token
        text = [self.tokenizer.encode(i, add_special_tokens=False) for i in text]
        logger.info(f"vllm one input:\n{self.tokenizer.decode(text[0])}")
        batch_num = math.ceil(len(text)/batch_size)
        logger.info(f"VLLM batch size: {batch_size}, length of datas num: {len(text)}, batch_num: {batch_num}")

        # 每次处理一个batch size
        outputs = []
        for b_idx,i in enumerate(range(0,len(text),batch_size)):
            s = i
            e = i + batch_size
            if e >= len(text):
                e = len(text)
            batch_ids = text[s: e]
            if (b_idx + 1) % 10 == 0:
                logger.info(f"batch id/batch num: {b_idx+1}/{batch_num}")
            batch_outputs = self.model.generate(prompt_token_ids = batch_ids, 
                                                sampling_params=sampling_params) 
            outputs = outputs + batch_outputs
        return outputs

Huggingface face 代码（其中 ModelUtils.load_model 是自定义的类，就是普通的float16加载模型）

class HuggingFaceInference():
    def __init__(self, 
            model_name_or_path:str, # base model 模型的路径或名称
            model_max_length:int=16384, # 模型允许的最大长度, 默认为 16384(即16k), 原模型长度不够的使用 PI 插值
            adapter_name_or_path:str=None,     # lora 适配器模型的路径或名称, None 表示不使用
            use_cache:bool=True,    # 推理时是否使用模型的缓存
            load_in_4bit:bool=False, # 是否使用 bnb 4bit进行推理，能够节省很多显存，但效果可能会有一定的下降
        ):
        self.model_name_or_path = model_name_or_path
        self.adapter_name_or_path = adapter_name_or_path
        self.model_max_length = model_max_length
        # 定义提示模版
        self.SYSTEM_PROMPT = SYSTEM_PROMPT
        logger.info(f"model_name_or_path: {model_name_or_path}")
        logger.info(f"adapter_name_or_path: {adapter_name_or_path}")
        logger.info(f"SYSTEM_PROMPT: {self.SYSTEM_PROMPT}")

        # 定义base model 和 lora 适配器模型的后缀名称
        self.model_name_suffix = os.path.split(model_name_or_path)[-1]
        if self.adapter_name_or_path is not None:
            self.adapter_name_suffix = os.path.split(adapter_name_or_path)[-1]
        else:
            self.adapter_name_suffix = 'no_adapter'
        
        logger.info(f"Loading model: {model_name_or_path} and adapter: {adapter_name_or_path}")
        self.config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=True)
        self.config.use_cache = use_cache  
        self.model_type = self.config.model_type
        # NOTE 从预训练模型配置中获取原始上下文长度, 
        # 如果设置的上下文窗口大小超过了原始长度,则需要计算 RoPE 缩放因子
        if self.config.model_type == 'llama':
            # 使用PI插值, 是否扩展RoPE的position
            orig_ctx_len = getattr(self.config, "max_position_embeddings", None)
            if orig_ctx_len and model_max_length > orig_ctx_len:
                scaling_factor = float(math.ceil(model_max_length / orig_ctx_len))
                self.config.rope_scaling = {"type": "linear", "factor": scaling_factor}
                logger.success(f'Use PI/NTK change {self.config.model_type} model_max_length from {orig_ctx_len} to {model_max_length}')
            else:
                logger.warning(f'Dont not PI/NTK, because {self.config.model_type} model_max_length {model_max_length} <= max_position_embeddings {orig_ctx_len}')
        elif self.config.model_type == 'qwen2':
            orig_ctx_len = getattr(self.config, "max_position_embeddings", None)
            logger.warning(f"{self.config.model_type} 'max_position_embeddings' is {orig_ctx_len}, {self.config.model_type} not support PI/NTK now. Default use '--model_max_length {model_max_length}'")
        else:
            raise ValueError(f"Only support 'llama or qwen2' now, but got '{self.config.model_type}'")

        # 加载模型
        self.model = ModelUtils.load_model(
                                    model_name_or_path,
                                    config=self.config,
                                    load_in_4bit=load_in_4bit,  # 是否 4bit 量化加载
                                    adapter_name_or_path=adapter_name_or_path,  # adapter_name_or_path 为 None 表示不使用lora适配器
                                    device_map='auto',
                                    ).eval()   # 开始 eval 模式

        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path,
                                                    trust_remote_code=True,
                                                    padding_side="left",   # 推理侧需要左pad
                                                    use_fast=True
                                                    # # llama不支持fast
                                                    # use_fast=False if model.config.model_type == 'llama' else True
                                                )
        logger.info(f"config:\n{self.config}")
        logger.info(f"model structure:\n{self.model}")
        logger.info(f"tokenizer:\n{self.tokenizer}")
        
    def _generate(self, 
                    texts:Union[List[str], str],   # 推理文本
                    model_type:str, # llama or qwen2
                    gen_kwargs:dict = {}, # 生成文本的超参
                ):
        """
        目前只能逐行推理一个文本, 暂时不支持并发批量推理
        """
        if isinstance(texts, str):
            texts = [texts]
        texts = [text.strip() for text in texts]

        # 生成超参配置
        if gen_kwargs == {}:
            # gen_kwargs = {