5090显卡+Triton，轻松玩转GPT-OSS-20B！

摘要：gpt-oss 提供了triton implementation供研究用，你可以使用单卡H100，也就是80 GB显存的Hopper or Blackwell 去跑gpt-oss-120b，其使用了 支持 MXFP4的 Triton MoE 算子来减少显存占用。我用5090 32G显存跑了20B的小模型，在此记录下，显存的使用在17G左右，吐出token非常快。

一、安装

项目地址：openai/gpt-oss

我的Python 版本是3.12，建议使用高版本。Torch是 2.7.0+cu128，Triton是自行编译安装的3.4.0+git450dabd3(Aug 6, 2025)，显卡是5090 32g。

1、模型下载

首先需要解决的是网络问题，我们需要把模型给下载下来。先安装huggingface_hub包，并下载模型。

pip install -U huggingface_hub
# or  ~/.local/bin/huggingface-cli, if warning `huggingface-cli: command not found`
huggingface-cli download openai/gpt-oss-20b --include "original/*" --local-dir gpt-oss-20b/

2、安装torch(可选)

在安装Triton前请确保torch已经安装好，安装torch会顺便帮你安装Triton，顺序错了需要重新安装triton。

pip install torch

3、源码安装Triton

clone Triton的源码，安装就好了，另外还需要安装triton_kernels。Triton会下载LLVM，1G以上的安装包，需要你解决网络问题。Triton目前的代码会产生triton.runtime.errors.OutOfResources: out of resource: shared memory 的问题，所以5090 你建议使用我的项目，我改小了 block_k。

# You need to install triton from source to use the triton implementation
# git clone https://github.com/triton-lang/triton
# 5090 need clone my repo
git clone https://github.com/OpenMLIR/triton
cd triton/
pip install -r python/requirements.txt
pip install -e . --verbose --no-build-isolation
pip install -e python/triton_kernels

4、源码安装gpt-oss

git clone https://github.com/openai/gpt-oss
cd gpt-oss

# Install the gpt-oss triton implementation
pip install -e .[triton]

二、运行

我们可以把gpt-oss-20b移动到gpt-oss文件夹下，也可以自己修改相关路径。示例命令是一个How are you?做为prompt的inference。

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
python -m gpt_oss.generate --backend triton gpt-oss-20b/original/

如果你torch.OutOfMemoryError了，请扩展allocator逻辑。

另外你可以简单改造gpt_oss/generate.py，使其作为你的独享大模型。比如f"Generated token: {repr(decoded_token)}, logprob: {logprob}"的输出实际我们使用print(decoded_token, end="")就好了，你也可以通过命令行参数来修改prompt及temperature，不清楚怎么调可以问大模型。

三、F&Q

1、可以跑在Ampere上运行吗

很遗憾不能，有kernel使用了mul.bf16x2，会报错Feature 'mul.bf16x2' requires .target sm_90 or higher，我在尝试将Triton implementation 扩展到更多的卡上。

2、可以不源码安装Triton吗

依赖main分支上的Kernel更新和部分修改。

我在尝试去掉main分支 Triton依赖，将编译好的cubin使用我开发的triton_runner 去简化Trirton源码安装这一流程。

3、IndexKernel

Triton implementation 还调用了torch的aten算子，可能会遇到如下错误

/pytorch/aten/src/ATen/native/cuda/IndexKernel.cu:175: operator(): block: [0,0,0], thread: [96,0,0] Assertion `idx >= 0 && idx < self_dim_size && "index_copy_(): index out of bounds"` failed.

4、AI模型发疯

在Sampling temperature为默认值0.0时，且不设置Limit on the number of token，会重复推理出以下内容，修改任意一个即可。也有可能是算子的问题。

It looks like your message might have been cut off. Could you let me know what you need help with? Whether it's more details on sorting algorithms, another programming question, or anything else, I'm here to help!

Sure! If you have any more questions or need further assistance, feel free to ask.

It seems like you might be looking for more information or a different topic. If you have a specific question or need help with a particular subject, just let me know!

It looks like you might have a question or need help with something. Feel free to ask! 

5、显存占用和功耗

以下是nvidia-smi的信息