当前位置：首页 > news >正文

Qwen2.5-1.5B GPU显存优化教程：torch.no_grad+清空对话按钮双策略详解

news 2026/5/12 19:18:41

Qwen2.5-1.5B GPU显存优化教程：torch.no_grad+清空对话按钮双策略详解

1. 项目背景与显存优化需求

Qwen2.5-1.5B是阿里通义千问推出的轻量级大语言模型，虽然参数量只有15亿，但在本地部署时仍然可能遇到GPU显存不足的问题。特别是在长时间对话或多轮交互场景下，显存占用会逐渐累积，最终导致程序崩溃或响应变慢。

本教程将详细解析两种实用的显存优化策略：使用torch.no_grad()禁用梯度计算和添加清空对话按钮。这两种方法结合使用，可以有效管理显存占用，让Qwen2.5-1.5B在资源受限的环境中稳定运行。

2. 环境准备与基础配置

2.1 安装必要依赖

首先确保你的环境中安装了必要的Python包：

pip install torch transformers streamlit

2.2 模型文件准备

将Qwen2.5-1.5B-Instruct模型文件存放在本地目录，例如/root/qwen1.5b。确保目录包含以下文件：

config.json
tokenizer相关文件
model权重文件（pytorch_model.bin或.safetensors）

3. 基础模型加载代码

让我们先看看基础的模型加载和推理代码：

import torch from transformers import AutoModelForCausalLM, AutoTokenizer import streamlit as st # 模型路径 MODEL_PATH = "/root/qwen1.5b" @st.cache_resource def load_model(): # 自动选择设备和数据类型 model = AutoModelForCausalLM.from_pretrained( MODEL_PATH, device_map="auto", torch_dtype="auto", trust_remote_code=True ) tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True) return model, tokenizer model, tokenizer = load_model()

这段代码使用了device_map="auto"和torch_dtype="auto"来自动选择最适合的计算设备和数据类型，这是第一层优化。

4. 核心优化策略一：torch.no_grad()

4.1 什么是torch.no_grad()

torch.no_grad()是PyTorch中的一个上下文管理器，它告诉PyTorch在代码块内部不计算梯度。在推理阶段，我们不需要计算梯度，因为不需要更新模型参数。

4.2 为什么能节省显存

在深度学习中，梯度计算需要保存前向传播的中间结果，这些中间结果会占用大量显存。通过禁用梯度计算，可以显著减少显存使用量。

4.3 实际代码实现

def generate_response(model, tokenizer, prompt, chat_history=[]): # 将对话历史转换为模型需要的格式 messages = chat_history + [{"role": "user", "content": prompt}] # 应用聊天模板 text = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) # 编码输入文本 inputs = tokenizer(text, return_tensors="pt").to(model.device) # 使用torch.no_grad()禁用梯度计算 with torch.no_grad(): outputs = model.generate( inputs.input_ids, max_new_tokens=1024, temperature=0.7, top_p=0.9, do_sample=True ) # 解码生成结果 response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True) return response

4.4 效果对比

使用torch.no_grad()前后，显存占用通常可以减少30-40%。对于Qwen2.5-1.5B模型，这意味著可以节省大约1-2GB的显存占用。

5. 核心优化策略二：清空对话按钮

5.1 为什么需要清空对话

在多轮对话中，对话历史会不断累积，导致显存占用逐渐增加。长时间运行后，即使有torch.no_grad()优化，显存也可能被占满。

5.2 Streamlit清空对话实现

import gc # 初始化对话历史 if "messages" not in st.session_state: st.session_state.messages = [] # 侧边栏清空对话按钮 with st.sidebar: if st.button("🧹 清空对话", help="清空对话历史并释放显存"): # 清空对话历史 st.session_state.messages = [] # 强制垃圾回收 gc.collect() # 清空GPU缓存 if torch.cuda.is_available(): torch.cuda.empty_cache() st.success("对话已清空，显存已释放") # 显示对话历史 for message in st.session_state.messages: with st.chat_message(message["role"]): st.markdown(message["content"]) # 用户输入 if prompt := st.chat_input("你好，我是Qwen2.5-1.5B，有什么可以帮你的？"): # 添加用户消息到历史 st.session_state.messages.append({"role": "user", "content": prompt}) # 显示用户消息 with st.chat_message("user"): st.markdown(prompt) # 生成回复 with st.chat_message("assistant"): with st.spinner("思考中..."): response = generate_response(model, tokenizer, prompt, st.session_state.messages) st.markdown(response) # 添加助手回复到历史 st.session_state.messages.append({"role": "assistant", "content": response})

5.3 清空机制详解

清空对话按钮执行三个关键操作：

清空对话历史：重置st.session_state.messages为空列表
强制垃圾回收：调用gc.collect()回收Python对象占用的内存
清空GPU缓存：调用torch.cuda.empty_cache()释放GPU显存

6. 完整优化代码示例

以下是结合两种优化策略的完整代码：

import torch import gc import streamlit as st from transformers import AutoModelForCausalLM, AutoTokenizer # 模型路径 MODEL_PATH = "/root/qwen1.5b" @st.cache_resource def load_model(): model = AutoModelForCausalLM.from_pretrained( MODEL_PATH, device_map="auto", torch_dtype="auto", trust_remote_code=True ) tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True) return model, tokenizer def generate_response(model, tokenizer, prompt, chat_history=[]): messages = chat_history + [{"role": "user", "content": prompt}] text = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) inputs = tokenizer(text, return_tensors="pt").to(model.device) # 使用torch.no_grad()节省显存 with torch.no_grad(): outputs = model.generate( inputs.input_ids, max_new_tokens=1024, temperature=0.7, top_p=0.9, do_sample=True ) response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True) return response # 页面设置 st.set_page_config(page_title="Qwen2.5-1.5B 智能助手", page_icon="🤖") # 加载模型 with st.spinner("🚀 正在加载模型，请稍候..."): model, tokenizer = load_model() # 侧边栏 with st.sidebar: st.title("设置") if st.button("🧹 清空对话", help="清空对话历史并释放显存"): st.session_state.messages = [] gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() st.success("对话已清空，显存已释放") # 初始化对话历史 if "messages" not in st.session_state: st.session_state.messages = [] # 显示对话历史 for message in st.session_state.messages: with st.chat_message(message["role"]): st.markdown(message["content"]) # 用户输入 if prompt := st.chat_input("你好，我是Qwen2.5-1.5B，有什么可以帮你的？"): st.session_state.messages.append({"role": "user", "content": prompt}) with st.chat_message("user"): st.markdown(prompt) with st.chat_message("assistant"): with st.spinner("思考中..."): response = generate_response(model, tokenizer, prompt, st.session_state.messages) st.markdown(response) st.session_state.messages.append({"role": "assistant", "content": response})