大模型镜像打包实战:CodeGeeX2为例

资源地址

docker torch镜像地址
CodeGeeX2 github

构建思路

查看CodeGeeX2项目，官方已经提供好启动脚本，配置好各种依赖应该就可以运行。

python ./demo/run_demo.py

usage: run_demo.py [-h] [--model-path MODEL_PATH] [--example-path EXAMPLE_PATH] [--quantize QUANTIZE]
                   [--chatglm-cpp] [--fastllm] [--n-gpus N_GPUS] [--gpu GPU] [--cpu] [--auth] [--username yourname]
                   [--password yourpassword]
                   [--port PORT] [--listen ADDRESS]

# 若要启用身份验证，请先启用--auth，然后定义--username与--password，如：
python run_demo.py --auth --username user --password password  # 若要监听所有地址请指定 --listen 0.0.0.0

镜像Dockerfile

写一个启动脚本bin/start.sh，判断是否启用身份验证，也可以在Dockerfile最后的CMD直接执行python脚本，sh bin/start.sh这样写更标准。

#!/usr/bin/env bash

if [[ $USERNAME != "" ]] && [[ $PASSWORD != "" ]]
then
    python /workspace/CodeGeeX2/demo/run_demo.py \
        --listen $LISTEN --port $PORT \
        --model-path $MODEL_PATH  \
        --gpu $GPU \
        --n-gpus $N_GPUS \
        $CHATGLM \
        --auth \
        --username $USERNAME \
        --password $PASSWORD
else
    python /workspace/CodeGeeX2/demo/run_demo.py \
        --listen $LISTEN --port $PORT \
        --model-path $MODEL_PATH  \
        --gpu $GPU \
        --n-gpus $N_GPUS \
        $CHATGLM
fi

FROM pytorch/pytorch:2.1.2-cuda11.8-cudnn8-runtime

RUN pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple \
    && pip install fire zmq sentencepiece gradio accelerate transformers==4.33.2 modelscope chatglm-cpp cpm_kernels tabulate \
    && mkdir /workspace \
    && cd /workspace

# 网好的同学直接拉远程的项目
# RUN git clone https://github.com/THUDM/CodeGeeX2.git

# 网不好可以先拉倒本地，然后用下面目录copy到指定路径
COPY CodeGeeX2 /workspace/CodeGeeX2/
# 原项目使用huggingface模型，此处使用的是国内modelscope仓库，代码中transformers的导包都需要改为modelscope，也可以上huggingface上拉倒本地，改成本地路径
COPY modelscope /root/.cache/modelscope/

ENV LISTEN=0.0.0.0
ENV PORT=7860
ENV MODEL_PATH=ZhipuAI/codegeex2-6b
ENV GPU=0
ENV N_GPUS=1
ENV CHATGLM=--chatglm-cpp
CMD sh /workspace/CodeGeeX2/bin/start.sh

执行命令构建镜像：docker build -t codegeex:v1 .
命令执行通过后查看此镜像ID，执行下面命令：
docker run -d --runtime=nvidia -p 7860:7860 --name pytorch2.0 -v /tmp:/tmp 镜像ID
访问localhost:7860查看是否成功

bug解决方案

如果有我未遇到的bug欢迎提问，一起解决问题总结在此处

• NameError: name 'round_up' is not defined
  当指定 --quantize 4 时出现，File "/root/.cache/huggingface/modules/transformers_modules/codegeex2-6b/quantization.py", line 78, in compress_int4_weight可以看到这个python是在拉模型权重中的一个文件，名字就知道是做量化的代码，from cpm_kernels.kernels.base import LazyKernelCModule, KernelFunction, round_up,进入文件后找到导入的语句，显然应该是缺少cpm_kernels包或者包版本不一致没有round_up函数，pip install cpm_kernels适合的版本即可。或者直接使用官方已经提供好的量化模型，就可以不做--quantize 4指定。