在 Transformers 中通过 ExecuTorch 导出边缘推理模型：torch.export 静态缓存导出实战指南

ExecuTorch 是 PyTorch 生态面向移动端与边缘设备（可穿戴设备、嵌入式设备、微控制器）的端到端推理解决方案，其工作流的第一步是把 PyTorch 模型通过 torch.export 导出为可继续被后端委派、编译器变换与内存规划等流程优化降级（lowering）的中间产物。本文以 docs/source/en/main_classes/executorch.md 为主线，结合 🤗 Transformers 仓库中的集成源码 src/transformers/integrations/executorch.py 与其配套测试 tests/test_executorch.py，系统讲解：为什么边缘推理需要静态 KV Cache 导出、TorchExportableModuleWithStaticCache 这个「可导出封装模块」如何工作、如何用 convert_and_export_with_cache 一步导出 ExportedProgram，以及导出的前置配置约束与注意事项。读完本文，你将能够配置模型并产出与 ExecuTorch LLM Runner 接口兼容的导出产物。

ExecuTorch 与 torch.export：模型上边缘设备的第一步

ExecuTorch 是 PyTorch 官方生态的一部分，它为可穿戴设备、嵌入式设备和微控制器等移动与边缘设备提供端上（on-device）推理的端到端能力，核心关注点是可移植性（portability）、生产力（productivity）与性能（performance）。ExecuTorch 定义了清晰的入口点，用于执行针对模型、设备或具体使用场景的优化，例如后端委派（backend delegation）、用户自定义的编译器变换、**内存规划（memory planning）**等。

这些优化都发生在导出之后：准备让一个 PyTorch 模型在边缘设备上运行的第一步，就是通过 PyTorch 的 torch.export API 将模型导出。torch.export 产生的是一个图级别（graph-level） 的中间表示——torch.export.ExportedProgram，它剥离了 Python 运行时控制流，具备静态、可序列化、可继续降级的特性，这正是 ExecuTorch 一系列 AOT（ahead-of-time）优化得以展开的前提。

在仓库文档中，这一集成被描述为正在开发中的能力点（integration point）：目标不仅是「能导出」，还要保证导出的产物能够被进一步降级并优化，在 ExecuTorch 上、尤其是移动与边缘使用场景下高效运行（见 docs/source/en/main_classes/executorch.md）。

Transformers 中的 ExecuTorch 集成全景

虽然本文主体文档很短，但其对应的实现相当完整。集成代码集中在 integrations.executorch 模块（源码位于 src/transformers/integrations/executorch.py），并在包顶层公开了两个导出入口：

• TorchExportableModuleWithStaticCache：面向静态缓存（StaticCache） 的解码器（decoder-only）模型可导出封装模块；
• convert_and_export_with_cache：一键式「封装 + 导出」便捷函数。

这两个符号在 src/transformers/init.py 中注册进 integrations.executorch 的导入结构，因此安装 torch 后可直接 from transformers import TorchExportableModuleWithStaticCache, convert_and_export_with_cache 使用。

除文档 autodoc 覆盖的这两个核心对象外，同一文件还提供了支撑更多模型形态的封装类，可作为扩展阅读：

类 / 函数	职责	源码位置
TorchExportableModuleForDecoderOnlyLM	解码器模型的统一可导出封装，按配置在静态 / 混合缓存间自动分发	executorch.py
TorchExportableModuleWithHybridCache	面向含滑窗 / 混合层结构的静态缓存封装	executorch.py
TorchExportableModuleForVLM	视觉语言模型（如 SmolVLM2）的分组件导出（视觉编码器 / 连接器 / 文本解码器）	executorch.py
Seq2SeqLMExportableModule 系列	编码器-解码器（Seq2Seq）模型的编码器 / 解码器分段导出	executorch.py
export_with_dynamic_cache	使用 DynamicCache 的导出路径	executorch.py

这些类的共同特征是：让 PreTrainedModel 在 torch.export 下可导出，同时保证导出产物与 ExecuTorch 运行时可兼容。

核心封装：TorchExportableModuleWithStaticCache

TorchExportableModuleWithStaticCache 是一个针对「带缓存的 decoder-only 语言模型」设计的 recipe 模块（源码定义于 executorch.py）。它之所以必要，是因为原生的 KV Cache 对象（会随推理动态增长的 DynamicCache）内含无法静态化的惰性初始化与动态控制流，不能直接进入 torch.export 的图。

前置约束：三重 sanity check

构造该封装时（__init__）会依次执行严格校验（executorch.py）：

1. 模型必须携带 generation_config，否则抛出 AssertionError；
2. generation_config.use_cache 必须为 True；
3. generation_config.cache_implementation 必须等于 "static"。

也就是说，导出一个模型前需要像下面这样改造其 generation_config：

from transformers import AutoModelForCausalLM, GenerationConfig

model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-LlamaForCausalLM")
model.eval()

model.generation_config = GenerationConfig(
    use_cache=True,
    cache_implementation="static",
    cache_config={"batch_size": 1, "max_cache_len": 32, "device": "cpu"},
)

上述写法来自配套测试 tests/test_executorch.py。注意 GenerationConfig 中 cache_implementation 的合法取值在生成配置中会统一校验，非法取值会直接报错（见 src/transformers/generation/configuration_utils.py）。

batch_size / max_cache_len / device 三参数的解析次序

封装构造函数的签名是 __init__(model, batch_size=None, max_cache_len=None, device=None)。三个参数均遵循同一套回退逻辑（executorch.py）：

• 显式传入则直接采用；
• 未传入则从 generation_config.cache_config 字典中取 batch_size、max_cache_len；
• 若仍未取到 batch_size 或 max_cache_len，抛出 ValueError；
• device 最后回退到 model.device。

对 ExecuTorch 场景而言，batch_size 通常固定为 1（TorchExportableModuleForVLM 的文档字符串也注明 ExecuTorch 下 batch 恒为 1），因为端侧推理以单条请求为主。max_cache_len 决定了预分配的静态 KV 缓冲区大小，直接约束模型能生成的最大序列长度。

导出前的缓存「实体化」与 buffer 注册

缓存若保持惰性初始化则无法导出，因此构造阶段做了两件关键工作：

1. 提前初始化：调用 StaticCache 的 early_initialization(batch_size, num_heads, head_dim, dtype, device)，一次性分配好全部层的 key/value 张量（executorch.py）。StaticCache 本身会依据 config 检查是否存在混合（hybrid）/滑窗层结构，为每一层选择对应的静态层类型（见 src/transformers/cache_utils.py，其内部通过 get_layer_types_and_kwargs + STATIC_LAYER_TYPE_MAPPING 完成层类型分发）。
2. 注册为非持久 buffer：将每一层的 key_cache_i、value_cache_i、cumulative_length_i 注册进模块，persistent=False 使其进入计算图但不进入 state_dict（executorch.py）。

此外，代码中有一个值得注意的兼容处理：如果某层是 StaticSlidingWindowLayer（滑窗静态层），会被替换成普通的 StaticLayer（executorch.py）。原因是滑窗缓存包含无法避免的动态控制流，不利于导出——代价是超出滑窗窗口长度之后的生成不受支持，这是一个需要在实际应用前评估的硬限制。

forward 适配器：与 ExecuTorch LLM Runner 对齐的接口

该封装之所以称得上「recipe（配方）」，其 forward 承担了两大职责（见 executorch.py 的注释）：

1. 让模型对 torch.export 友好：把 Cache 这类不支持的对象从图的输入输出中隐藏起来，只暴露 input_ids / inputs_embeds 与 cache_position；
2. 与 ExecuTorch 运行期兼容：使前向签名与 executorch/extension/llm/runner 保持一致，让导出产物可以在 ExecuTorch 的 LLM Runner 中「开箱即用」。

其 forward 流程大致为（executorch.py）：

# 每次 forward 先把 cache_position[0] 拷入各层 cumulative_length，相当于软重置缓存，
# 否则同一个导出程序会被原地无限期地改写，而导出后的程序无法在两次 generate 之间调用 reset
for layer in self.static_cache.layers:
    layer.cumulative_length.copy_(cache_position[0])

outs = self.model(
    input_ids=input_ids,
    inputs_embeds=inputs_embeds,
    attention_mask=None,
    past_key_values=self.static_cache,
    use_cache=True,
)
# 有 logits 则返回 CausalLMOutputWithPast.logits，否则返回 BaseModelOutputWithPast.last_hidden_state

这套「每步手动前移 cache_position、逐 token 推理」的运行方式，恰好对应 ExecuTorch 中静态图逐 token 解码的执行模型，也决定了在端上做并行 prefill 之外的顺序解码是主路径。

一键导出函数：convert_and_export_with_cache

文档 autodoc 列出的第二个对象是 convert_and_export_with_cache（executorch.py），它把「实例化静态缓存封装 → 构造示例输入 → 调用 torch.export」收敛成一个函数，最小调用方式为：

from transformers import convert_and_export_with_cache

exported_program = convert_and_export_with_cache(model)

其完整签名为：

convert_and_export_with_cache(
    model: PreTrainedModel,
    example_input_ids: torch.Tensor | None = None,
    example_cache_position: torch.Tensor | None = None,
    dynamic_shapes: dict | None = None,
    strict: bool | None = None,
) -> torch.export.ExportedProgram

要点说明：

• 默认示例输入：当未显式提供时，example_input_ids 默认为 torch.tensor([[1]])、example_cache_position 默认为 torch.tensor([0])（executorch.py）。源码注释明确指出默认输入目前只适用于文本模型，视觉 / 音频模型的默认输入支持仍在规划中——多模态模型请使用前文表格中的专用封装类。
• 版本分支：在 torch ≥ 2.6.0 时走公开 API torch.export.export，并透传 dynamic_shapes 与 strict；在更早版本上则出于 PyTorch 兼容性（对应 issue 128394）回退到内部 API torch.export._trace._export，此时 pre_dispatch=False、strict=True 固定，传入的 dynamic_shapes 与 strict 会被忽略并打印警告（executorch.py）。
• 导出全程包在 torch.no_grad() 内，避免在图中留下梯度计算。
• 返回值为 torch.export.ExportedProgram，可继续送入 ExecuTorch 的降级（lowering）与优化流程。

如果需要对序列长度做动态维度（dynamic shapes），可仿照集成源码中解码器的写法，用 torch.export.Dim 描述维度并传入 dynamic_shapes（例如对 input_ids 的第 1 维与 cache_position 的第 0 维绑定同一个带 max 上限的维度对象，见 executorch.py 中 VLM 文本解码器导出示例），或在 convert_and_export_with_cache 的 dynamic_shapes 参数中直接指定。

底层原理：为什么 KV Cache 必须「静态化」

所有上述封装都建立在 StaticCache 之上。在 Transformers 中，StaticCache 是专为 torch.compile 与 torch.export 设计的缓存类（见 src/transformers/cache_utils.py 的 docstring）。它与传统的 DynamicCache 的本质区别在于：

• DynamicCache 随解码长度动态扩展，张量形状在运行期变化，且存在惰性初始化的分支——这类行为在 torch.export 严格追踪（strict tracing）下无法静态化；
• StaticCache 在开始时按 max_cache_len 一次性预分配 key/value 存储，形状固定、控制流简单，天然适合导出为静态计算图；每层的 StaticLayer/StaticSlidingWindowLayer 定义分别位于 cache_utils.py 与 cache_utils.py。

集成代码正是依赖这一特性：封装模块在构造时调用 StaticCache 的 early_initialization（cache_utils.py）把每个 KV 张量实体化，再以非持久 buffer 形式注册进 torch.nn.Module。这样导出时 KV 缓存就成为图内的可复用状态 buffer，而非图的输入/输出，从而绕开了「Cache 对象不能进出 graph」的限制。

同时需要注意，被替换成普通 StaticLayer 的滑窗层意味着：若原模型配置了 sliding_window 与多层混合结构，导出后生成长度一旦超过窗口即不受支持。因此混合缓存有专门封装 TorchExportableModuleWithHybridCache（executorch.py），并在 TorchExportableModuleForDecoderOnlyLM 中依据 config 是否含 layer_types 且 sliding_window 非空来自动选择混合封装或静态封装（executorch.py）。

在 CPU 上验证导出结果：直接封装导出与 generate 测试工具

在真实跑 ExecuTorch 工具链之前，可以用纯 PyTorch 在本机完成「导出 → 推理一致性」验证，方法与仓库测试保持一致。tests/test_executorch.py 中的 test_decoder_only_lm_export 展示了验证链路（test_executorch.py）：

from transformers import AutoModelForCausalLM, GenerationConfig
from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM

model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-LlamaForCausalLM").eval()
model.generation_config = GenerationConfig(
    use_cache=True,
    cache_implementation="static",
    cache_config={"batch_size": 1, "max_cache_len": 32, "device": "cpu"},
)

input_ids = torch.tensor([[1, 2, 3]], dtype=torch.long)
cache_position = torch.arange(3, dtype=torch.long)

module = TorchExportableModuleForDecoderOnlyLM(model)
exported_program = module.export(input_ids=input_ids, cache_position=cache_position)

# eager 结果与导出产物结果在 1e-4 精度内一致
eager_output = model(input_ids=input_ids, use_cache=False).logits
exported_output = exported_program.module()(input_ids=input_ids, cache_position=cache_position)
torch.testing.assert_close(eager_output, exported_output, atol=1e-4, rtol=1e-4)

对应 TorchExportableModuleWithStaticCache 的数值一致性、input_ids 与 inputs_embeds 两种输入通道、以及混合缓存路径（将 sliding_window=16、layer_types 设为全 full_attention 后与 eager 结果对比），均可参照同文件的 test_static_cache_module_forward、test_hybrid_cache_module_forward（test_executorch.py）。

封装还提供了两个顺序解码的 generate 工具，用于在导出后模拟完整生成流程：

• TorchExportableModuleWithStaticCache.generate(exported_program, prompt_token_ids, max_new_tokens)：逐 token 处理 prompt（不做并行 prefill），并依据导出程序中的 key_cache buffer 形状推算出可生成的最大长度，取 prompt_token_len + max_new_tokens 与该上限的较小者（executorch.py）。注释明确说明该函数仅用于测试导出模型，并非要取代原 generate 方法。
• TorchExportableModuleForDecoderOnlyLM.generate(...)：一个功能更完整的静态方法，支持传入 tokenizer 与 prompt 字符串，并可配置 max_new_tokens、do_sample、temperature、top_k、top_p（完整实现采样与 top-k/top-p 过滤逻辑），遇 EOS 即停止并解码返回文本（executorch.py）。

from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM

text = TorchExportableModuleForDecoderOnlyLM.generate(
    exported_program=exported_program,
    tokenizer=tokenizer,
    prompt="The capital of France is",
    max_new_tokens=20,
    device="cpu",
)

注意此类 generate 属于为验证导出产物而实现的推理循环，其导出前的前向封装才决定 ExecuTorch 运行时的接口契约。

版本、依赖与限制小结

基于当前仓库源码可以确认的适用前提与限制如下：

• 依赖 torch：TorchExportableModuleWithStaticCache 与 convert_and_export_with_cache 在 src/transformers/init.py 中注册，顶层导入需要 torch 后端可用；无 torch 环境的占位对象见 src/transformers/utils/dummy_pt_objects.py。
• torch 版本行为差异：convert_and_export_with_cache 对 torch < 2.6.0 会忽略 dynamic_shapes 与 strict 并使用内部导出 API（executorch.py）；DynamicCache 的 pytree 导出注册在 torch 2.6.0+ 上验证（executorch.py）。
• 模型配置约束：导出要求 use_cache=True、cache_implementation="static"，并显式提供 batch_size 与 max_cache_len（或写入 cache_config）；滑窗层会被替换为普通静态层，生成长度超出窗口后不受支持。
• 模型类型边界：convert_and_export_with_cache 的默认示例输入仅覆盖文本模型；视觉语言、Seq2Seq 等架构需要走同文件内的专用导出类（如 TorchExportableModuleForVLM、Seq2SeqLMExportableModule）。
• 集成成熟度：主文档将该集成描述为正在开发中的能力点（"An integration point is being developed"），因此在使用时建议跟随当前仓库版本校验 API 细节。

如需进一步了解围绕该能力的上层使用方式，可查阅 docs/source/en/community_integrations/executorch.md（介绍通过 optimum-executorch 一键导出、以及在 ExecuTorch 原生 C++ Runner 中加载 .pte 文件进行流式生成的完整链路）。综合本文的源码级讲解，你已经可以完成从「配置静态缓存 → 封装导出 → 本机数值校验」的全部闭环，为真正进入 ExecuTorch 的端侧部署铺平道路。