如何在低显存GPU上高效运行大型NLP Transformers模型

Shrink.media

Shrink.media是当今市场上最快、最直观、最智能的图像文件缩减工具

83

查看详情

# 安装transformers和accelerate库 !pip install -q transformers accelerate  # 安装AutoAWQ库。请注意，这里的版本可能需要根据您的CUDA环境进行调整。 # 对于Colab，有时需要指定一个与当前CUDA版本兼容的旧版本。 # 示例：针对CUDA 11.8环境 !pip install -q -U https://github.com/casper-hansen/AutoAWQ/releases/download/v0.1.6/autoawq-0.1.6+cu118-cp310-cp310-linux_x86_64.whl # 如果您遇到安装问题，请查阅AutoAWQ的官方文档或github仓库以获取最新兼容版本。

import torch from awq import AutoAWQForCausalLM from transformers import AutoTokenizer  # 指定量化模型的名称 model_name = 'TheBloke/neural-chat-7B-v3-1-AWQ'  # 使用AutoAWQ的from_quantized方法加载模型 # 这将加载一个已经过AWQ量化的模型，显著减少显存占用 model = AutoAWQForCausalLM.from_quantized(model_name)  # 加载对应的分词器 tokenizer = AutoTokenizer.from_pretrained(model_name)  print(f"模型 '{model_name}' 已成功加载。")

def generate_response(system_input, user_input):     """     根据系统和用户输入生成模型响应。     """     # 格式化输入提示，遵循模型预期的模板     prompt = f"### System:n{system_input}n### User:n{user_input}n### Assistant:n"      # 将提示词编码为张量，并确保其被移动到GPU设备上     # .cuda() 方法将张量从CPU移动到GPU     inputs = tokenizer.encode(prompt, return_tensors="pt", add_special_tokens=False).cuda()      # 使用模型生成响应     # max_length 控制生成文本的最大长度     # num_return_sequences 控制返回的序列数量     outputs = model.generate(inputs, max_length=1000, num_return_sequences=1)      # 解码生成的张量为可读文本     response = tokenizer.decode(outputs[0], skip_special_tokens=True)      # 提取并返回助手部分的响应     return response.split("### Assistant:n")[-1].strip()  # 示例用法 system_input = "You are a math expert assistant. Your mission is to help users understand and solve various math problems. You should provide step-by-step solutions, explain reasonings and give the correct answer." user_input = "calculate 100 + 520 + 60" response = generate_response(system_input, user_input) print("n--- 模型生成的响应 ---") print(response)  # 预期响应示例（模型实际输出可能略有不同，但逻辑应一致） """ To calculate the sum of 100, 520, and 60, we will follow these steps:  1. Add the first two numbers: 100 + 520 2. Add the result from step 1 to the third number: (100 + 520) + 60  Step 1: Add 100 and 520 100 + 520 = 620  Step 2: Add the result from step 1 to the third number (60) (620) + 60 = 680  So, the sum of 100, 520, and 60 is 680. """

大家都在看：