Transformers 集成 OpenAI GPT-OSS 新特性

OpenAI 最近发布了 GPT-OSS 系列模型。这些模型采用了 MXFP4 量化、高效内核、全新聊天格式等新颖技术。为了让开发者能通过 transformers 库使用 gpt-oss，我们对库进行了重大升级。这些更新让模型的加载、运行和微调都变得非常高效。

这篇博客将深入探讨所有升级内容，以及它们如何成为 transformers 工具包的一部分，让其他模型（包括现有和未来的）也能从中受益。在 transformers 中提供新方法的清晰实现，也有助于社区快速理解和采用它们。像 MLX、llama.cpp 或 vLLM 这样的框架，也可以参考 transformers 的代码来构建自己的实现。

本次发布，我们主要做了以下工作：

零构建内核，可从 Hub 下载
MXFP4 量化
张量并行
专家并行
动态滑动窗口层与缓存
连续批处理与分页注意力
更快加载大模型

最棒的是：这些特性大多能在 transformers 中的所有主流模型上运行！

零构建内核，可从 Hub 下载

内核（Kernel）是运行在加速器上的专用、紧凑的程序，用于执行矩阵乘法、激活函数或归一化等任务。在 PyTorch 的即时执行模式（eager mode）下，操作会顺序触发单个内核，这种方式简单直接，但可能带来额外的内存传输和启动开销。PyTorch 2.0 的 torch.compile 配合 TorchInductor 等后端，通过自动融合和优化内核来解决这个问题，能带来 2–10 倍的性能提升。

此外，社区还为频繁出现的操作组合（而不仅仅是像 matmul 这样的单个 PyTorch 操作）创建了自定义内核。例如，Flash Attention 就是为了优化定义 Transformer 架构的关键注意力块而创建的，它存在于包括大多数大语言模型在内的许多模型中。通过精心地将所有注意力操作组合在单个内核中，可以最大限度地减少内存传输、降低内存使用并实现加速。

问题在于，所有这些不同的内核都分散在不同的库中，如果把它们都加到 transformers 库里，会造成依赖膨胀。而且，这些内核不仅仅是 Python 代码，它们由底层 CUDA 代码组成，通过 C++ 粘合在一起，并通过 Python 层暴露出来。这意味着它们必须在目标系统上编译，而这又需要每个内核库所需的任何构建系统。

内核包通过从 Hub 下载预构建的、受支持内核的二进制文件来解决这个问题。你只需指定要使用的内核，kernels 就会查找与你的系统兼容的版本，并在首次使用时下载它。

GPT-OSS 的自定义内核

GPT-OSS 是一个专家混合（MoE）模型，它大量使用了来自 Hub 的内核。它利用了多个自定义内核：

Liger RMSNorm，通过 @use_kernel_forward_from_hub("RMSNorm") 使用。
Megablocks MoE 内核：通过 @use_kernel_forward_from_hub("MegaBlocksMoeMLP") 使用。
支持注意力汇聚（attention sinks）的 Flash Attention 3。
MXFP4 triton 内核（稍后介绍）。

我们来看看前两个。

在幕后，这些装饰器（1 和 2）只是指向社区贡献的内核。例如，RMSNorm 来自 liger_kernels，而 MegaBlocksMoeMLP 内核来自 megablocks。根据你的设备（CUDA 或 ROCm）以及你是在训练还是运行推理，系统会自动拉取正确的内核。

这种设计既具体又通用：Liger RMSNorm 内核已经在多个模型中复用，MoE 内核也可以应用于未来的 MoE 模型。

因为 kernels 从 Hub 拉取代码，所以你需要通过在你的模型实例化中传递 use_kernels=True 来选择启用此功能，如下所示。我们在示例中启用了 INFO 级别的日志记录，以便你轻松验证是否正在使用可下载的内核。

这些内核与 mxfp4 不兼容，因此如果你使用它们，推理将以 bfloat16 精度进行。请根据你的项目需求，在内存和吞吐量方面进行基准测试，以找到最佳组合！

python

from transformers import AutoTokenizer, AutoModelForCausalLM

import logging
logging.basicConfig(level=logging.INFO)

model_id = "openai/gpt-oss-20b"
tokenizer = AutoTokenizer.from_pretrained(model_id)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype="auto",
    device_map="auto",
    use_kernels=True,
)

运行一个简单的生成任务，你会看到类似这样的日志信息：

code

INFO:root:Using layer `LigerRMSNorm` from repo `kernels-community/liger_kernels`
INFO:root:Using layer `MegaBlocksMoeMLP` from repo `kernels-community/megablocks`

图 1 显示，在我们测试的系统中，这些内核在较大的批处理大小时表现最佳。我们始终建议尽可能接近你的生产条件，对任何与性能相关的更改进行基准测试。

图 1：自定义内核的基准测试结果

你可以在这里探索和试用基准测试脚本。

Flash Attention 3

OpenAI 的 gpt-oss 模型使用了注意力汇聚，这提高了模型质量并有助于使用更长的上下文。vLLM 团队将这一特性添加到了最新版本的 Flash Attention（Flash Attention 3）中，生成的自定义内核已在 Hub 上提供。目前，该内核与 Hopper 架构兼容。如果你有该架构的硬件，可以通过以下方式启用它：

python

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype="auto",
    device_map="auto",
    # 带注意力汇聚的 Flash Attention
    attn_implementation="kernels-community/vllm-flash-attn3",
)

MXFP4 量化

大语言模型非常消耗内存。量化通过以较低精度格式存储权重（有时也包括激活值）来减少内存占用。作为参考，FP32 每个数字使用 32 位，BF16 使用 16 位。通过减少位宽，我们牺牲一些精度来换取更小的模型和更快的内存移动。

如果你想了解量化权衡的直观入门，Maarten Grootendorst 的文章非常出色：《量化的视觉指南》。

什么是 MXFP4

图 2：MXFP4 格式中使用的 E2M1 格式

MXFP4 是一种 4 位浮点格式，采用 E2M1 布局：1 个符号位、2 个指数位和 1 个尾数位，如图 2 所示。单看 E2M1 本身非常粗糙。MXFP4 通过分块缩放来补偿：

向量被分成 32 个元素一组。
每个块存储一个共享的缩放因子，用于在反量化时恢复动态范围。
在每个块内部，4 位值代表相对于该缩放因子的数字。

这种分块方案让 MXFP4 在只使用很少位数的情况下保持范围。实际上，当启用 MXFP4 时，GPT-OSS 20B 大约占用 16 GB 的显存，GPT-OSS 120B 大约占用 80 GB，这决定了模型是“无法加载”还是“可以在单个 GPU 上运行”。关键在于，矩阵乘法现在必须考虑块缩放因子。要在规模上高效地做到这一点，需要专用的内核。

MXFP4 在 `transformers` 中

transformers 现在原生支持 MXFP4，利用优化的 triton（MXFP4）内核来提升性能。这建立在之前讨论的社区驱动的内核分发之上，利用 Hub 上的预编译内核来简化部署。

关键实现细节：

量化器逻辑：位于 MXFP4 量化器文件中，处理 MXFP4 的核心量化过程。
集成钩子：MXFP4 集成文件使得 MXFP4 能在 transformers 框架内无缝使用。

要检查模型是否支持 MXFP4，可以查看其配置：

python

from transformers import GptOssConfig

model_id = "openai/gpt-oss-120b"
cfg = GptOssConfig.from_pretrained(model_id)
print(cfg.quantization_config)

# 示例输出：
# {
#   'modules_to_not_convert': [
#     'model.layers.*.self_attn',
#     'model.layers.*.mlp.router',
#     'model.embed_tokens',
#     'lm_head'
#   ],
#   'quant_method': 'mxfp4'
# }

如果存在 'quant_method': 'mxfp4'，模型将在支持时自动使用带有 Triton 内核的 MXFP4 路径。

多亏了这个拉取请求，你现在可以微调 gpt-oss 模型，并直接以 MXFP4 格式保存到 Hub，简化了具有优化性能的部署流程。

运行要求与降级方案

想在 GPU 上跑 MXFP4 量化，你需要满足两个条件：

装好 accelerate、kernels 和 triton>=3.4。注意 PyTorch 2.8 已经自带 triton 3.4，所以如果你在用 PyTorch 2.7，才需要手动安装 triton。
一块计算能力 ≥ 7.5 的 NVIDIA GPU。这个要求其实不高，从 Tesla 架构的卡就支持了。这意味着你可以在 Google Colab 和 Kaggle 的免费层上跑 gpt-oss-20b，很多消费级显卡也够用。

如果条件不满足，transformers 会自动降级到更高精度的路径（默认用 bfloat16），但这需要大约 4 倍于 MXFP4 的内存。

这个代码片段会在 CUDA 上加载两次 GPT-OSS：一次用 Mxfp4Config(dequantize=True)（内存消耗大），一次用默认的量化路径（内存高效）。图 3 展示了每次加载后使用的显存，让你直观看到节省了多少。

图 3：量化模型与非量化模型的内存需求对比

MXFP4 专用内核

高效的 MXFP4 运算需要能理解 32 元素块及其缩放因子的内核，尤其是在做 GEMM 和融合操作时。这时候又得靠 Hub 上的内核了。当你加载一个需要 MXFP4 的模型时，transformers 会自动从社区仓库拉取支持 MXFP4 的 Triton 内核。这个仓库会出现在你的本地缓存里，并在前向传播时使用。对于 MXFP4 内核，你不需要像之前那样设置 use_kernels=True 参数，它在 transformers 里默认就是开启的。

用 Hugging Face 缓存 CLI 快速检查一下，在兼容 Triton MXFP4 内核的 GPU 上运行 gpt-oss-20b 后：

code

hf cache scan

示例输出：

code

REPO ID                          REPO TYPE SIZE ON DISK
-------------------------------- --------- ------------
kernels-community/triton_kernels model           536.2K
openai/gpt-oss-20b               model            13.8G

这表明 MXFP4 内核已经被拉取，可以用于执行了。

跑个基准测试看看 MXFP4 内核表现如何。从图 4 可以看到，对于更大的批次，MXFP4 内核的表现甚至比定制的 MoE 和 RMSNorm 内核还要好。

图 4：MXFP4 内核基准测试

你可以在这里探索和试用这个基准测试脚本。

张量并行

图 5：张量并行（Tensor Parallelism）原理说明。

张量并行（TP）将层内的张量拆分到多个 GPU 上（如图 5 所示）。每个 GPU 并行计算自己分到的部分，然后通过 all-gather 或 all-reduce 操作收集部分结果。这降低了单 GPU 的内存占用，并让所有 GPU 同时处理同一层，从而在序列长度或批次大小增长时提升吞吐量。TP 通信密集，通常在单机内通过高速节点内链路效果最好。

这在 `transformers` 里怎么用

transformers 直接在 from_pretrained 里实现了 TP。你可以从预定义的计划开始：

code

# 运行命令：torchrun --nproc-per-node 4 tp_gpt_oss.py
import torch
from transformers import PreTrainedTokenizerFast, GptOssForCausalLM

model_id = "openai/gpt-oss-120b"
tokenizer = PreTrainedTokenizerFast.from_pretrained(model_id)
model = GptOssForCausalLM.from_pretrained(
    model_id,
    tp_plan="auto", # 内置 TP 支持
    dtype="auto",
).eval()

messages = [
    {"role": "system", "content": "Be concise."},
    {"role": "user", "content": "Explain KV caching briefly."},
]
inputs = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    return_tensors="pt",
    return_dict=True,
    reasoning_effort="low",
).to(model.device)

with torch.inference_mode():
    generations = model.generate(**inputs, max_new_tokens=128)

print(tokenizer.decode(generations[0][inputs["input_ids"].shape[-1]:]))

如果你没有运行上述代码的基础设施，可以直接用 Hugging Face Jobs 在我们的 GPU 上启动一个进程：

code

hf jobs run --detach --flavor l4x4 ghcr.io/astral-sh/uv:debian /bin/bash -c \
  "uv venv .venv --python 3.12 && \
  source .venv/bin/activate && \
  uv pip install --upgrade torch numpy transformers accelerate triton kernels && \
  wget https://huggingface.co/datasets/ariG23498/distributed/raw/main/tp_gpt_oss.py && \
  torchrun --nproc-per-node=4 tp_gpt_oss.py"

hf jobs 对所有 Hugging Face PRO 和企业用户开放。

在底层，tp_plan="auto" 会为每一层选择一个预定义的分片方案，并配置必要的集合通信。如果你想确认正在分片的是什么，可以用 print(model._tp_plan) 查看当前激活的计划。

什么时候该用 TP

当模型太大，单 GPU 放不下，并且你希望获得并行计算能力，而不仅仅是内存分布时，就该用 TP 了。TP 往往能随着 GPU 数量的增加而提升吞吐量，尤其是在处理长序列或大批次时。

如果你好奇 TP 和 device_map="auto"（内存分布）有什么区别，这个简短的 Stack Overflow 回答解释了两者的区别以及各自的适用场景。

想了解更多关于 TP 的内容，这两个是必读资源：

transformers 指南：张量并行、支持的模型、计划和扩展点。
超大规模模型训练手册：TP 的背景知识及其与其他并行模式的关系。

专家并行

专家并行（EP）将 MoE 层内的专家 拆分到多个 GPU 上。每个 token 会被路由到一个或少数几个专家，因此只有这些专家会执行其前馈传播。由于专家是独立的 MLP，我们可以将不同的专家放在不同的计算节点上，只交换被路由 token 的隐藏状态。这让每个节点上的矩阵乘法保持完整，并用路由和集合通信代替了张量切片。

用 torchrun 启动多个进程来运行。EP 通过分布式配置启用，在 transformers 里对 GPT-OSS 的 MoE 层是开箱即用的。

code

# 运行命令：torchrun --nproc-per-node 4 ep_gpt_oss.py
import torch
from transformers import PreTrainedTokenizerFast, GptOssForCausalLM
from transformers.distributed import DistributedConfig

model_id = "openai/gpt-oss-120b"
tokenizer = PreTrainedTokenizerFast.from_pretrained(model_id)
model = GptOssForCausalLM.from_pretrained(
    model_id,
    distributed_config=DistributedConfig(enable_expert_parallel=True), # 启用 EP
    dtype="auto",
).eval()

messages = [
    {"role": "system", "content": "Be concise."},
    {"role": "user", "content": "Explain KV caching briefly."},
]
inputs = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    return_tensors="pt",
    return_dict=True,
    reasoning_effort="low",
).to(model.device)

with torch.inference_mode():
    generations = model.generate(**inputs, max_new_tokens=128)

print(tokenizer.decode(generations[0][inputs["input_ids"].shape[-1]:]))

用 hf jobs 运行的命令如下：

code

hf jobs run --detach --flavor l4x4 ghcr.io/astral-sh/uv:debian /bin/bash -c \
  "uv venv .venv --python 3.12 && \
  source .venv/bin/activate && \
  uv pip install --upgrade torch numpy transformers accelerate triton kernels && \
  wget https://huggingface.co/datasets/ariG23498/distributed/raw/main/ep_gpt_oss.py && \
  torchrun --nproc-per-node=4 ep_gpt_oss.py"

当你启用专家并行时，张量并行也会被激活。这意味着你可以同时享受两者的优势！

动态滑动窗口层与缓存

许多最近的 LLM 使用滑动窗口注意力，或者滑动与全局注意力层的组合，作为节省内存、减少那些随序列长度呈二次方增长的昂贵矩阵乘法的手段。然而，transformers 里原来的动态 KV 缓存实现，会继续根据序列长度分配空间，而不考虑各个注意力层。你当然可以通过编译（意味着固定形状）来优化内存，但那完全是另一个场景了。

现在 transformers 有了一个 DynamicSlidingWindowLayer 和一个配置感知的 DynamicCache。如果模型配置声明了滑动窗口或混合注意力（同时使用滑动和全局注意力层），那么对于滑动层，缓存在超过窗口大小后就会停止增长。如果你不传递配置，行为则保持原样（KV 会随着序列长度增长而持续增长）。

对于只使用滑动窗口层的模型，比如 Mistral 7B，当序列达到窗口大小（这里是 4096）时，缓存内存就会停止增长。这很合理，因为滑动层本来也看不到 4K 个 token 之前的内容。

OpenAI gpt-oss 交替使用滑动和全局注意力层，这导致总的 KV 缓存内存随着序列长度增加而减半，我们马上会看到。这为我们带来了：

大幅降低的 KV 缓存内存，适用于具有滑动或混合注意力的模型（例如 GPT-OSS）。一旦达到窗口大小（例如，Mistral 是 4K；GPT-OSS 的滑动层是 128），缓存增长就会趋于平稳，而不是随着生成的总 token 数线性增长。（GitHub, Transformers）
长提示/长生成的速度/延迟优势：更小的 KV 张量意味着更轻量的注意力读写操作和更少的内存带宽压力，尤其是在窗口被命中之后。（这正是滑动窗口/混合 LLM 背后的核心动机。）（AI21, vLLM Blog）

怎么用

优化后的缓存是默认设置的，这意味着你不需要对现有代码做任何改动。如果你想显式地创建 DynamicCache，可以这样做：

code

from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache

model_id = "openai/gpt-oss-20b"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype="auto",
    device_map="auto",
).eval()

messages = [
    {"role": "system", "content": "Always respond in riddles"},
    {"role": "user", "content": "What is the weather like in Madrid?"},
]

inputs = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    return_tensors="pt",
    return_dict=True,
    reasoning_effort="low",
).to(model.device)

cache = DynamicCache(config=model.config) # 使用模型的配置创建缓存

generated = model.generate(
    **inputs,
    max_new_tokens=500,
    past_key_values=cache
)
print(tokenizer.decode(generated[0][inputs["input_ids"].shape[-1]:]))

图 6 展示了使用动态 KV 缓存配合滑动窗口注意力能带来多大的差异。

图 6：滑动窗口注意力下动态缓存的内存分析

连续批处理与分页注意力

典型的自回归生成过程如图 7所示：先输入预填充（Prefill）token，模型逐个预测新 token，直到生成结束序列（EOS）token。

图 7：自回归 token 生成

当我们传入一批输入时，生成过程会变成什么样？图 8中可以看到，有些序列的生成比其他的更早结束。这种长度不匹配会导致 GPU 利用率不足。

图 8：序列的静态批处理（Static Batching）

这种批处理方式称为静态批处理。虽然简单易懂，但天生就有低效问题——必须等整批句子全部生成完毕，才能处理下一批。

为了解决这个问题，我们可以用动态批处理（也叫连续批处理）。不再等待所有生成完成，而是把新请求调度到已完成的生成任务上。这样，只要批次中有一个生成任务结束，就能立刻用下一个请求填充这个空位。整个过程如图 9所示。

图 9：序列的连续批处理（Continuous Batching）

Transformers 库通过 generate_batch API 支持连续批处理。这个 API 不是为生产级模型服务设计的（vLLM、SGLang 等框架更适合），但对评估和实验很有帮助。这里有一个示例脚本，可以在 Qwen/Qwen3-4B-Instruct-2507 模型上运行完整的连续批处理流程。

我们还用 100 个样本对连续批处理和静态批处理做了基准测试。从图 9 可以看出，连续批处理比静态批处理快不少。

图 9：连续批处理 vs 静态批处理（token/秒）

你可以在这里体验基准测试：静态批处理、连续批处理

更快加载大模型

当你把一个大模型加载到 GPU 时，PyTorch 需要为每一层的权重预留 GPU 内存。每个这样的请求（每层一次）都需要时间，对于数十亿参数的模型来说，可能意味着成千上万次细小的内存分配，加起来会让模型准备好之前等待很久。

与其每次都向 GPU 申请新内存，不如一次性申请一大块内存，然后快速从中分配切片。PyTorch 的内存分配器就能做到这一点。但有个前提：分配器只有在拿到一些内存后才会变快。如果你不先“备好库存”，还是得跑很多趟慢速的“市场采购”。

这个 PR（🎉 #36380）教会了 transformers 在开始复制模型权重之前先备好库存。具体做法是：

查看 device_map（确定每层放在哪个设备上）
在每个 GPU 上预分配足够大的内存块
然后，当各层被复制进来时，它们就能整齐地放入这个预留空间

你不需要修改现有代码，因为这已经是 transformers 的默认行为。如果你使用 device_map="auto" 或提供自己的设备映射，你的模型现在会自动加载得更快。如果你用张量并行（tp_plan="auto"）和 torchrun，配套的改动也能让多 GPU 加载更智能。

结语

transformers 发展迅速，始终以社区为先。这个库随着领域的发展而演进，因为贡献者在开放环境中塑造它。为新模型添加的组件会成为工具包的一部分，并在未来的集成中复用。

这种速度使得像 GPT-OSS 系列这样的“零日集成”成为可能。随着技术栈越来越以 PyTorch 为先，它精简了冗余，并加倍投入实践中重要的 PyTorch 路径。结果是一个更清晰的核心，通过社区内核、量化和并行计划解锁新能力，同时标准化模型定义，让 transformers 支持的架构成为参考，并扩展到更广泛的生态系统中。

这篇文章只是我们反复迭代过程中的一次快照，方向始终如一：满足社区需求。要了解 transformers 的最新进展，请查看文档和发布说明。也请继续分享你的反馈，并在 transformers 中发布你的模型，让社区受益 🤗

延伸阅读

如果你想深入了解特定主题，这里有一些值得访问的链接：

Transformers 集成 OpenAI GPT-OSS 新特性

零构建内核，可从 Hub 下载

GPT-OSS 的自定义内核

Flash Attention 3

MXFP4 量化

什么是 MXFP4

MXFP4 在 `transformers` 中

运行要求与降级方案

MXFP4 专用内核

张量并行

这在 `transformers` 里怎么用

什么时候该用 TP

专家并行

动态滑动窗口层与缓存

怎么用

连续批处理与分页注意力

更快加载大模型

结语

延伸阅读

概念速查

相关文章

pgEdge 推出 Postgres MCP 服务器，AI 智能体连接数据库新方案

Gemini API 新增 Flex 与 Priority 推理层级

评论

Transformers 集成 OpenAI GPT-OSS 新特性

零构建内核，可从 Hub 下载

GPT-OSS 的自定义内核

Flash Attention 3

MXFP4 量化

什么是 MXFP4

MXFP4 在 transformers 中

运行要求与降级方案

MXFP4 专用内核

张量并行

这在 transformers 里怎么用

什么时候该用 TP

专家并行

动态滑动窗口层与缓存

怎么用

连续批处理与分页注意力

更快加载大模型

结语

延伸阅读

概念速查

相关文章

pgEdge 推出 Postgres MCP 服务器，AI 智能体连接数据库新方案

Gemini API 新增 Flex 与 Priority 推理层级

评论

MXFP4 在 `transformers` 中

这在 `transformers` 里怎么用