揭秘大语言模型实践：分布式推理的工程化落地才是关键！

engine=DeepSpeedoption.parallel_loading=trueoption.tensor_parallel_degree=2option.model_loading_timeout=600option.model_id=model/LLM/bloom-7b1/deepspeed/bloom-7b1option.data_type=fp16option.max_new_tokens=100

mport osimport torchfrom typing import Optional
import deepspeedimport logginglogging.basicConfig(format='[%(asctime)s] %(filename)s %(funcName)s():%(lineno)i [%(levelname)s] %(message)s', level=logging.DEBUG)from djl_python.inputs import Inputfrom djl_python.outputs import Outputfrom transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
predictor = None

def get_model(properties: dict):    model_dir = properties.get("model_dir")    model_id = properties.get("model_id")    mp_size = int(properties.get("tensor_parallel_degree", "2"))    local_rank = int(os.getenv('OMPI_COMM_WORLD_LOCAL_RANK', '0'))    logging.info(f"process [{os.getpid()}  rank is [{local_rank}]]")    if not model_id:        model_id = model_dir    logging.info(f"rank[{local_rank}] start load model")    model = AutoModelForCausalLM.from_pretrained(model_id)    tokenizer = AutoTokenizer.from_pretrained(model_id)    logging.info(f"rank[{local_rank}] success load model")
    model = deepspeed.init_inference(model,                                     mp_size=mp_size,                                     dtype=torch.float16,                                     replace_method='auto',                                     replace_with_kernel_inject=True)    logging.info(f"rank[{local_rank}] success to convert model to deepspeed kernel")
    return pipeline(task='text-generation',                    model=model,                    tokenizer=tokenizer,                    device=local_rank)

def handle(inputs: Input) -> Optional[Output]:    global predictor    if not predictor:        predictor = get_model(inputs.get_properties())
    if inputs.is_empty():        # Model server makes an empty call to warmup the model on startup        return None
    data = inputs.get_as_string()    output = Output()    output.add_property("content-type", "application/json")    result = predictor(data, do_sample=True, max_new_tokens=50)    return output.add(result)

arena serve custom \    --name=bloom7b1-deepspeed \    --gpus=2 \    --version=alpha \    --replicas=1 \    --restful-port=8080 \    --data=bloom7b1-pvc:/model \    --image=ai-studio-registry-vpc.cn-beijing.cr.aliyuncs.com/kube-ai/djl-serving:2023-05-19 \    "djl-serving -m "

# 使用 kubectl 启动port-forwardkubectl  -n default-group port-forward svc/bloom7b1-deepspeed-alpha 9090:8080

# 打开新的终端，执行下列命令$ curl -X POST http://127.0.0.1:9090/predictions/deepspeed -H "Content-type: text/plain" -d "I'm very thirsty, I need"[  {    "generated_text":"I'm very thirsty, I need some water.\nWhat are you?\n- I'm a witch.\n- I thought you'd say that.\nI know a great witch.\nShe's right in here.\n- You know where we can go?\n- That's right, in one moment.\n- You want to"  }]

% curl -X POST http://deepspeed-bloom7b1.c78d407e5fa034a5aa9ab10e577e75ae9.cn-beijing.alicontainer.com/predictions/deepspeed -H "Content-type: text/plain" -d "I'm very thirsty, I need"[  {    "generated_text":"I'm very thirsty, I need to drink!\nI want more water.\nWhere is the water?\nLet me have the water, let me have the water...\nWait!\nYou're the father aren't you?\nDo you have water?\nAre you going to let me have some?\nGive me the"  }]