[Python]BLOOM模型推理简介

BLOOM模型是Huggingface开发的，在transformers库中提供了支持：

• 通过AutoTokenizer从模型中加载Tokenizer；
• 通过AutoModelForCausalLM加载bloom模型；
• 通过generate进行推理：推理的结果包含输入，所以在计算新token数量或输出时，需要去掉头部的输入。

pytorch推理

pyTorch是由Facebook基于Torch开发的。

在pytorch中：

• 通过env（CUDA_VISIBLE_DEVICES）来设定所需的GPU卡；
• 通过cuda.is_available来判断是否存在cuda显卡，获取卡设备后，可通过to(device)把数据放到卡上；
  device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

import os
import gc
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

os.environ["CUDA_VISIBLE_DEVICES"] = '0'
input_text = "你好"

batch_size = 1
gpu_mem_map = [30]
data_type = torch.float16
model_path = "/workspace/bloom-7b"

# load tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_path, truncation_side='left')

# load model
# 通过device_map与max_memory可把模型自动加载到多卡上（并且限制每个卡上加载参数大小）
model_params = {"low_cpu_mem_usage": True, "device_map": 'auto', "torch_dtype": data_type}
max_memory = {}
for i in range(len(gpu_mem_map)):
    max_memory[i] = f'{gpu_mem_map[i]}GiB'
model_params['max_memory'] = max_memory

model = AutoModelForCausalLM.from_pretrained(model_path, **model_params)
print(model.gen_config)

##############################
# inference
# gc.collect()
# torch.cuda.empty_cache()

# build input tokens
inputs = [input_text]
inputs *= batch_size

input_tokens = tokenizer(inputs, return_tensors="pt")
# 输入token要与模型在同样设备上（GPU卡），需要通过cuda把token全部放到卡上
for t in input_tokens:
    if torch.is_tensor(input_tokens[t]):
        input_tokens[t] = input_tokens[t].cuda()

# generate
gen_params = dict(max_new_tokens=1000, do_sample=False, top_k=1)
with torch.no_grad():
    output_tokens = model.generate(**input_tokens, **gen_params)

# 生成的输出包含输入，所以需要去掉头部的输入信息
new_tokens = len(output_tokens[0]) - len(input_tokens[0])
print("new tokens:", new_tokens)

output_text = tokenizer.batch_decode(output_tokens, skip_special_tokens=True)
outputs = output_text[0][len(inputs[0]):].strip()
print(outputs)

ds_zero推理

DeepSpeed是微软推出的大规模模型分布式训练的工具，主要实现了ZeRO并行训练算法。

在加载模型前，一定要先通过HfDeepSpeedConfig来设定配置信息。模型加载完成后通过torch.cuda.current_device()获取当前所使用的设备；后续在推理前，把数据放到此设备上。

在获取输出单词数量时，中英文使用不同的方式计算（中文每个字为一个，英文根据单词数量判断）。

import argparse
import gc
import time
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, GenerationConfig
from transformers.deepspeed import HfDeepSpeedConfig
import deepspeed

# usage: deepspeed --num_gpus 1 inference.py

def isChinese(words):
    for w in words:
        if '\u4e00' <= w <= '\u9fff':
            return True
    return False


def main(args):
    data_type = torch.float16
    train_micro_batch_size_per_gpu = 1

    world_size = args.world_size
    train_batch_size = train_micro_batch_size_per_gpu*world_size

    model_path = args.model_path
    input_text = args.input_text

    deepspeed.init_distributed("nccl")

    # load model
    config = AutoConfig.from_pretrained(model_path)
    ds_config = {
        "fp16": {
            "enabled": data_type == torch.float16,
        },
        "bf16": {
            "enabled": data_type == torch.bfloat16,
        },
        "zero_optimization": {
            "stage": 3,
            "overlap_comm": True,
            "contiguous_gradients": True,
            "reduce_bucket_size": config.hidden_size * config.hidden_size,
            "stage3_prefetch_bucket_size": 0.9 * config.hidden_size * config.hidden_size,
            "stage3_param_persistence_threshold": 0,
        },
        "steps_per_print": 2000,
        "train_batch_size": train_batch_size,
        "train_micro_batch_size_per_gpu": train_micro_batch_size_per_gpu,
        "wall_clock_breakdown": False,
    }
    
    # 设定此配置信息，以保证模型直接在GPU上加载
    hfConf = HfDeepSpeedConfig(ds_config)
    print("**Config:", hfConf)

    model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=data_type)
    model = model.eval()

    # 把模型转换为使用Zero方式
    model = deepspeed.initialize(model=model, config_params=ds_config)[0]

    model.module.eval()
    model = model.module

    # 获取当前模型使用的设备，以便后续把数据放到此设备上
    device = torch.cuda.current_device()
    print("***device:", device)

    tokenizer = AutoTokenizer.from_pretrained(model_path, truncation_side='left')

    # inference
    #gc.collect()
    #torch.cuda.empty_cache()
    inputs = [input_text]

    generate_kwargs = dict(max_new_tokens=200, do_sample=False, top_k=1)
    gen_config = GenerationConfig.from_model_config(AutoConfig.from_pretrained(model_path))
    gen_config.update(**generate_kwargs)

    outputs, words = gen_tokens(device, gen_config, inputs, model, tokenizer)
    print(outputs)


def gen_tokens(device, gen_config, inputs, model, tokenizer):
    input_tokens = tokenizer(inputs, return_tensors="pt")
    for t in input_tokens:
        if torch.is_tensor(input_tokens[t]):
            input_tokens[t] = input_tokens[t].to(device)

    output_tokens = model.generate(**input_tokens, generation_config=gen_config)
    output_text = tokenizer.batch_decode(output_tokens, skip_special_tokens=True)
    outputs = output_text[0][len(inputs[0]):].strip()

    if isChinese(outputs):
        words = len(outputs)
    else:
        words = len(outputs.split())
    return outputs, words


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="deepspeed ds_zero infer")
    parser.add_argument('--model_path',
                        type=str,
                        default="/models/bloom-7b",
                        help="The path of model.")
    parser.add_argument('--input_text',
                        type=str,
                        default="你好",
                        help="The input text for infer.")
    parser.add_argument('--world_size',
                        type=int,
                        default=1,
                        help="The world size")
    parser.add_argument('--local_rank',
                        type=str,
                        default="0",
                        help="The local rank")

    args = parser.parse_args()
    main(args)