批处理？推理速度快一万倍！！！！ 大模型批量推理（LLM Generate in Parallel）

众所周知，CPU只能多进程串行计算，而GPU是天然能进行多进程并行计算。数据标注作为一个现在常用的方式，往往使用大模型对图像进行标注，利用大模型的理解能力重新合成图像所对应的文本。

huggingface.co作为目前最大的人工智能社区，推出的transformers库已经成为通用大模型的一种代码规范。我们能很轻易的使用transformers完成对大模型的使用即推理。

使用llama3进行推理

使用pipline

import transformers
import torch

model_id = "meta-llama/Meta-Llama-3-8B"

pipeline = transformers.pipeline(
    "text-generation", model=model_id, model_kwargs={"torch_dtype": torch.bfloat16}, 	     device_map="auto"
)
pipeline("Hey how are you doing today?")

使用AutoModelForCausalLM

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
model = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Meta-Llama-3-8B-Instruct",
        torch_dtype=torch.bfloat16,
    )
content = "Hey how are you doing today?"   
messages = [
                {
                    "role": "user",
                    "content": content
                },
            ]
input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
terminators = [
                tokenizer.eos_token_id,
                tokenizer.convert_tokens_to_ids("<|eot_id|>")
]
outputs = model.generate(
                input_ids,
                max_new_tokens=256,
                eos_token_id=terminators,
                do_sample=True
            )
response = outputs[0][input_ids.shape[-1]:]
decoded = tokenizer.decode(response, skip_special_tokens=True)

显存占用如下

image-20240728134118957

image-20240728134141529

运行速度如下，可以看到1万卡图像推理大概需要7小时

image-20240728134243420

既然显卡显存未完全占用，再加上显卡本身支持批处理，能不能批量计算？从理论上分析是批量计算完全没有任何问题，但未查到相关代码。

使用llama3进行批量推理

批量推理最重要的就是构造批量数据，一般来说，模型都是采用右填充，但对大语言模型来说，通常是做预测下一个token的任务，如果采用右填充会一直重复输出不相关的词汇，所以在推理时通常采用左填充。对于llama3来说未设置padding token，我们将eos_token作为padding_token。

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
model = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Meta-Llama-3-8B-Instruct",
        torch_dtype=torch.bfloat16,
    )

if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = 'left'
input_ids = tokenizer(
                texts_list, return_tensors="pt", padding="longest", truncation=True
            ).to(model.device)

terminators = [
                tokenizer.eos_token_id,
                tokenizer.convert_tokens_to_ids("<|eot_id|>")
            ]

outputs = model.generate(
                max_new_tokens=256,
                eos_token_id=terminators,
                do_sample=True,
                **input_ids
            )
for index, output in enumerate(outputs):
                response = output[len(input_ids[index].attention_mask):]
                decoded = tokenizer.decode(response, skip_special_tokens=True)
                try:
                    decoded = decoded.split('\n')[-1]
                except:
                    pass
                captions.append([decoded])

显存占用如下

image-20240728135317459

运行速度如下

image-20240724150212730

可以看到运算速度从7小时变成了7分钟，后期计算把batch拉大，速度继续提升，最终几千万样本数据生成只花费2天。