Hugging Face 详解 transformers 为支持 OpenAI GPT-OSS 做出的多项性能升级

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

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

对于本次发布，我们致力于：

最棒的是：这些功能中的大多数应该适用于 transformers 中的所有主要模型！

零构建内核，可从 Hub 下载

内核是一种专用的紧凑程序，在加速器上运行以执行矩阵乘法、激活或归一化等任务。在 eager PyTorch 中，操作会依次触发各个内核，这很直接，但可能会带来额外的内存传输和启动开销。PyTorch 2.0 的 torch.compile 配合 TorchInductor 等后端，通过自动融合和优化内核来解决这个问题，带来 2–10× 的性能提升。

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

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

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

GPT-OSS 的自定义内核

GPT-OSS 是一个混合专家（MoE）模型，是 Hub 中内核的重度用户。它利用了多个自定义内核：

Liger RMSNorm，用作 @use_kernel_forward_from_hub("RMSNorm")`
Megablocks MoE 内核：@use_kernel_forward_from_hub("MegaBlocksMoeMLP")
Flash Attention 3，支持注意力汇。
MXFP4 triton 内核（稍后介绍）

让我们来看看前两个。

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

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

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

这些内核与 mxfp4 不兼容，因此如果你使用它们，推理将在 bfloat16 中进行。请对你的系统进行基准测试，以找到最适合你项目的内存和吞吐量组合！

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,
)

运行一次快速生成会产生如下日志消息

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 架构兼容。如果你有 Hopper 架构，可以这样启用它：

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype="auto",
    device_map="auto",
+    # Flash Attention with Sinks
+    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 上运行”之间的区别。问题在于，矩阵乘法现在必须遵守块缩放因子。要大规模高效地做到这一点，需要专门的内核。

`transformers` 中的 MXFP4

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

关键实现细节：

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

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

from transformers import GptOssConfig

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

# Example output:
# {
#   '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 的免费层以及许多消费级 GPU 上运行 gpt-oss-20b。

如果不满足这些限制，transformers 会回退到更高精度的路径（默认使用 bfloat16），这需要大约 MXFP4 的 4 倍内存。

代码片段在 CUDA 上加载 GPT-OSS 两次：一次使用 Mxfp4Config(dequantize=True)（内存密集），一次使用默认量化路径（内存高效）。图 3 显示了每次加载后使用的 VRAM 量，以便你可视化节省的内存。


图 3：量化模型和反量化模型的内存需求

MXFP4 的内核

高效的 MXFP4 需要在内核中理解 32 元素块及其缩放因子，在 GEMM 和融合操作期间。这就是 来自 Hub 的内核 再次发挥作用的地方。当你加载需要它们的模型时，transformers 会自动从社区仓库拉取感知 MXFP4 的 Triton 内核。该仓库将出现在你的本地缓存中，并在前向传递期间使用。对于 MXFP4 内核，不需要像以前那样使用 use_kernels=True 参数，它在 transformers 中默认设置。

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

hf cache scan

示例输出：

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：张量并行的解释。

张量并行 (TP) 将 层内的张量 拆分到多个 GPU 上（如 图 5 所示）。每个 GPU 并行地乘以其分片，然后使用 all-gather 或 all-reduce 操作收集部分结果。这减少了每个 GPU 的内存占用，并让所有 GPU 都在 同一层 上工作，从而随着序列长度或批次大小的增长提高吞吐量。TP 是通信密集型的，通常在 具有快速节点内链路的单台机器 上效果最佳。

这在 `transformers` 中实现了什么

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

# run with: 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", # built in TP support
    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 上启动一个进程！

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 & Enterprise 用户可用。

在底层，tp_plan="auto" 为每一层选择一个预定义的分片方案，并连接必要的 collectives。如果你想验证正在分片的内容，可以使用 print(model._tp_plan) 检查当前活动的计划。

何时使用 TP

当模型太大而无法放入单个 GPU，并且你想要并行计算，而不仅仅是内存放置时，请使用 TP。TP 往往会随着 GPU 数量的增加而扩展吞吐量，尤其是对于长序列或更大的批次。

如果你好奇 TP 与 device_map="auto"（内存放置）有何不同，这个简短的 Stack Overflow 回答解释了区别以及何时使用每种方法。

要了解有关 TP 的更多信息，这里有两个必读资源：

transformers 指南：张量并行、支持的模型、计划和扩展点。
Ultra-Scale Playbook：TP 的背景及其与其他并行模式的关系。

专家并行

专家并行（EP）将 MoE 层内的专家 分片到多个 GPU 上。每个 token 被路由到一个或几个专家，因此只有那些专家运行它们的前馈传递。由于专家是独立的 MLP，我们可以将不同的专家放在不同的 rank 上，并且只交换被路由 token 的隐藏状态。这保持了每个 rank 上的矩阵乘法完整，并用路由和集合通信取代了张量切片。

使用 torchrun 运行多个进程。EP 通过分布式配置启用，并且在 transformers 中开箱即用地支持 GPT-OSS MoE 层。

# run with: 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), # enabling 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 运行

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，可以这样做：

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) # create the cache with the model's 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所示。你输入预填充 token，模型逐个预测每个新 token，直到预测出 EOS（序列结束）token。


图 7：自回归 token 生成

让我们看看传入一批输入时生成过程是什么样子。在图 8中你会注意到，有些生成比其他生成更早结束。这种长度不匹配会导致 GPU 利用率不足。


图 8：序列的静态批处理

这种批处理序列的方式称为静态批处理。虽然它简单易懂，但本身存在效率低下的问题。只有等每个句子完全生成后，我们才能进入下一批。

为了绕过这个问题，我们使用动态批处理（也称为连续批处理）。我们不再等待所有生成完成，而是将新到达的请求调度到已完成的生成位置。这样，一旦一批中的某个生成完成，我们就用下一个请求预填充该批次。该过程如图 9所示。


图 9：序列的连续批处理

Transformers 通过 generate_batch API 支持连续批处理。这并非用于生产级模型服务——vLLM 和 SGLang 等框架在这方面表现出色——但对评估和实验非常有帮助。这里有一个示例脚本，可在 Qwen/Qwen3-4B-Instruct-2507 上端到端运行 CB。

我们还使用 100 个样本对连续批处理和静态批处理进行了基准测试。在图 9 中，我们注意到 CB 比 SB 快得多。


图 9：连续批处理与静态批处理的每秒 token 数对比

你可以在这里试用该基准测试：SB，CB

更快加载更大的模型

当你将大模型加载到 GPU 时，PyTorch 需要为每一层的权重预留 GPU 内存。这些（每层的）请求每一个都需要时间，而对于数十亿参数的模型，这可能意味着成千上万次微小的内存分配，累积起来就是模型就绪前的漫长等待。与其每次都向 GPU 申请新内存，不如一次性持有一大块内存，然后从中快速分配切片。

PyTorch 分配器恰好能做到这一点。问题在于，只有在你给它一些内存之后，分配器才会变快。如果你不先“把食品储藏室备满”，你仍然要多次缓慢地跑市场。这个 PR（🎉 #36380）教会了 transformers 在开始复制模型权重之前预先备满储藏室。

它：

查看 device_map（每一层将位于何处）。
在每个 GPU 上预分配一个足够大的块。
然后，当层被复制进来时，它们就整齐地放入这个预保留的空间。

你无需对现有代码做任何更改，因为这是 transformers 中的默认行为。如果你使用 device_map="auto" 或提供自己的设备映射，你的模型现在会自动加载得更快。如果你在 Tensor Parallel（tp_plan="auto"）和 torchrun 下运行，你也会受益于使多 GPU 加载更智能的配套更改。

结论

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

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

这篇文章是我们朝着同一方向反复迭代的一个过程的单次快照：服务社区的需求。要了解 transformers 的最新新增内容，请查看文档和发行说明。并且，请继续分享你的反馈，并在 transformers 中发布你的模型，供社区享用 🤗