【假期 AI 充电】揭秘大语言模型实践：分布式推理的工程化落地才是关键！

engine=DeepSpeed
option.parallel_loading=true
option.tensor_parallel_degree=2
option.model_loading_timeout=600
option.model_id=model/LLM/bloom-7b1/deepspeed/bloom-7b1
option.data_type=fp16
option.max_new_tokens=100

mport os
import torch
from typing import Optional
import deepspeed
import logging
logging.basicConfig(format='[%(asctime)s] %(filename)s %(funcName)s():%(lineno)i [%(levelname)s] %(message)s', level=logging.DEBUG)
from djl_python.inputs import Input
from djl_python.outputs import Output
from 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 get pod | grep bloom7b1-deepspeed-alpha-custom-serving
bloom7b1-deepspeed-alpha-custom-serving-766467967d-j8l2l    1/1     Running     0          8s
# 查看启动日志
kubectl logs bloom7b1-deepspeed-alpha-custom-serving-766467967d-j8l2l -f

INFO  ModelServer Starting model server ...
INFO  ModelServer Starting djl-serving: 0.23.0-SNAPSHOT ...
INFO  ModelServer
INFO  PyModel Loading model in MPI mode with TP: 2.
INFO  PyProcess [1,0]<stdout>:process [92  rank is [0]]
INFO  PyProcess [1,0]<stdout>:rank[0] start load model
INFO  PyProcess [1,1]<stdout>:process [93  rank is [1]]
INFO  PyProcess [1,1]<stdout>:rank[1] start load model
INFO  PyProcess [1,0]<stdout>:rank[0] success to convert model to deepspeed kernel
INFO  PyProcess [1,1]<stdout>:rank[1] success to convert model to deepspeed kernel
INFO  PyProcess [1,0]<stdout>:rank[0] success load model
INFO  PyProcess [1,1]<stdout>:rank[1] success load model
INFO  PyProcess Model [deepspeed] initialized.
INFO  PyProcess Model [deepspeed] initialized.
INFO  PyModel deepspeed model loaded in 297083 ms.
INFO  ModelServer Initialize BOTH server with: EpollServerSocketChannel.
INFO  ModelServer BOTH API bind to: http://0.0.0.0:8080

# 使用 kubectl 启动port-forward
kubectl  -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"
  }
]