A100用transformers推理gpt-oss

发布于:2025-08-10 ⋅ 阅读:(18) ⋅ 点赞:(0)

A100本地用transformers推理gpt-oss

GPT-oss试用
gpt-oss有两个原生配置是目前(2025-8-8)Ampere系列显卡不支持的,分别是默认的MXFP4量化,以及Flash-attn V3。tranformers推理是比较慢的,可以用于研究模型本身,如果是本地部署,最好用vllm,见三步完成,A100本地vLLM部署gpt-oss,并启动远程可访问api

下载模型参数到本地

国内用魔塔更快,下载到本地xxxx目录下:

pip install modelscope
modelscope download --model openai-mirror/gpt-oss-120b --local_dir xxxxxx

创建conda环境

conda create -n vllm python=3.12 -y

pip install -U transformers 
pip install -U accelerate 
pip install -U torch 
pip install -U kernels

注意到,这里跟官方给的安装步骤不太一样。官方安装多一个triton,里面包含了Triton kernels for MXFP4 compatibility,也就是gpt-oss专用的一个量化技术 MXFP4,这个量化在A100 GPU上是不被支持的。因此不能安装,安装反而会报错。

模型generate推理

相信大家一般用的最多的还是generate推理,这里给个python示例:

from transformers import AutoModelForCausalLM, AutoTokenizer
 
model_name = "/mnt/models/gpt-oss-120b" # /mnt/models/gpt-oss-20b
 
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto", # auto
    device_map="auto",
    attn_implementation="sdpa", # flash_attention_2 sdpa
)
 
messages = [
    {"role": "system", "content": "Reasoning: low"},
    {"role": "user", "content": "Explain what MXFP4 quantization is."},
]
 
inputs = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    return_tensors="pt",
    return_dict=True,
).to(model.device)
 
outputs = model.generate(
    **inputs,
    max_new_tokens=1024,
    temperature=0.7
)
 
print(tokenizer.decode(outputs[0]))

其中model_name就是下载的模型参数目录,attn_implementation这里,不能选择flash-attn。因为默认使用flash-attn V3,这不被a100支持,会报错Sinks are only supported in FlashAttention 3,用flash-attn 2也不行,即使是最新release的2.8.2(2025-8-8),也会报错一些参数无法识别。messages的第一行{"role": "system", "content": "Reasoning: low"}用来指定模型思考深度,gpt-oss有三个可选思考深度,但不能关闭。

最后,可以用transformers启动服务,进行终端命令行对话:

transformers serve 
transformers chat localhost:8000 --model-name-or-path /mnt/models/gpt-oss-120b

网站公告

今日签到

点亮在社区的每一天
去签到