摘要： 本文将厘清 Eager Mode、FX Mode、JIT (TorchScript)、PT2E 以及 NVIDIA ModelOpt 等核心概念的关系，帮助开发者梳理现代深度学习部署的技术选型路线。

---

1. 概念澄清：Eager、FX 与 JIT 的底层逻辑

在探讨具体的量化工具之前，首先需要明确 PyTorch 中不同运行模式与图捕获机制的区别。

1.1 Eager Mode vs. Graph Mode

• Eager Mode（动态图模式）：PyTorch 默认的逐行执行模式。极度灵活、易于调试，但缺乏全局计算图信息，无法自动完成跨层算子融合（如 $\text{Conv}+\text{BN}+\text{ReLU}$）。
• Graph Mode（静态图模式）：将 Python 计算过程提取并表示为中间表示图（IR Graph）。由于掌握了全图结构，图模式能够进行全局优化、算子融合以及自动化 QDQ 节点插入。

1.2 FX Graph Mode 和 JIT (TorchScript) 的关系

很多开发者容易混淆 FX Graph Mode 与 JIT，实际上两者没有任何直接关系：

• JIT (TorchScript)：PyTorch 1.x 时代的 C++ 静态图编译器，通过 torch.jit.trace 或 torch.jit.script 捕获模型。其语法限制极多，遇到 Python 控制流（如 if/else）或复杂数据结构时极易报出难以调试的 C++ 异常。
• FX (torch.fx)：PyTorch 1.8+ 推出的纯 Python 级别符号追踪（Symbolic Tracing）工具。FX Graph Mode 指的是“基于 FX 捕获纯 Python 语言图进行的量化”，整个图的提取、修改与节点插入完全由 FX 驱动，摆脱了对 C++ JIT 的依赖。

---

2. 传统量化工具链的局限与替代方案

2.1 传统 torch.ao.quantization 的现状

PyTorch 官方早期的 torch.ao.quantization（包含 Eager 和 FX Graph Mode）已进入维护状态，并计划在后续版本中被全面废弃。

• JIT 路线停滞：官方的核心编译与部署重心已全面转移至 PyTorch 2.x 的 torch.compile、torch.export 以及 ExecuTorch，TorchScript 不再添加新特性。
• 新一代标准（PT2E）：对于必须使用 Graph Mode 导出部署的场景，官方推荐采用 PT2E (PyTorch 2 Export) 路线（通过 torch.export 提取无 Python 依赖的 ATen 计算图）结合 TorchAO (torchao) 进行量化。

2.2 NVIDIA 工具链演进：从 pytorch-quantization 到 ModelOpt

对于部署至 NVIDIA TensorRT 的传统流水线，NVIDIA 的官方工具链也完成了重大迭代：

• 旧方案（pytorch-quantization）：采用 Eager Mode 进行动态图重换（Module Replacement），虽然在 Python 交互中稳定易用，但缺乏对 Transformer/ViT、大语言模型以及 FP8/NVFP4/INT4 等新精度的原生支持，目前已被列入遗留/替代路线。
• 新方案（NVIDIA ModelOpt）：NVIDIA 统一的模型优化旗舰库（nvidia-modelopt），彻底摒弃了旧 JIT 路线，全面拥抱 PyTorch 2.x (Dynamo / FX / PT2E) 与现代化 Eager 接口。

---

3. PyTorch 图架构与量化生态演进脉络

下表总结了 PyTorch 编译图与量化范式的三代演进：

演进阶段	图捕获机制	代表性量化工具/范式	核心特点与局限
第一代	Eager Mode / JIT (TorchScript)	torch.ao.quantization (Eager)
		pytorch-quantization	Eager 模式灵活但需手动改代码；JIT 捕获图极易断裂、调试困难。
第二代	FX (Symbolic Tracing)	torch.ao.quantization (FX Graph)	纯 Python 级别的图分析，较 JIT 更易调试，但对复杂控制流支持仍有受限。
第三代	PyTorch 2.0 Dynamo / torch.export (PT2E)	TorchAO (torchao)
		NVIDIA ModelOpt (modelopt)	完全原生支持复杂 Python 逻辑，统一 Server/Edge 端导出与多精度（INT8/FP8/NVFP4）支持。

---

1. PyTorch Quantization (Eager vs. FX Mode)

这是 PyTorch 官方（旧版 torch.ao.quantization）提供的两种模型量化工作流：

• Eager Mode 量化 (Imperative Mode)

  • 机制：基于动态图（Eager Execution）。需要手动修改模型结构，显式插入 QuantStub 和 DeQuantStub 节点，并将需要量化的模块替换为量化版本（如 QuantizedConv2d）。
  • 特点：对 Python 灵活控制流（如 if/else、循环）支持好，但需要侵入式修改代码，无法自动做复杂的算子融合（Fuse）。

• FX Graph Mode 量化 (Declarative Mode)

  • 机制：基于 PyTorch FX 的静态图分析。通过 torch.fx.symbolic_trace 自动捕获模型的计算图，解析节点关系，并自动插入量化节点和完成算子融合。
  • 特点：完全非侵入式（无需改模型代码），自动化程度高。但受限于 FX 符号追踪的能力，对包含复杂动态控制流或非 PyTorch 原生 Python 操作的代码支持较差。

---

2. JIT / Graph Mode

• JIT (TorchScript)：PyTorch 早期用于消除 Python 运行时依赖、将模型打包为 C++ 部署格式的静态图编译器。可以通过 torch.jit.script 或 torch.jit.trace 将 Python 模型编译为 ScriptModule 图结构。
• Graph Mode：指将模型的 Python 动态执行过程转换并表示为 静态中间表示图（IR Graph） 的运行模式。无论是 TorchScript 图还是 FX 图，都属于 Graph Mode。

---

3. PrepareMethod.JIT_STRIP

PrepareMethod.JIT_STRIP 是 PyTorch 内部或相关图变换/量化工具链（如某些 TorchScript/FX 量化准备过程）中的一个枚举标记，主要用于清理与剥离（Strip）编译图中的 JIT 冗余元数据。

• 背景与作用：当模型从 Eager 模式通过 JIT 编译转为 Graph 模式并准备进行量化时，编译图（IR Graph）中会夹杂大量 Python 运行时的调试信息、JIT 追踪元数据、冗余的节点包装或无用的 Hook。
• 行为：JIT_STRIP 作为准备阶段（Prepare Phase）的方法之一，会遍历计算图，剥离/净化（Strip） 这些专属于 JIT 运行时的附加节点与冗余指令，将图简化为纯净的线性/分支算子图，以便后续的量化分析器（Quantizer）能够精准地插入 Fake-Quantization（伪量化）节点或 QDQ 节点。

4. NVIDIA 官方开源的 pytorch-quantization 工具包

（包含在 NVIDIA TensorRT 仓库中，主要用于 QAT 和 PTZ 导出 ONNX 到 TensorRT），它的底层运行机制是标准的 Eager Mode。

为什么是 Eager Mode？

• 动态图重换（Module Replacement）：pytorch_quantization 的核心逻辑（如 quant_modules.initialize()）是通过 Python 的动态特性，在 Eager 模式下将 PyTorch 原生的 nn.Conv2d、nn.Linear 等模块替换为包装好的 quant_nn.QuantConv2d、quant_nn.QuantLinear。
• 自定义 Python 逻辑：QuantDescriptor、calib（校准器，如 Max/Histogram/MSE 校准）以及 tensor_quant 内部的伪量化（Fake Quantization）逻辑，都是作为普通的 PyTorch nn.Module 和 Python 函数在 Eager 动态图中执行的。

---

它与 JIT / TensorRT 的关系

虽然它运行在 Eager Mode，但它的终极目标是导出（Export）：

1. 训练/校准阶段（Eager Mode）：
   在 PyTorch Python 交互式或训练脚本中运行，利用 Eager 模式的灵活性收集 Activation 的统计信息（校准），或者进行 Quantization-Aware Training (QAT) 微调。
2. 导出阶段（ONNX Export）：
   校准或训练结束后，通过 torch.onnx.export() 将带有伪量化算子（如 TensorQuantizer 生成的 QuantizeLinear / DequantizeLinear，即 QDQ 节点）的模型导出为 ONNX 图。
3. 推理阶段（TensorRT Engine）：
   最后将 ONNX 解析进 NVIDIA TensorRT，由 TensorRT 读取这些 QDQ 节点并完成 INT8 算子融合与硬件加速（此时才脱离 Python 进入编译好的 Native 引擎）。

---

5. 为什么弃用 JIT (TorchScript)？

• 官方战略转移：从 PyTorch 2.0 开始，官方的技术路线已全面由旧版的 JIT (TorchScript) 转向 torch.compile 和 torch.export。
• 维护状态：TorchScript 处于维护状态（不再添加新特性），且旧版的 torch.ao.quantization（包含 Eager 和 FX Graph Mode）计划被废弃移除。

---

根据最新的 PyTorch (TorchAO / PyTorch 2.x) 架构推荐：

• 首选：PT2E Graph Mode（torch.export + torchao）

• 机制：通过 torch.export 提取无 Python 依赖的完整的 ATen 计算图（Exported Program），并在图上进行校准、算子融合与伪量化插入（prepare_pt2e / convert_pt2e）。

• 适用场景：静态量化（Static PTQ）、量化感知训练（QAT）、以及需要跨端部署（导出至 ExecuTorch 或 Inductor C++ 运行时）的场景。

• 灵活动态量化/LLM 场景：Eager/Compile Mode (torchao)

• 如果做权重仅量化（Weight-Only，如 INT4/INT8 W4A16）或大模型动态量化，直接使用新一代 torchao 提供的 Eager/Dynamo 原生 API，配合 torch.compile 即可获得性能提升，无需繁琐的图捕获。

---

6. 部署方向的统一

目标部署环境	推荐的量化与导出方式	替代的旧方案
云端 / 服务端 (CPU/GPU)	torchao + torch.compile / Inductor	旧版 FX Graph Mode / TorchScript JIT
边缘端 / 移动端	PT2E Graph Mode (torch.export) $\to$ ExecuTorch	TorchScript JIT (.pt 文件)
NVIDIA TensorRT	pytorch-quantization (Eager 校准) $\to$ ONNX (QDQ) $\to$ TensorRT	—

结论：在新版 PyTorch 体系中，应该采用 Graph Mode，但实现载体由旧版的 JIT/TorchScript 升级为基于 PyTorch 2.0 规范的 torch.export (PT2E)。

7 ModelOpt 的底层工作原理

ModelOpt 的量化工作流主要采用以下两种机制，均与旧版 JIT 无关：

• 基于 Dynamo / FX Graph Mode (PT2E)

  • 对于复杂模型（如 ViT、LLM、多模态），ModelOpt 利用 PyTorch 2.0 的 torch.compiler / torch.export (Dynamo) 机制来捕获并分析模型的计算图。
  • 它在图的层面上实现高阶的自动算子融合、节点改写、自动插入 Fake-Quantization (QDQ) 节点，并精确匹配 NVIDIA 硬件（如 Hopper/Blackwell 上的 Transformer Engine）的执行图。

• 基于 Eager 动态图模式（灵活动态插入）

  • 对于权重量化（Weight-Only, 如 INT4/FP8 AWQ/GPTQ）或传统的 CNN/BEV 静态校准，ModelOpt 提供了极大简化的 Eager 接口。
  • 你只需传入原生的 nn.Module，ModelOpt 会直接在 Eager 动态图内存中完成模块的 Hook 或替换，无需经过任何 JIT script/trace 步骤。