qidao123.com技术社区-IT企服评测·应用市场

标题: 多张GPU卡 [打印本页]

作者: 不到断气不罢休 时间: 2024-9-18 20:16
标题: 多张GPU卡
from transformers import pipeline
from accelerate import init_empty_weights, infer_auto_device_map
from transformers import AutoModelForCausalLM, AutoTokenizer
初始化加速器

from accelerate import Accelerator
accelerator = Accelerator()
加载模型和 tokenizer

tokenizer = AutoTokenizer.from_pretrained(“/app/model/”)
model = AutoModelForCausalLM.from_pretrained(“/app/model/”)
将模型分布到多张 GPU 上

device_map = infer_auto_device_map(model, max_memory={0: “16GB”, 1: “16GB”, 2: “16GB”, 3: “16GB”, 4: “16GB”, 5: “16GB”, 6: “16GB”, 7: “16GB”})
将模型加载到设备上

model = model.from_pretrained(“/app/model/”, device_map=device_map)
使用 pipeline

generator = pipeline(“text-generation”, model=model, tokenizer=tokenizer, device_map=device_map)
进行推理

output = generator(“Your input text”, max_length=30)
print(output)

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。

欢迎光临 qidao123.com技术社区-IT企服评测·应用市场 (https://dis.qidao123.com/)