来自OpenAI gpt-oss的技巧，你🫵在transformers中也可以使用

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

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

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

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

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

零构建内核，可从Hub下载

内核(kernel)是一种专门的紧凑程序，可在加速器上运行，用于执行矩阵乘法、激活或归一化等任务。在 Eager PyTorch 中，操作会依次触发单个内核，这很直接但可能会产生额外的内存传输和启动开销。PyTorch 2.0 的 torch.compile 和 TorchInductor 等后端通过自动融合和优化内核解决了这个问题，从而实现了 2 到 10 倍 的性能提升。

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

问题在于，所有这些不同的内核都在单独的库中可用，如果将它们添加到 Transformers 库中，会导致依赖项膨胀。此外，这些内核不仅仅是 Python 代码，它们还包含底层的 cuda 代码，这些代码与 C++ 粘合在一起，并通过 Python 层公开。这意味着它们必须在目标系统中进行编译，这反过来需要每个内核库所需的构建系统。

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

GPT-OSS的自定义内核

GPT-OSS，是一个混合专家 (MoE) 模型，大量使用了 Hub 中的内核。它利用了以下几个自定义内核：

1. Liger RMSNorm，用作@use_kernel_forward_from_hub("RMSNorm")
2. Megablocks MoE 内核：@use_kernel_forward_from_hub("MegaBlocksMoeMLP")
3. Flash Attention 3，支持注意力接收器。
4. 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显示，在我们测试的系统中，这些内核在较大的批次大小下效果最好。我们始终建议, 尽可能接近您的生产条件对任何与性能相关的变化进行基准测试。

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

Flash Attention 3

OpenAI gpt-oss模型使用 _attention sinks_，这提高了质量并促进了更长上下文的使用。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

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集成文件在transformers框架内实现MXFP4的无缝使用。

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

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，你需要：

1. 安装accelerate、kernels和triton>=3.4。注意Pytorch 2.8已经带有triton 3.4，所以如果使用Pytorch 2.7，你只需要手动安装triton。
2. 计算能力≥ 7.5的NVIDIA GPU。这一直追溯到Tesla，所以你可以在Google Colab和Kaggle的免费层以及许多消费级GPU上运行gpt-oss-20b。

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

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

MXFP4的内核

高效的 MXFP4 需要能够在 GEMM 和融合操作期间理解 32 元素块及其尺度的内核。这时，来自 Hub 的内核 再次发挥作用。当您加载需要 MXFP4 的模型时，transformers 会自动从社区代码库中提取支持 MXFP4 的 Triton 内核。该代码库将出现在您的本地缓存中，并在前向传播过程中使用。对于 MXFP4 内核，无需像以前一样使用 use_kernels=True 参数，它在 transformers 中已设置为默认值。

在与 triton MXFP4 内核兼容的 GPU 上运行gpt-oss-20b后，使用Hugging Face缓存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 内核表现更佳。

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

张量并行

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

这在transformers中实现了什么

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

# 运行命令：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上生成一个进程！

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"会为每一层选择一个预定义的分片方案，并连接必要的集合操作。如果您想验证分片的具体内容，可以使用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，我们可以将不同的专家置于不同的等级上，并仅交换路由 token 的隐藏状态。这可以保持每个等级上的矩阵乘法不变，并用路由和集合操作取代张量切片。

使用 torchrun 以多进程方式运行。EP 通过分布式配置启用，并可在 Transformer 中与 GPT-OSS MoE 层开箱即用。

# 运行命令：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运行的方法

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 使用 滑动窗口 注意力机制（或滑动和全局注意力层的组合）来节省内存，并减少那些随着序列长度增长而增长的昂贵二次矩阵乘法运算。然而，Transformer 中的动态键值缓存实现过去会根据序列长度继续分配空间，而不查看单个注意力层。您始终可以使用编译（即固定形状）来优化内存，但那完全是另一种情况。

transformers现在有一个DynamicSlidingWindowLayer和一个配置感知的DynamicCache。如果模型配置声明滑动窗口或混合注意力（同时使用滑动和全局注意力层），缓存在超过滑动层的窗口大小后会停止增长。如果不声明该配置，则行为将保持不变（随着序列长度的增长，键值对 (KV) 会保持满且不断增长）。

对于只使用滑动窗口层的模型，如Mistral 7B，当序列达到窗口大小（在本例中为 4096）时，缓存内存停止增长。这是合理的，因为滑动层无论如何都无法超过之前的 4K 个 token。

OpenAI gpt-oss在滑动和全局注意力层之间交替，这导致总KV缓存内存 减半 ，正如我们将看到的，随着序列长度的增加。 这为我们提供了：

• 对于具有滑动或混合注意力的模型（例如GPT-OSS），KV缓存内存要低得多。一旦达到窗口（例如，Mistral为4K；GPT-OSS滑动层为128），缓存增长就会趋于平稳，而不是随着总生成token数线性扩展。（GitHub，Transformers）
• 长提示/长生成的速度/延迟优势：较小的KV张量意味着更轻的注意力读/写操作，以及更小的内存带宽压力，特别是在达到窗口后。（这是滑动窗口/混合LLM背后的核心动机。）（AI21，vLLM博客）

如何使用它

优化的缓存是默认设置的，这意味着你不必对现有代码进行任何更改。如果你想显式创建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) # 使用模型配置创建缓存

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

图6展示了使用带有滑动窗口注意力的动态KV缓存对我们来说有多大的区别。

连续批处理和分页注意力

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

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

这种批处理序列的类型称为静态批处理。虽然这简单易懂，但它本质上带来了低效率。只有在每个句子完全生成后，我们才能继续下一批。

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

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

我们还对100个样本进行了连续批处理和静态批处理之间的基准测试。在图9中，我们注意到CB(Continuous Batching)比SB(Static Batching)快得多。

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

更快地加载更大的模型

当你将大型模型加载到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 中发布你的模型，供社区使用🤗

延伸阅读

如果您想进一步了解特定主题，可以访问以下链接列表：

1. Hugging Face GPT-OSS配方仓库
2. 欢迎GPT OSS：OpenAI 全新开源模型系列
3. OpenAI Cookbook：GPT-OSS主题
4. Transformers文档：基于多 GPU 的分布式推理
5. Matthew Carrigan 的X 推文: 关于GPT OSS创新
6. YouTube视频：OpenAI GPT OSS公告
7. Transformers PR #36380：在加速器上更快的模型加载
8. Transformers PR #36335：为张量并行更新from_pretrained
9. Transformers PR #40039：新的动态滑动窗口层和缓存
10. HAN Lab博客：注意力机制如何保持语言模型稳定