cann/cann-recipes-infer DeepSeek-R1推理优化

DeepSeek-R1或Kimi-K2模型在NPU实现高性能推理【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer概述DeepSeek-R1和Kimi-K2都是2025年开源的大语言模型二者结构类似代码可以复用。本样例基于Deepseek开源代码进行迁移并完成对应的优化适配。支持的产品型号Atlas A3 系列产品环境准备安装CANN软件包。本样例的编译执行依赖CANN开发套件包cann-toolkit与CANN二进制算子包cann-kernels支持的CANN软件版本为CANN 8.5.0。请从软件包下载地址下载Ascend-cann-toolkit_${version}_linux-${arch}.run与Ascend-cann-A3-ops_version_linux-arch.run软件包并参考CANN安装文档进行安装。${version}表示CANN包版本号如8.5.0。${arch}表示CPU架构如aarch64、x86_64。安装Ascend Extension for PyTorchtorch_npu。Ascend Extension for PyTorchtorch_npu为支撑PyTorch框架运行在NPU上的适配插件本样例支持的Ascend Extension for PyTorch版本为v7.3.0PyTorch版本为2.8.0。 请从软件包下载地址下载torch_npu-2.8.0.post2-cp311-cp311-manylinux_2_28_${arch}.whl安装包参考torch_npu安装文档进行安装。${arch}表示CPU架构如aarch64、x86_64。下载项目源码并安装依赖的python库。# 下载项目源码以master分支为例 git clone https://gitcode.com/cann/cann-recipes-infer.git # 安装依赖的python库仅支持python 3.11 cd cann-recipes-infer pip3 install -r ./models/deepseek_r1/requirements.txt配置样例运行所需环境信息。修改executor/scripts/set_env.sh中的如下字段:IPs配置所有节点的IP按照rank id排序多个节点的ip通过空格分开例如(xxx.xxx.xxx.xxx xxx.xxx.xxx.xxx)。cann_path: CANN软件包安装路径例如/usr/local/Ascend/ascend-toolkit/latest。说明HCCL相关配置如HCCL_SOCKET_IFNAME、HCCL_OP_EXPANSION_MODE可以参考集合通信文档并在executor/scripts/function.sh中自定义配置。权重准备请根据所使用的模型类型自行下载原始权重到本地路径例如/data/models/origin/。Deepseek-R1与Kimi-K2的原始权重下载地址如下Deepseek-R1权重Kimi-K2权重权重转换在各个节点上使用weight_convert.sh脚本完成fp8到bfloat16/int8权重转换。入参介绍input_fp8_hf_path原始fp8权重路径output_hf_path转换后输出的权重路径quant_mode量化模式如果权重转换的运行环境为NPU需要先执行cann_path/usr/local/Ascend/ascend-toolkit/latest # cann包安装路径 source ${cann_path}/bin/setenv.bash权重转换拉起示例# 转换为bfloat16权重适用于DeepSeek-R1和Kimi-K2。 bash utils/weight_convert.sh --input_fp8_hf_path /data/models/origin/DeepSeek-R1-FP8 --output_hf_path /data/models/origin/DeepSeek-R1-Bfloat16 --quant_mode bfloat16 # 转换为W8A8C16权重适用于DeepSeek-R1和Kimi-K2。 bash utils/weight_convert.sh --input_fp8_hf_path /data/models/origin/DeepSeek-R1-FP8 --output_hf_path /data/models/origin/DeepSeek-R1-W8A8C16 --quant_mode w8a8c16 # 转换为W8A8C8权重仅适用于DeepSeek-R1。 bash utils/weight_convert.sh --input_fp8_hf_path /data/models/origin/DeepSeek-R1-FP8 --output_hf_path /data/models/origin/DeepSeek-R1-W8A8C8 --quant_mode w8a8c8注意仅DeepSeek-R1支持转W8A8C8权重。推理执行配置推理执行需要加载的权重文件以及YAML文件。修改YAML文件中model_path参数。关于YAML文件中的更多配置说明可参见YAML参数描述。在models/deepseek_r1/config目录下已提供了较优性能的YAML样例供您参考您可以根据模型类型、集群规模以及量化类型选择对应的YAML文件本文以models/deepseek_r1/config/decode_r1_rank_16_16ep_a8w8.yaml文件为例修改其中的model_path参数将其设置为权重转换阶段准备好的权重文件存储路径例如/data/models/origin/DeepSeek-R1-W8A8。修改models/deepseek_r1/infer.sh脚本中YAML_FILE_NAME参数。将YAML_FILE_NAME设置为config文件夹下YAML文件名称例如decode_r1_rank_16_16ep_a8w8.yaml。准备输入prompt。使用内置prompt。本样例已在dataset/default_prompt.json中内置了输入prompt若您直接使用内置prompt本步骤可直接跳过。当然您也可以在dataset/default_prompt.json文件中自定义prompt输入。使用长序列prompt。本样例默认使用内置prompt若您需要使用长序列prompt可以选择LongBench数据集或者InfiniteBench数据集。需要执行以下操作修改YAML文件中的dataset参数将其修改为dataset: LongBench/dataset: InfiniteBench使用LongBench数据集或InfiniteBench数据集作为长序列prompt。若您选择LongBench数据集且机器无法联网需要您从huggingface手动下载数据集至dataset/LongBench目录下LongBench文件夹需手工创建目录中包含LongBench.py和data目录并需要在LongBench.py中修改数据集加载路径若您的机器可正常联网样例执行过程中会自动在线读取LongBench数据集您无需手工下载。若您选择InfiniteBench数据集需要从链接中下载长序列输入数据集longbook_qa_eng并上传到各个节点上新建的路径 dataset/InfiniteBench下。mkdir -p dataset/InfiniteBench说明在使用LongBench/InfiniteBench数据集或其他自定义数据集时默认执行文本摘要任务可在cann-recipes-infer/executor/utils/data_utils.py的build_dataset_input函数里修改默认的system prompt。长序列请求执行中若出现out of memory问题可参见附录中的长序列请求out of memory问题处理。执行推理脚本。cd models/deepseek_r1 bash infer.sh说明如果是多机环境需要在每个节点上执行。优化点参考本样例prefill阶段采用的详细优化点介绍可参见基于Atlas A3集群的DeepSeek-R1模型prefill阶段推理性能优化实践。本样例decode阶段采用的详细优化点介绍及性能Benchmark可参见基于Atlas A3集群的DeepSeek-R1模型decode阶段推理性能优化。Benchmark基于Atlas A3本实践使用config/decode_r1_rank_128_128ep_a8w8c8_mtp_benchmark.yaml作为运行配置文件对DeepSeek-R1 W8A8C8 量化版本进行了性能Benchmark测试。 |Quant Mode| Global Batch Size | Seq Length | Chips | TPOT (ms) | Throughput (tokens/p/s) | |-------| ----------------- | ---------- | ----- | --------- | ----------------------- | |W8A8C8 | 6144 | 4096 | 64 | 44.9 |2138 |注性能数据基于 MTP1 与 perfect eplb 配置采集平均 1 个 draft token 中 accept token 为 0.7 个。附录常见问题处理长序列请求out of memory问题处理长序列请求可能导致device内存out of memory尤其是在prefill阶段:Attention的Softmax操作通常为float32计算其内存大小为batch_size * num_heads * q_s * kv_s * (2Bytes 4Bytes)。MoE的Routing分发可能存在极端负载不均导致个别卡上的grouped_matmul算子占用较大内存。为了缓解这两处峰值带来的OOM问题可分别采用以下方法使能Paged Attention进行内存管理并调用Flash Attention融合算子算子内会切块计算Attention避免了q_s * kv_s的峰值内存产生。Prefill内存通常与batch_size大小成正比当decode需要推理的global batch size过大时prefill可能会由于OOM而无法在一轮推理中处理完所有的batch因此我们可进行多次小batch串行推理从而降低峰值内存。当前executor-core的prefill默认采用packed sequence执行可通过配置YAML中的max_prefill_tokens限制单次prefill batch的总prompt token数从而控制prefill阶段的峰值内存。为了缓解MoE负载不均带来的峰值内存我们可进行Chunk MoE推理即在MoE切Chunk串行推理降低极端场景下的峰值内存可通过YAML中的moe_chunk_max_len开关设置chunk的大小。当前该开关只针对prefill生效开启后由于MoE部分将串行计算各chunk会对prefill的性能产生相应的影响。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考