昇思+昇腾开发板：DeepSeek-R1-Distill-Qwen-1.5B 模型推理部署与 JIT 优化实践

引言

模型推理部署

环境准备

安装 MindSpore

查看当前 mindspore 版本

安装 MindNLP

模型与分词器加载

导入必要的库

加载分词器

加载模型

对话功能实现

设置系统提示词

构建对话历史输入

推理函数实现

交互界面实现

推理JIT优化

基础环境安装

JIT 优化配置

核心优化组件实现

Top-p 采样函数

模型加载与 JIT 化

JIT 加速的推理函数

自回归生成与优化效果验证

静态缓存与生成流程

优化效果测试

引言

本文将详细介绍如何对 DeepSeek-R1-Distill-Qwen-1.5B 模型进行推理部署，构建可交互的对话机器人，并利用 MindSpore 的 JIT（Just-In-Time）编译技术进行推理优化，以提升模型的响应速度和用户体验。本教程适用于昇思大模型平台的单卡环境，相关操作在昇腾开发板上的实际应用可参考示例代码。

模型推理部署

环境准备

在进行模型部署前，需要先配置合适的运行环境，主要包括安装指定版本的 MindSpore 深度学习框架和 MindNLP 自然语言处理工具库。

安装 MindSpore

MindSpore 是华为推出的开源深度学习框架，本教程使用 2.6.0 版本，安装命令如下：

%%capture captured_output
# 实验环境已经预装了mindspore==2.6.0，如需更换mindspore版本，可更改下面 MINDSPORE_VERSION 变量
!pip uninstall mindspore -y
%env MINDSPORE_VERSION=2.6.0
!pip install https://ms-release.obs.cn-north-4.myhuaweicloud.com/${MINDSPORE_VERSION}/MindSpore/unified/aarch64/mindspore-${MINDSPORE_VERSION}-cp39-cp39-linux_aarch64.whl --trusted-host ms-release.obs.cn-north-4.myhuaweicloud.com -i https://pypi.tuna.tsinghua.edu.cn/simple

查看当前 mindspore 版本

# 查看当前 mindspore 版本
!pip show mindspore

运行结果：

Name: mindspore
Version: 2.6.0
Summary: MindSpore is a new open source deep learning training/inference framework that could be used for mobile, edge and cloud scenarios.
Home-page: https://www.mindspore.cn
Author: The MindSpore Authors
Author-email: contact@mindspore.cn
License: Apache 2.0
Location: /home/mindspore/miniconda/envs/jupyter/lib/python3.9/site-packages
Requires: asttokens, astunparse, dill, numpy, packaging, pillow, protobuf, psutil, safetensors, scipy
Required-by: mindnlp

安装 MindNLP

MindNLP 是基于 MindSpore 的自然语言处理库，提供了丰富的预训练模型和处理工具，本教程使用 0.4.1 版本：

%%capture captured_output
# 安装mindnlp 0.4.1 版本
!pip uninstall mindnlp -y
!pip install https://xihe.mindspore.cn/coderepo/web/v1/file/MindSpore/mindnlp/main/media/mindnlp-0.4.1-py3-none-any.whl

模型与分词器加载

环境配置完成后，需要加载 DeepSeek-R1-Distill-Qwen-1.5B 模型及其对应的分词器。

导入必要的库

from mindnlp.transformers import AutoModelForCausalLM, AutoTokenizer
from mindnlp.transformers import TextIteratorStreamer
from mindnlp.peft import PeftModel
from threading import Thread

运行结果：

加载分词器

分词器的作用是将自然语言文本转换为模型可理解的 token 序列：


# 开启同步，在出现报错，定位问题时开启
# mindspore.set_context(pynative_synchronize=True)# Loading the tokenizer and model from Modelers's model hub.
tokenizer = AutoTokenizer.from_pretrained("MindSpore-Lab/DeepSeek-R1-Distill-Qwen-1.5B-FP16", mirror="modelers")
# 设置pad_token为eos_token
if tokenizer.pad_token is None:tokenizer.pad_token = tokenizer.eos_token

运行结果：

加载模型

使用 AutoModelForCausalLM 加载预训练的因果语言模型：

model = AutoModelForCausalLM.from_pretrained("MindSpore-Lab/DeepSeek-R1-Distill-Qwen-1.5B-FP16", mirror="modelers")
# adapter_model path
# model = PeftModel.from_pretrained(model, "./output/DeepSeek-R1-Distill-Qwen-1.5B/adapter_model_for_demo/")

加载过程中可能会出现一些警告信息，如关于模型未继承 GenerationMixin 类和滑动窗口注意力机制的提示，这些通常不影响基本功能的使用。

运行结果：

对话功能实现

为了实现与用户的交互，需要构建处理对话历史、生成模型响应的相关函数。

设置系统提示词

系统提示词用于定义模型的行为和角色：

system_prompt = "你是一个智能聊天机器人，以最简单的方式回答用户问题"

构建对话历史输入

该函数将历史对话和当前用户输入整理成模型所需的格式：

def build_input_from_chat_history(chat_history, msg: str):messages = [{'role': 'system', 'content': system_prompt}]for info in chat_history:role, content = info['role'], info['content']messages.append({'role': role, 'content': content})messages.append({'role': 'user', 'content': msg})return messages

推理函数实现

推理函数负责将用户输入转换为模型输入，调用模型生成响应，并处理流式输出：

def inference(message, history):messages = build_input_from_chat_history(history, message)input_ids = tokenizer.apply_chat_template(messages,add_generation_prompt=True,return_tensors="ms",tokenize=True)streamer = TextIteratorStreamer(tokenizer, timeout=300, skip_prompt=True, skip_special_tokens=True)generate_kwargs = dict(input_ids=input_ids,streamer=streamer,max_new_tokens=1024,use_cache=True,)t = Thread(target=model.generate, kwargs=generate_kwargs)t.start()  # Starting the generation in a separate thread.partial_message = ""for new_token in streamer:partial_message += new_tokenprint(new_token, end="", flush=True)messages.append({'role': 'assistant', 'content': partial_message})return messages[1:]

交互界面实现

为了方便用户使用，构建一个简单的命令行交互界面：

import os
import platformos_name = platform.system()
clear_command = 'cls' if os_name == 'Windows' else 'clear'
welcome_prompt = '欢迎使用 DeepSeek-R1-Distill-Qwen-1.5B 模型，输入内容即可进行对话，clear 清空对话历史，stop 终止程序'
print(welcome_prompt)
history = []
while True:query = input("\n用户：")if query.strip() == "stop":breakif query.strip() == "clear":os.system(clear_command)print(welcome_prompt)continueprint("\nDeepSeek-R1-Distill-Qwen-1.5B：", end="")history = inference(query, history)print("")

运行结果：

输入框中输入：qianduanjidi回车

推理JIT优化

在大语言模型的实际应用中，推理速度直接影响用户体验。本文将详细介绍如何基于 MindSpore 框架的 JIT（Just-In-Time）编译技术，对 DeepSeek-R1-Distill-Qwen-1.5B 模型进行推理优化，通过降低单次推理耗时提升对话响应速度。

基础环境安装

JIT 优化依赖特定版本的 MindSpore 和 MindNLP，需先完成环境配置：

# 卸载现有MindSpore版本
!pip uninstall mindspore -y
# 指定MindSpore版本为2.6.0并安装
%env MINDSPORE_VERSION=2.6.0
!pip install https://ms-release.obs.cn-north-4.myhuaweicloud.com/${MINDSPORE_VERSION}/MindSpore/unified/aarch64/mindspore-${MINDSPORE_VERSION}-cp39-cp39-linux_aarch64.whl --trusted-host ms-release.obs.cn-north-4.myhuaweicloud.com -i https://pypi.tuna.tsinghua.edu.cn/simple# 安装适配的MindNLP 0.4.1版本
!pip uninstall mindnlp -y
!pip install https://xihe.mindspore.cn/coderepo/web/v1/file/MindSpore/mindnlp/main/media/mindnlp-0.4.1-py3-none-any.whl

安装完成后，可通过pip show mindspore确认版本为 2.6.0，确保环境一致性。

JIT 优化配置

MindSpore 的 JIT 编译需要通过set_context进行参数配置，开启图算融合和指定优化级别：

import mindspore
from mindnlp.transformers import AutoTokenizer, AutoModelForCausalLM, StaticCache
from mindnlp.core import ops
from mindnlp.configs import set_pyboost
import time
import numpy as np# 开启O2级别的jit优化，开启图算融合
mindspore.set_context(enable_graph_kernel=True,mode=mindspore.GRAPH_MODE,jit_config={"jit_level": "O2",},
)

运行结果：

核心优化组件实现

Top-p 采样函数

采样环节是生成式模型的关键步骤，为提升效率，采用基于 NumPy 的实现（在边缘设备如香橙派上表现更优）：

def sample_top_p(probs, p=0.9):"""Top-p采样函数，用于生成文本时选择下一个token。此处优先采用基于numpy而不是原生MindSpore的实现方式，因为在香橙派上运行效率更高"""probs_np = probs.asnumpy()# 按概率降序排序sorted_indices = np.argsort(-probs_np, axis=-1)sorted_probs = np.take_along_axis(probs_np, sorted_indices, axis=-1)# 计算累积概率并创建掩码cumulative_probs = np.cumsum(sorted_probs, axis=-1)mask = cumulative_probs - sorted_probs > psorted_probs[mask] = 0.0sorted_probs = sorted_probs / np.sum(sorted_probs, axis=-1, keepdims=True)# 转换回MindSpore Tensorsorted_probs_tensor = mindspore.Tensor(sorted_probs, dtype=mindspore.float32)sorted_indices_tensor = mindspore.Tensor(sorted_indices, dtype=mindspore.int32)next_token_idx = ops.multinomial(sorted_probs_tensor, 1)batch_size = probs.shape[0]batch_indices = ops.arange(0, batch_size, dtype=mindspore.int32).reshape(-1, 1)# 此处采用基于mindspore.ops的实现方式，在香橙派上兼容性最好# next_token = sorted_indices_tensor[batch_indices, next_token_idx]next_token = mindspore.ops.gather(sorted_indices_tensor, next_token_idx, axis=1, batch_dims=1)# next_token = mindspore.mint.gather(sorted_indices_tensor, dim=1, index=next_token_idx)return next_token

模型加载与 JIT 化

加载模型并通过model.jit()实现全图静态化，为编译优化奠定基础：

# 该任务将使用DeepSeek-R1-Distill-Qwen-1.5B模型，对给定的prompt进行补齐
prompts = ["请介绍一下自己。<think>","My favorite all time favorite condiment is ketchup.",
]# 生成参数配置
NUM_TOKENS_TO_GENERATE = 40  # 每个输入要生成的token数量
TEMPERATURE = 0.8            # 温度参数（控制生成多样性）
TOP_P = 0.8                  # Top-p采样阈值model_id = "MindSpore-Lab/DeepSeek-R1-Distill-Qwen-1.5B-FP16"
tokenizer = AutoTokenizer.from_pretrained(model_id, mirror="modelers")
model = AutoModelForCausalLM.from_pretrained(model_id, low_cpu_mem_usage=True, mirror="modelers")# 使用model.jit()将全图静态图化
model.jit()inputs = tokenizer(prompts, return_tensors="ms", padding=True)
set_pyboost(False)

运行结果：

JIT 加速的推理函数

通过@mindspore.jit装饰器对核心推理函数进行编译优化，减少重复计算开销：

# 使用@mindspore.jit装饰器封装模型推理函数
@mindspore.jit(jit_config=mindspore.JitConfig(jit_syntax_level='STRICT'))
def get_decode_one_tokens_logits(model, cur_token, input_pos, cache_position, past_key_values, temperature=TEMPERATURE, top_p=TOP_P):"""单个token的解码函数，返回logits，可以使用jit进行优化"""logits = model(cur_token,position_ids=input_pos,cache_position=cache_position,past_key_values=past_key_values,return_dict=False,use_cache=True)[0]return logits
def decode_one_tokens(model, cur_token, input_pos, cache_position, past_key_values, temperature=TEMPERATURE, top_p=TOP_P):"""单个token的解码函数，由logits、温度和Top_p选择合适的token"""logits = get_decode_one_tokens_logits(model, cur_token, input_pos, cache_position, past_key_values, temperature, top_p)if temperature > 0:probs = mindspore.mint.softmax(logits[:, -1] / temperature, dim=-1)new_token = sample_top_p(probs, top_p)else:new_token = mindspore.mint.argmax(logits[:, -1], dim=-1)[:, None]return new_token

自回归生成与优化效果验证

静态缓存与生成流程

使用StaticCache缓存注意力计算结果，减少自回归生成中的重复计算：

batch_size, seq_length = inputs["input_ids"].shape# 创建静态缓存（用于加速自回归生成）
past_key_values = StaticCache(config=model.config, max_batch_size=2, max_cache_len=512, dtype=model.dtype
)
cache_position = ops.arange(seq_length)
generated_ids = ops.zeros(batch_size, seq_length + NUM_TOKENS_TO_GENERATE + 1, dtype=mindspore.int32
)
generated_ids[:, cache_position] = inputs["input_ids"].to(mindspore.int32)# 初始前向传播获取首个logits
logits = model(**inputs, cache_position=cache_position, past_key_values=past_key_values,return_dict=False, use_cache=True
)[0]

优化效果测试

通过循环生成并记录单步耗时，验证 JIT 优化效果：

# 生成第一个新token
if TEMPERATURE > 0:probs = mindspore.mint.softmax(logits[:, -1] / TEMPERATURE, dim=-1)next_token = sample_top_p(probs, TOP_P)
else:next_token = mindspore.mint.argmax(logits[:, -1], dim=-1)[:, None]generated_ids[:, seq_length] = next_token[:, 0]# 自回归生成循环
cache_position = mindspore.tensor([seq_length + 1])
for i in range(1, NUM_TOKENS_TO_GENERATE):s = time.time()next_token = decode_one_tokens(model, next_token, None, cache_position, past_key_values)generated_ids[:, cache_position] = next_token.int()cache_position += 1t = time.time()# 打印单步生成耗时print("[%d]:" % i, t - s)text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
print(text)

打印时间：