模型和环境
构建环境
# 本地获取ngc pytorch原生镜像
docker pull nvcr.io/nvidia/pytorch:23.02-py3
# 启动容器
docker run --gpus all -itd --rm --name lyrallms_cu12 nvcr.io/nvidia/pytorch:23.02-py3
docker exec -it lyrallms_cu12 bash
获取代码后安装依赖
pip install -r requirements.txt
将lyralib下对应cuda版本的so文件 复制到/usr/lib/lyralib下。
推理使用
使用核心片段
from lyra_baichuan import lyraBaichuan7B, lyraBaichuan13B
model_path = 'XXX' # 包含转换后的模型参数，配置，tokenizer文件目录
tokenizer_path = 'XXX'
data_type = 'fp16' # 推理精度
memopt_mode = 1
# 加载加速后的模型，C++ 底层已经掩盖，依赖加速编译的 /usr/lib/ftlib 下的 so 库，已经打在镜像中
# 模型加载需要花一些时间，建议把下载的模型参数解压到本地磁盘
# 如需使用Baichuan1/2-7B模型，下方更换为：lyraBaichuan7B(model_path, tokenizer_path, data_type, memopt_mode)
model = lyraBaichuan13B(model_path, tokenizer_path, data_type, memopt_mode)
# 输入, 若有多个输入，可batch 推理，prompts 支持列表，这里为模拟多个输入，直接复制 32 分，batchsize 达到32
prompts = "登鹳雀楼->王之涣\n夜雨寄北->"
prompts = [prompts,]*32
# 生成, 最大长度可自行调整，这里设置 64，模型遇到 end token 或者达到最大计算长度时会停止当前批次计算.
output_texts = model.generate(prompts, output_length=64, do_sample=False, top_k=30, top_p=0.85, temperature=1.0, repetition_penalty=1.0)
# 输出查看， 虽然 输入字符串也是参与计算，用于评估模型吞吐量和计算速度。
# 这里为了显示应用方便， output_texts 中每个样本的输出已经去掉了输入字符串
print(output_texts)
# 输出示例
>>> Inputs: 登鹳雀楼->王之涣
夜雨寄北->
>>> Outputs:
李商隐
望洞庭->刘禹锡
黄鹤楼送孟浩然之广陵->李白
登岳阳楼->杜甫
秋词->刘禹锡
枫桥夜泊->张继
饮湖上初晴后雨->苏轼
浪淘沙->刘禹锡
demo 脚本
examples/batch_demo.py 中有类似上面的使用示例，做了简单的跑速测试，考虑大家对 token 的理解各有不同，我们这里直接按字符数来评估，不同 token 的理解可以自行根据生成结果字符数去观测。
更多测试脚本及用法详见参考 examples 下的 README.md ，如：
- Batch推理
- 不等长Batch推理
- Batch流式推理