From cd8534fb2f96b10534a8613d1739307b961b1bc7 Mon Sep 17 00:00:00 2001 From: Uxito-Ada <414416158@qq.com> Date: Wed, 29 Jul 2026 05:41:43 +0000 Subject: [PATCH 1/4] Add Quantization Introduction and Bug Fix --- docs/en/quantization.md | 178 ++++++++++++++++++++++++---- docs/javascripts/mathjax.js | 14 +++ docs/zh/quantization.md | 175 +++++++++++++++++++++++---- mkdocs.yml | 6 +- telefuser/core/config.py | 13 +- telefuser/ops/torchao_fp8_linear.py | 9 +- tools/convert/converter.py | 26 ++-- 7 files changed, 347 insertions(+), 74 deletions(-) create mode 100644 docs/javascripts/mathjax.js diff --git a/docs/en/quantization.md b/docs/en/quantization.md index 328bd7e..8d0260e 100644 --- a/docs/en/quantization.md +++ b/docs/en/quantization.md @@ -1,46 +1,172 @@ # Quantization -## TeleFuser FP8 deployment for Qwen-Image +Quantization stores or computes selected tensors at lower precision. In TeleFuser, it is used for model weights, Linear inputs, attention operands, and KV caches. These paths are independent: enabling one does not quantize the entire pipeline. -TeleFuser uses TorchAO as the backend for FP8 weight-only linear quantization. +## Core ideas -First, install Telefuser as [here](https://github.com/Tele-AI/TeleFuser#install). +`W8A16` means 8-bit weights and 16-bit activations; `W8A8` means both are 8-bit. The accumulator or output can still use BF16 or FP32, so this notation does not describe every tensor in the operation. -Next, download the `Qwen-Image-2512` model to `TF_MODEL_ZOO_PATH` (or pass `--model_root`), and run the FP8 example: +For symmetric integer quantization, TeleFuser's offline converter uses: + +$$ +s = \frac{\max |x|}{127}, \qquad +q = \mathrm{clamp}\left(\mathrm{round}\left(\frac{x}{s}\right), -128, 127\right), +\qquad \hat{x} = s q. +$$ + +The scale can cover a tensor, an output channel, a token, or a small block. Smaller groups usually reduce error, but require more scale values and more complicated kernels. + +| Granularity | TeleFuser example | +| --- | --- | +| Per-tensor | ComfyUI INT8/FP8 conversion | +| Per-output-channel | Offline INT8/FP8 weights and `LinearFP8` weights | +| Per-token | `LinearFP8`, LiveAct, and LingBot FP8 activations | +| Per-block | MXFP and NVFP4 conversion kernels | + +## Available paths + +| Path | Precision | Entry point | +| --- | --- | --- | +| TorchAO online quantization | W8A8 or weight-only W8A16 | `QuantType.TORCHAO_FP8` | +| bitsandbytes online quantization | NF4 weight-only, W4A16 | `QuantType.BNB_NF4` | +| Scaled FP8 checkpoint | FP8 weights and dynamic FP8 activations, W8A8 | `torch_dtype=torch.float8_e4m3fn` | +| Offline checkpoint conversion | INT8, FP8, MXFP4/6/8, or NVFP4 weights | `tools/convert/converter.py` | + +`QuantType` also contains formats that are not connected to generic online model loading. An enum value alone is not evidence that a model implements that path. + +## Online Linear quantization + +### TorchAO FP8: version-dependent + +TeleFuser probes TorchAO's dynamic-activation and weight-only FP8 APIs. The first available API determines the mode: +dynamic activation and weight FP8 is W8A8, while `Float8WeightOnlyConfig` with BF16 inputs is W8A16. Check the +conversion log and run a real forward instead of inferring the mode from `QuantType.TORCHAO_FP8`. + +The supported TeleFuser models are Wan, Qwen-Image, and LTX transformer blocks. The default filter skips names such as `head`, `time_embedding`, `time_projection`, and `patch_embedding`. + +```python +import torch + +from telefuser.core.config import QuantConfig, QuantKernelBackend, QuantType +from telefuser.core.module_manager import ModuleManager + +quant_config = QuantConfig( + enabled=True, + quant_type=QuantType.TORCHAO_FP8, + kernel_backend=QuantKernelBackend.TORCHAO, +) +manager = ModuleManager(torch_dtype=torch.bfloat16, device="cpu") +manager.load_model( + dit_paths, + device="cuda", + torch_dtype=torch.bfloat16, + quant_config=quant_config, +) +``` + +Run the complete Qwen-Image example with: ```bash -PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ python examples/qwen_image/qwen_image_t2i_telefuser_fp8_h100.py \ --prompt "A cat playing piano" \ - --aspect_ratio 1:1 \ - --num-inference-steps 16 \ - --seed 42 \ --output qwen_image_fp8.png ``` -Then generated image is saved as `qwen_image_fp8.png` at current directory. +FP8 reduces weight memory traffic; W8A8 additionally quantizes Linear inputs. Whether it improves latency depends on +the selected mode, matrix shapes, GPU, TorchAO version, and `torch.compile` behavior. + +TorchAO and PyTorch must be version-compatible. Check the +[TorchAO release compatibility table](https://github.com/pytorch/ao/releases) instead of installing the newest +TorchAO release blindly; an import warning or failure means that configuration has not been validated. + +### bitsandbytes NF4: W4A16 + +NF4 uses a non-uniform 4-bit codebook designed for approximately normal weight distributions. TeleFuser replaces selected Linear layers with `bitsandbytes.nn.Linear4bit`, uses BF16 compute, and enables compressed quantization statistics. + +```python +quant_config = QuantConfig( + enabled=True, + quant_type=QuantType.BNB_NF4, + kernel_backend=QuantKernelBackend.BITSANDBYTES, +) +``` + +The full example is `examples/qwen_image/qwen_image_t2i_telefuser_nf4_h100.py`. NF4 usually saves more weight memory than FP8, but 4-bit decoding does not guarantee lower latency. + +## Scaled FP8 checkpoints: W8A8 + +This path is different from TorchAO. A compatible checkpoint already contains E4M3FN weights and a scale for each output channel. `LinearFP8` dynamically quantizes every input row to FP8, then calls a scaled GEMM through `tf_kernel` or vLLM/CUTLASS. The output returns to BF16 or FP16. + +For each row, scaled FP8 follows the same absmax idea: -## TeleFuser NF4 deployment for Qwen-Image +$$ +s = \frac{\max |x|}{\mathrm{max}(\mathrm{E4M3FN})}, \qquad +q = \mathrm{cast}_{\mathrm{E4M3FN}}\left( +\mathrm{clamp}\left(\frac{x}{s}, f_{\min}, f_{\max}\right)\right). +$$ -Telefuser users bitsandbytes as its backend of NF4 weight-only linear quantization. +Load only a checkpoint that follows TeleFuser's expected weight and scale layout: -First, install Telefuser as [here](https://github.com/Tele-AI/TeleFuser#install). +```python +manager.load_model( + fp8_checkpoint, + device="cuda", + torch_dtype=torch.float8_e4m3fn, +) +``` + +Changing `torch_dtype` does not turn an arbitrary BF16 checkpoint into a scaled FP8 checkpoint. Start from the supplied Qwen-Image or Wan FP8 examples. + +## Offline conversion -Next, download `Qwen-Image-2512` model to your `TF_MODEL_ZOO_PATH` (or specify the model path to `--model_root`), and run the NF4 `Qwen-Image-2512` example: +The converter quantizes selected two-dimensional weights and writes the scales beside them. It creates an artifact; it does not add an inference kernel. ```bash -# set the below cuda versions according to your environment -export BNB_CUDA_VERSION=128 -export CUDA_HOME=/usr/local/cuda-12.8 -export LD_LIBRARY_PATH=/usr/local/cuda-12.8/lib64:${LD_LIBRARY_PATH:-} - -PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ -python examples/qwen_image/qwen_image_t2i_telefuser_nf4_h100.py \ - --prompt "A cat playing piano" \ - --aspect_ratio 1:1 \ - --num-inference-steps 16 \ - --seed 42 \ - --output qwen_image_nf4.png +python tools/convert/converter.py \ + --source /path/to/source \ + --output /path/to/output \ + --model_type wan_dit \ + --quantized \ + --linear_dtype fp8 \ + --non_linear_dtype torch.bfloat16 \ + --single_file ``` -Then generated image is saved as `qwen_image_nf4.png` at current directory. +`--linear_dtype` accepts `int8`, `fp8`, `mxfp4`, `mxfp6`, `mxfp8`, and `nvfp4`. + +- `int8` and `fp8` use one absmax scale per output row. ComfyUI mode uses one scale per tensor. +- MXFP4/6/8 use 32-value blocks and E8M0 scales through `lightx2v_kernel`. +- NVFP4 packs two values per byte and uses one E4M3 scale per 16 values plus a tensor-wide global scale. +- Non-quantized tensors are converted to `--non_linear_dtype`. + +The default scale keys are `_scale`; NVFP4 also writes `_global_scale`. MXFP and NVFP4 conversion +requires CUDA and `lightx2v_kernel`. TeleFuser's generic loader does not execute these artifacts; a matching consumer +must implement the same layout and GEMM. + +## Other quantized data paths + +- **LiveAct FP8:** wraps Linear layers with vLLM-style dynamic W8A8 GEMM. Weights are cached in FP8 and activations are quantized per token. +- **LingBot-Video MoE FP8:** quantizes expert weights per output channel and routed activations per row, then uses `torch._scaled_mm`. +- **SageAttention:** all three TeleFuser variants quantize Q/K to INT8. `2_8_16` uses FP16 P/V, while `2_8_8` and its SM90 variant use FP8 P/V. These kernels do not change model weights. +- **LiveAct FP8 KV cache:** stores K/V as E4M3FN with one FP32 scale per last-dimension vector, then dequantizes to the requested attention dtype on load. + +## Validation + +Use the same prompt, input, seed, scheduler, and inference steps as the BF16 baseline. Check all of the following: + +1. The log reports a non-zero number of converted Linear layers. +2. Optional backends pass a real `torch.inference_mode()` forward; a successful import alone is insufficient. +3. Peak loading memory and steady-state VRAM are measured separately. +4. Latency is measured after warmup with CUDA synchronization or CUDA events. +5. Generated images or videos are compared with the baseline; successful execution alone is not an accuracy test. + +For operator comparisons, useful error measures are: + +$$ +E_{\max} = \max |y_q-y|, \qquad +E_{\mathrm{rel}} = \frac{\lVert y_q-y \rVert_2}{\lVert y \rVert_2}. +$$ + +Relevant tests include `tests/unit/quantize/test_quantized_linear.py` and `tests/unit/models/test_lingbot_video_moe.py`. + +For TorchAO's distinction between weight-only and dynamic FP8, see its [quantized inference guide](https://docs.pytorch.org/ao/stable/workflows/inference.html). diff --git a/docs/javascripts/mathjax.js b/docs/javascripts/mathjax.js new file mode 100644 index 0000000..87891fb --- /dev/null +++ b/docs/javascripts/mathjax.js @@ -0,0 +1,14 @@ +window.MathJax = { + tex: { + inlineMath: [["\\(", "\\)"]], + displayMath: [["\\[", "\\]"], ["$$", "$$"]], + }, + options: { + ignoreHtmlClass: ".*|", + processHtmlClass: "arithmatex", + }, +}; + +document$.subscribe(() => { + MathJax.typesetPromise(); +}); diff --git a/docs/zh/quantization.md b/docs/zh/quantization.md index 5f1591e..2877605 100644 --- a/docs/zh/quantization.md +++ b/docs/zh/quantization.md @@ -1,47 +1,170 @@ # 量化 -## 使用 TeleFuser 为 Qwen-Image 部署 FP8 量化 +量化用较低精度存储或计算部分张量。TeleFuser 分别对模型权重、Linear 输入、注意力张量和 KV cache 提供了量化路径；启用其中一项不等于整个 pipeline 都使用低精度。 -TeleFuser 使用 TorchAO 作为 FP8 仅权重线性量化的后端。 +## 基本原理 -首先，按照[此处说明](https://github.com/Tele-AI/TeleFuser#install)安装 TeleFuser。 +`W8A16` 表示 8 bit 权重和 16 bit 激活，`W8A8` 表示权重与激活都是 8 bit。累加器和输出仍可使用 BF16 或 FP32，因此该记号并不描述算子里的所有张量。 -接下来，将 `Qwen-Image-2512` 模型下载到 `TF_MODEL_ZOO_PATH` 目录中（也可以通过 `--model_root` 指定模型路径），然后运行 FP8 版本的 `Qwen-Image-2512` 示例： +TeleFuser 离线转换器的对称 INT8 量化为： + +$$ +s = \frac{\max |x|}{127}, \qquad +q = \mathrm{clamp}\left(\mathrm{round}\left(\frac{x}{s}\right), -128, 127\right), +\qquad \hat{x} = s q. +$$ + +一个 scale 可以对应整个 tensor、一个输出通道、一个 token 或一个小 block。分组越细通常误差越小，但 scale 数量和 kernel 复杂度也越高。 + +| 粒度 | TeleFuser 中的例子 | +| --- | --- | +| Per-tensor | ComfyUI INT8/FP8 转换 | +| Per-output-channel | 离线 INT8/FP8 权重、`LinearFP8` 权重 | +| Per-token | `LinearFP8`、LiveAct、LingBot FP8 激活 | +| Per-block | MXFP 和 NVFP4 转换 kernel | + +## 已有路径 + +| 路径 | 精度 | 入口 | +| --- | --- | --- | +| TorchAO 在线量化 | W8A8 或仅权重 W8A16 | `QuantType.TORCHAO_FP8` | +| bitsandbytes 在线量化 | NF4 仅权重，W4A16 | `QuantType.BNB_NF4` | +| scaled FP8 checkpoint | FP8 权重和动态 FP8 激活，W8A8 | `torch_dtype=torch.float8_e4m3fn` | +| 离线 checkpoint 转换 | INT8、FP8、MXFP4/6/8 或 NVFP4 权重 | `tools/convert/converter.py` | + +`QuantType` 中还有尚未接入通用在线加载的格式。枚举值存在，不代表具体模型已经实现该路径。 + +## 在线 Linear 量化 + +### TorchAO FP8：取决于版本 + +TeleFuser 会探测 TorchAO 的动态激活与仅权重 FP8 API，首先找到的 API 决定实际模式：动态激活和权重 +FP8 是 W8A8；BF16 输入配合 `Float8WeightOnlyConfig` 是 W8A16。应检查转换日志并运行真实前向， +不能只根据 `QuantType.TORCHAO_FP8` 推断模式。 + +当前接入 Wan、Qwen-Image 和 LTX 的 transformer blocks。默认过滤器会跳过 `head`、`time_embedding`、`time_projection` 和 `patch_embedding` 等名称。 + +```python +import torch + +from telefuser.core.config import QuantConfig, QuantKernelBackend, QuantType +from telefuser.core.module_manager import ModuleManager + +quant_config = QuantConfig( + enabled=True, + quant_type=QuantType.TORCHAO_FP8, + kernel_backend=QuantKernelBackend.TORCHAO, +) +manager = ModuleManager(torch_dtype=torch.bfloat16, device="cpu") +manager.load_model( + dit_paths, + device="cuda", + torch_dtype=torch.bfloat16, + quant_config=quant_config, +) +``` + +完整的 Qwen-Image 示例： ```bash -PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ python examples/qwen_image/qwen_image_t2i_telefuser_fp8_h100.py \ --prompt "A cat playing piano" \ - --aspect_ratio 1:1 \ - --num-inference-steps 16 \ - --seed 42 \ --output qwen_image_fp8.png ``` -生成的图像将以 `qwen_image_fp8.png` 为文件名保存在当前目录中。 +FP8 会减少权重显存和读取流量；W8A8 还会量化 Linear 输入。是否降低延迟取决于实际模式、矩阵形状、 +GPU、TorchAO 版本以及 `torch.compile`。 -## 使用 TeleFuser 为 Qwen-Image 部署 NF4 量化 +TorchAO 必须与 PyTorch 版本兼容。不要直接安装最新版，应先查看 +[TorchAO release 兼容表](https://github.com/pytorch/ao/releases)；出现 import warning 或失败，说明该组合尚未通过验证。 + +### bitsandbytes NF4：W4A16 + +NF4 使用针对近似正态分布权重设计的非均匀 4 bit 码本。TeleFuser 把选中的 Linear 替换为 `bitsandbytes.nn.Linear4bit`，使用 BF16 计算，并压缩量化统计量。 + +```python +quant_config = QuantConfig( + enabled=True, + quant_type=QuantType.BNB_NF4, + kernel_backend=QuantKernelBackend.BITSANDBYTES, +) +``` -TeleFuser 使用 bitsandbytes 作为 NF4 仅权重线性量化的后端。 +完整示例是 `examples/qwen_image/qwen_image_t2i_telefuser_nf4_h100.py`。NF4 通常比 FP8 更省权重显存，但 4 bit 解码不保证延迟更低。 -首先，按照[此处说明](https://github.com/Tele-AI/TeleFuser#install)安装 TeleFuser。 +## Scaled FP8 checkpoint：W8A8 -接下来，将 `Qwen-Image-2512` 模型下载到 `TF_MODEL_ZOO_PATH` 目录中（也可以通过 `--model_root` 指定模型路径），然后运行 NF4 版本的 `Qwen-Image-2512` 示例： +这条路径与 TorchAO 不同。兼容的 checkpoint 已经包含 E4M3FN 权重和逐输出通道 scale。`LinearFP8` 在每次 forward 时把输入逐行量化为 FP8，再通过 `tf_kernel` 或 vLLM/CUTLASS 执行 scaled GEMM，输出恢复为 BF16 或 FP16。 + +每一行的 scaled FP8 同样使用 absmax： + +$$ +s = \frac{\max |x|}{\mathrm{max}(\mathrm{E4M3FN})}, \qquad +q = \mathrm{cast}_{\mathrm{E4M3FN}}\left( +\mathrm{clamp}\left(\frac{x}{s}, f_{\min}, f_{\max}\right)\right). +$$ + +只能加载符合 TeleFuser 权重和 scale 布局的 checkpoint： + +```python +manager.load_model( + fp8_checkpoint, + device="cuda", + torch_dtype=torch.float8_e4m3fn, +) +``` + +只修改 `torch_dtype` 不能把任意 BF16 checkpoint 变成 scaled FP8 checkpoint。应从仓库提供的 Qwen-Image 或 Wan FP8 示例开始。 + +## 离线转换 + +转换器量化选中的二维权重，并把 scale 一同写入 checkpoint。它只生成 artifact，不会自动提供推理 kernel。 ```bash -# 请根据你的运行环境设置以下 CUDA 版本 -export BNB_CUDA_VERSION=128 -export CUDA_HOME=/usr/local/cuda-12.8 -export LD_LIBRARY_PATH=/usr/local/cuda-12.8/lib64:${LD_LIBRARY_PATH:-} - -PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ -python examples/qwen_image/qwen_image_t2i_telefuser_nf4_h100.py \ - --prompt "A cat playing piano" \ - --aspect_ratio 1:1 \ - --num-inference-steps 16 \ - --seed 42 \ - --output qwen_image_nf4.png +python tools/convert/converter.py \ + --source /path/to/source \ + --output /path/to/output \ + --model_type wan_dit \ + --quantized \ + --linear_dtype fp8 \ + --non_linear_dtype torch.bfloat16 \ + --single_file ``` -生成的图像将以 `qwen_image_nf4.png` 为文件名保存在当前目录中。 +`--linear_dtype` 支持 `int8`、`fp8`、`mxfp4`、`mxfp6`、`mxfp8` 和 `nvfp4`。 + +- `int8` 和 `fp8` 默认每个输出行使用一个 absmax scale；ComfyUI 模式每个 tensor 使用一个 scale。 +- MXFP4/6/8 通过 `lightx2v_kernel` 使用 32 元素 block 和 E8M0 scale。 +- NVFP4 每个字节打包两个值，每 16 个值使用一个 E4M3 scale，另有一个全 tensor 的 global scale。 +- 未量化 tensor 会转换为 `--non_linear_dtype`。 + +默认 scale key 是 `_scale`；NVFP4 还会写入 `_global_scale`。MXFP 和 NVFP4 转换依赖 CUDA +与 `lightx2v_kernel`。TeleFuser 通用 loader 不执行这些 artifact；消费端必须实现相同的布局和 GEMM。 + +## 其他量化数据路径 + +- **LiveAct FP8：** 使用 vLLM 风格的动态 W8A8 GEMM；权重缓存为 FP8，激活逐 token 量化。 +- **LingBot-Video MoE FP8：** expert 权重逐输出通道量化，路由后的激活逐行量化，再调用 `torch._scaled_mm`。 +- **SageAttention：** TeleFuser 的三个变体都把 Q/K 量化为 INT8；`2_8_16` 使用 FP16 P/V，`2_8_8` 及其 SM90 变体使用 FP8 P/V。这些 kernel 不修改模型权重。 +- **LiveAct FP8 KV cache：** K/V 保存为 E4M3FN，每个末维向量使用一个 FP32 scale；加载时反量化为注意力请求的 dtype。 + +## 验证方法 + +固定与 BF16 baseline 相同的 prompt、输入、seed、scheduler 和推理步数，并检查： + +1. 日志中的 Linear 转换数量不为零。 +2. 可选后端通过真实的 `torch.inference_mode()` 前向；仅 import 成功不够。 +3. 分开记录加载峰值显存和稳态显存。 +4. warmup 后使用 CUDA 同步或 CUDA event 测量延迟。 +5. 对比生成图像或视频；进程成功退出不代表精度合格。 + +算子对比可记录： + +$$ +E_{\max} = \max |y_q-y|, \qquad +E_{\mathrm{rel}} = \frac{\lVert y_q-y \rVert_2}{\lVert y \rVert_2}. +$$ + +相关测试包括 `tests/unit/quantize/test_quantized_linear.py` 和 `tests/unit/models/test_lingbot_video_moe.py`。 +TorchAO 对 weight-only 与动态 FP8 的定义见其[量化推理文档](https://docs.pytorch.org/ao/stable/workflows/inference.html)。 diff --git a/mkdocs.yml b/mkdocs.yml index ab66034..2673c3e 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -101,6 +101,10 @@ markdown_extensions: - pymdownx.tasklist: custom_checkbox: true +extra_javascript: + - javascripts/mathjax.js + - https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js + # Plugins plugins: - search: @@ -130,13 +134,13 @@ plugins: Stream Scheduler: 流式调度器 TeleFuser and AIPerf: TeleFuser 与 AIPerf Service Metadata: 服务元数据 + Quantization: 量化 Adding New Example: 新增示例 Model Loading: 模型加载 TF-Kernel: TF-Kernel CPU Offloading: CPU卸载 Attention: 注意力机制 Feature Cache: 特征缓存 - Quantization: 量化 Latent Cache: 潜空间缓存 Adding New Model: 新增模型 Adding New Stage: 新增阶段 diff --git a/telefuser/core/config.py b/telefuser/core/config.py index 9e9394e..96587af 100644 --- a/telefuser/core/config.py +++ b/telefuser/core/config.py @@ -334,7 +334,8 @@ class QuantType(Enum): MXFP4: Microscaling FP4 (OCP standard). NVFP4: NVIDIA FP4 format (Blackwell+). BNB_NF4: bitsandbytes weight-only NF4 linear path. - TORCHAO_FP8: TorchAO dynamic-activation FP8 linear path. + TORCHAO_FP8: TorchAO FP8 linear path; activation precision depends on + the TorchAO API selected at runtime. """ FP8 = auto() @@ -360,15 +361,11 @@ class QuantKernelBackend(Enum): @dataclass class QuantConfig: - """Configuration for online quantization during model loading. + """Configuration for model quantization during loading. .. warning:: - This is an interface definition only. The actual quantization - functionality is NOT yet implemented. This config serves as a - placeholder for future online quantization support. - - Online quantization converts bf16/fp16 weights to lower precision (FP8/INT8/MXFP4/etc.) - at load time, reducing memory footprint without requiring pre-quantized checkpoint files. + Supported values depend on the model. Wan, Qwen-Image, and LTX currently + implement ``TORCHAO_FP8`` and ``BNB_NF4`` for transformer Linear layers. Attributes: enabled: Whether to enable online quantization. diff --git a/telefuser/ops/torchao_fp8_linear.py b/telefuser/ops/torchao_fp8_linear.py index 4cf93d4..beedc0c 100644 --- a/telefuser/ops/torchao_fp8_linear.py +++ b/telefuser/ops/torchao_fp8_linear.py @@ -1,8 +1,8 @@ """TorchAO FP8 helpers for TeleFuser DiT linear layers. -This backend applies TorchAO dynamic-activation FP8 + FP8 weight quantization -to selected ``nn.Linear`` modules. It targets W8A8 inference on Hopper/H100 -and keeps the integration close to TorchAO's native ``quantize_`` API. +The selected mode depends on the APIs exported by the installed TorchAO +version. Dynamic activation and weight FP8 uses W8A8; the weight-only fallback +uses W8A16 for BF16 inputs. """ from __future__ import annotations @@ -87,7 +87,8 @@ def replace_linear_layers_with_torchao_fp8( """Quantize selected ``nn.Linear`` modules with TorchAO FP8. Returns the number of selected Linear layers. TorchAO performs in-place - conversion through ``quantize_``. + conversion through ``quantize_``. The first available dynamic or + weight-only API determines the activation precision. """ _check_torchao_fp8_available() diff --git a/tools/convert/converter.py b/tools/convert/converter.py index 2d134d2..cb83190 100755 --- a/tools/convert/converter.py +++ b/tools/convert/converter.py @@ -314,7 +314,7 @@ def quantize_model( adapter_keys=None, key_idx=2, ignore_key=None, - linear_type="int8", + linear_dtype="int8", non_linear_dtype=torch.float, comfyui_mode=False, comfyui_keys=[], @@ -382,7 +382,7 @@ def quantize_model( original_size += original_tensor_size # Quantize tensor and store results - quantizer = CONVERT_WEIGHT_REGISTER[linear_type](tensor) + quantizer = CONVERT_WEIGHT_REGISTER[linear_dtype](tensor) w_q, scales, extra = quantizer.weight_quant_func(tensor, comfyui_mode) weight_global_scale = extra.get("weight_global_scale", None) # For nvfp4 @@ -590,9 +590,10 @@ def convert_key(key): if args.quantized: if args.full_quantized and args.comfyui_mode: logger.info("Quant all tensors...") - assert args.linear_dtype, "Error: only support 'torch.int8' and 'torch.float8_e4m3fn'." + target_dtype = dtype_mapping.get(args.linear_dtype) + assert target_dtype, "Error: full quantization only supports int8 and fp8." for k in converted_weights.keys(): - converted_weights[k] = converted_weights[k].float().to(args.linear_dtype) + converted_weights[k] = converted_weights[k].float().to(target_dtype) else: converted_weights = quantize_model( converted_weights, @@ -601,7 +602,7 @@ def convert_key(key): adapter_keys=args.adapter_keys, key_idx=args.key_idx, ignore_key=args.ignore_key, - linear_type=args.linear_type, + linear_dtype=args.linear_dtype, non_linear_dtype=args.non_linear_dtype, comfyui_mode=args.comfyui_mode, comfyui_keys=args.comfyui_keys, @@ -734,6 +735,14 @@ def copy_non_weight_files(source_dir, target_dir): logger.info("Non-weight files and subdirectories copied") +def _normalize_linear_dtype(value: str) -> str: + aliases = { + "torch.int8": "int8", + "torch.float8_e4m3fn": "fp8", + } + return aliases.get(value, value) + + def main(): parser = argparse.ArgumentParser(description="Model weight format converter") parser.add_argument("-s", "--source", required=True, help="Input path (file or directory)") @@ -783,10 +792,10 @@ def main(): help="Device to use for quantization (cpu/cuda)", ) parser.add_argument( - "--linear_type", - type=str, + "--linear_dtype", + type=_normalize_linear_dtype, choices=["int8", "fp8", "nvfp4", "mxfp4", "mxfp6", "mxfp8"], - help="Quant type for linear", + help="Data type used to quantize linear weights", ) parser.add_argument( "--non_linear_dtype", @@ -848,7 +857,6 @@ def main(): logger.warning("--chunk_size is ignored when using --single_file option.") if args.quantized: - args.linear_dtype = dtype_mapping.get(args.linear_type, None) args.non_linear_dtype = eval(args.non_linear_dtype) model_type_keys_map = { From 4c508494995682711488cad5e62ba48cd2f2175f Mon Sep 17 00:00:00 2001 From: Heyang Sun <60865256+Uxito-Ada@users.noreply.github.com> Date: Wed, 29 Jul 2026 13:44:26 +0800 Subject: [PATCH 2/4] Update torchao_fp8_linear.py --- telefuser/ops/torchao_fp8_linear.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/telefuser/ops/torchao_fp8_linear.py b/telefuser/ops/torchao_fp8_linear.py index beedc0c..d0dccf5 100644 --- a/telefuser/ops/torchao_fp8_linear.py +++ b/telefuser/ops/torchao_fp8_linear.py @@ -1,7 +1,7 @@ """TorchAO FP8 helpers for TeleFuser DiT linear layers. -The selected mode depends on the APIs exported by the installed TorchAO -version. Dynamic activation and weight FP8 uses W8A8; the weight-only fallback +TorchAO auto switches between two modes depending on hardware paltform: +Dynamic activation and weight FP8 uses W8A8; the weight-only fallback uses W8A16 for BF16 inputs. """ From 2af01b53b77277ff09b91747447783c85123d60b Mon Sep 17 00:00:00 2001 From: Heyang Sun <60865256+Uxito-Ada@users.noreply.github.com> Date: Wed, 29 Jul 2026 13:47:18 +0800 Subject: [PATCH 3/4] Update quantization.md --- docs/zh/quantization.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/zh/quantization.md b/docs/zh/quantization.md index 2877605..48f2aaa 100644 --- a/docs/zh/quantization.md +++ b/docs/zh/quantization.md @@ -36,9 +36,9 @@ $$ ## 在线 Linear 量化 -### TorchAO FP8：取决于版本 +### TorchAO FP8: -TeleFuser 会探测 TorchAO 的动态激活与仅权重 FP8 API，首先找到的 API 决定实际模式：动态激活和权重 +TorchAO会根据硬件平台选择支持的内核：动态激活/仅权重 FP8 。TeleFuser会根据TorchAO选择的 API 决定实际模式：动态激活和权重 FP8 是 W8A8；BF16 输入配合 `Float8WeightOnlyConfig` 是 W8A16。应检查转换日志并运行真实前向， 不能只根据 `QuantType.TORCHAO_FP8` 推断模式。 From 3f3f7a55fa51c5d134139d186b32be885126299f Mon Sep 17 00:00:00 2001 From: Heyang Sun <60865256+Uxito-Ada@users.noreply.github.com> Date: Wed, 29 Jul 2026 13:50:44 +0800 Subject: [PATCH 4/4] Update quantization.md --- docs/en/quantization.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/en/quantization.md b/docs/en/quantization.md index 8d0260e..7467846 100644 --- a/docs/en/quantization.md +++ b/docs/en/quantization.md @@ -36,9 +36,9 @@ The scale can cover a tensor, an output channel, a token, or a small block. Smal ## Online Linear quantization -### TorchAO FP8: version-dependent +### TorchAO FP8: -TeleFuser probes TorchAO's dynamic-activation and weight-only FP8 APIs. The first available API determines the mode: +TorchAO automatically selects FP8 kernels depending on your hardware paltform, between dynamic-activation and weight-only FP8s. TeleFuser determines the mode according to the selection of TorchAO: dynamic activation and weight FP8 is W8A8, while `Float8WeightOnlyConfig` with BF16 inputs is W8A16. Check the conversion log and run a real forward instead of inferring the mode from `QuantType.TORCHAO_FP8`.