Hugging Face 分享 gpt-oss 在 transformers 中可用的优化技巧

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 的文章非常出色：A Visual Guide to Quantization。

什么是 MXFP4


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

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

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

这种分块方案让 MXFP4 在使用非常少位数的同时保持范围。在实践中，当启用 MXFP4 时，GPT-OSS 20B 约占 16 GB 的 VRAM，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 需要能够在 GEMM 和融合操作期间理解 32 元素块及其缩放因子的内核。这正是 来自 Hub 的内核 再次发挥作用的地方。当加载需要这些内核的模型时，transformers 会自动从社区仓库拉取可感知 MXFP4 的 Triton 内核。该仓库将出现在你的本地缓存中，并会在前向传播期间使用。对于 MXFP4 内核，不再像以前那样需要使用 use_kernels=True 参数，它在 transformers 中已被设为默认值。

在使用与 triton MXFP4 内核兼容的 GPU 运行 gpt-oss-20b 后，使用 HF Mirror 缓存 CLI 快速进行完整性检查：

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]:]))

如果你没有运行上述内容所需的基础设施，你可以直接在我们的 GPU 上使用 HF Mirror 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://hf-mirror.com/datasets/ariG23498/distributed/raw/main/tp_gpt_oss.py && \
  torchrun --nproc-per-node=4 tp_gpt_oss.py"

hf jobs 面向所有 HF Mirror PRO 和企业版用户开放。

在底层，tp_plan="auto" 为每一层选择一个预定义的分片方案，并连接必要的集合通信。如果你想验证实际发生了哪些分片，可以用 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://hf-mirror.com/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‑cache 内存占用，适用于采用滑动或混合注意力的模型（如 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 展示了使用带滑动窗口注意力的 Dynamic KV Cache 为我们带来了多大的差异。


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

连续批处理与分页注意力

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


图 7：自回归 token 生成

让我们看看当我们传入一批（batch）输入时生成过程是什么样子。在图 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" 或提供自己的 device map，你的模型现在会自动加载得更快。如果你在 Tensor Parallel（tp_plan="auto"）和 torchrun 下运行，你还会受益于使多 GPU 加载更智能的配套更改。

结论

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

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

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