昇思+昇腾开发板：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_token
        print(new_token, end="", flush=True)

    messages.append({'role': 'assistant', 'content': partial_message})
    return messages[1:]

交互界面实现

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

import os
import platform


os_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":
        break
    if query.strip() == "clear":
        os.system(clear_command)
        print(welcome_prompt)
        continue
    print("\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 > p
    sorted_probs[mask] = 0.0
    sorted_probs = sorted_probs / np.sum(sorted_probs, axis=-1, keepdims=True)
    # 转换回MindSpore Tensor
    sorted_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 += 1
    t = time.time()
    # 打印单步生成耗时
    print("[%d]:" % i, t - s)

text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
print(text)

打印时间：