昇思+昇騰開發板：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)

打印時間：