From fd41337e3571d2b68e654f6fe2c1a251a8518e0d Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Fri, 17 Jul 2026 15:31:36 -0700 Subject: [PATCH 1/3] Document custom PyTorch NVFP4 ComfyUI export Signed-off-by: Jingyu Xin --- examples/diffusers/README.md | 2 + .../quantization/CUSTOM-TORCH-TO-COMFYUI.md | 84 +++++++++++++++++++ 2 files changed, 86 insertions(+) create mode 100644 examples/diffusers/quantization/CUSTOM-TORCH-TO-COMFYUI.md diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index a9efb5fc3a3..90525973361 100644 --- a/examples/diffusers/README.md +++ b/examples/diffusers/README.md @@ -177,6 +177,8 @@ To additionally apply NVFP4 scale swizzle and padding , add: **You can use the generated checkpoint directly in PyTorch, export a Hugging Face checkpoint (`--hf-ckpt-dir`) to deploy the model on SGLang/vLLM/TRTLLM, or follow the ONNX/TensorRT workflow in [`quantization/ONNX-TRT-Deployment.md`](./quantization/ONNX-TRT-Deployment.md).** +For a custom, non-Diffusers PyTorch model, see [Export a custom PyTorch model to a ComfyUI NVFP4 checkpoint](./quantization/CUSTOM-TORCH-TO-COMFYUI.md). + ## Quantization Aware Training (QAT) Quantization-Aware Training (QAT) is a powerful technique for optimizing your models, particularly when post-training quantization (PTQ) methods fail to meet the requirements for your tasks. By simulating the effects of quantization during training, QAT allows the model to learn to minimize the quantization error, ultimately delivering better accuracy. diff --git a/examples/diffusers/quantization/CUSTOM-TORCH-TO-COMFYUI.md b/examples/diffusers/quantization/CUSTOM-TORCH-TO-COMFYUI.md new file mode 100644 index 00000000000..94f6c4c315b --- /dev/null +++ b/examples/diffusers/quantization/CUSTOM-TORCH-TO-COMFYUI.md @@ -0,0 +1,84 @@ +# API guide: custom PyTorch to ComfyUI NVFP4 + +The [`quantize.py`](./quantize.py) example loads a Hugging Face Diffusers pipeline and exports it with +`export_hf_checkpoint`. This note instead assumes that the model remains a plain PyTorch `torch.nn.Module`: it +does not inherit from Diffusers and its source checkpoint is a regular torch state dict. + +For a custom checkpoint, first instantiate its exact model class, load the state dict, and calibrate it: + +```python +import copy + +import torch + +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq + +model = build_model() +state_dict = torch.load("model.pt", map_location="cpu", weights_only=True) +model.load_state_dict(state_dict, strict=True) +model.cuda().eval() + + +def forward_loop(model): + for sample in calibration_data: + model(**sample) + + +mtq.quantize(model, copy.deepcopy(mtq.NVFP4_DEFAULT_CFG), forward_loop) + +# Optional: save a resumable ModelOpt checkpoint before deployment packing. +mto.save(model, "model_nvfp4_modelopt.pt") +``` + +## Which export API to call + +| Model object | API | +| --- | --- | +| A Diffusers `ModelMixin` or `DiffusionPipeline` | Public `modelopt.torch.export.export_hf_checkpoint` | +| A plain `torch.nn.Module` | Internal `_export_diffusers_checkpoint` shown below | + +The internal Diffusers exporter already accepts a plain `nn.Module` and writes packed NVFP4 tensors to +`model.safetensors`. This is an intermediate packed checkpoint; the adapter in the next section is what makes it +compatible with a particular ComfyUI loader. + +```python +from pathlib import Path + +from modelopt.torch.export.unified_export_hf import _export_diffusers_checkpoint + +_export_diffusers_checkpoint( + pipe=model, + dtype=torch.bfloat16, + export_dir=Path("packed_model"), + components=None, + max_shard_size="1000GB", # Keep the deployment checkpoint in one file. +) +``` + +This is a private, version-sensitive API. It currently requires the `diffusers` package to be installed, although +the model itself remains a plain `nn.Module`. It packs recognized quantized modules in place, so call it on a +disposable model instance after `mto.save`. To avoid the Diffusers dependency, build the same staging file by +calling `_process_quantized_modules` and `hide_quantizers_from_state_dict` directly. + +## Make the packed file ComfyUI-specific + +There is no universal ComfyUI key layout. After packing, add a small adapter for the exact ComfyUI loader or +custom node that will consume the model: + +1. Rename every tensor key to the loader's expected namespace. Rename `.weight`, `.weight_scale`, + `.weight_scale_2`, `.input_scale`, and bias keys consistently. +2. If the loader expects a full checkpoint, merge the required VAE, text encoder, or other tensors from a base + safetensors file. +3. Optionally call `pad_nvfp4_weights`, followed by `swizzle_nvfp4_scales`, only when the target runtime expects + the cuBLAS/comfy_kitchen scale layout. +4. Read `quantization_config` from the generated `packed_model/config.json`. After the final key mapping, call + `build_layerwise_quant_metadata` and write both values into the single-file safetensors header. + +These post-processing helpers live in +[`modelopt/torch/export/diffusers_utils.py`](../../../modelopt/torch/export/diffusers_utils.py). The existing LTX-2 +implementation, `_merge_ltx2`, is a useful example, but its `model.diffusion_model.` prefix and component merge +rules are LTX-2-specific and should not be copied blindly. + +No Diffusers wrapper is required for this path. Until ModelOpt exposes a public plain-module deployment exporter, +pin the ModelOpt revision when using `_export_diffusers_checkpoint` or promote that path to a public API. From a0b42331dcb6c6289c0e857d76156ff5672c13cf Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Fri, 17 Jul 2026 15:38:29 -0700 Subject: [PATCH 2/3] Document FP8 custom ComfyUI export Signed-off-by: Jingyu Xin --- examples/diffusers/README.md | 2 +- .../quantization/CUSTOM-TORCH-TO-COMFYUI.md | 18 +++++++++++------- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index 90525973361..f43d4776a99 100644 --- a/examples/diffusers/README.md +++ b/examples/diffusers/README.md @@ -177,7 +177,7 @@ To additionally apply NVFP4 scale swizzle and padding , add: **You can use the generated checkpoint directly in PyTorch, export a Hugging Face checkpoint (`--hf-ckpt-dir`) to deploy the model on SGLang/vLLM/TRTLLM, or follow the ONNX/TensorRT workflow in [`quantization/ONNX-TRT-Deployment.md`](./quantization/ONNX-TRT-Deployment.md).** -For a custom, non-Diffusers PyTorch model, see [Export a custom PyTorch model to a ComfyUI NVFP4 checkpoint](./quantization/CUSTOM-TORCH-TO-COMFYUI.md). +For a custom, non-Diffusers PyTorch model, see [Export a custom PyTorch model to a ComfyUI FP8 or NVFP4 checkpoint](./quantization/CUSTOM-TORCH-TO-COMFYUI.md). ## Quantization Aware Training (QAT) diff --git a/examples/diffusers/quantization/CUSTOM-TORCH-TO-COMFYUI.md b/examples/diffusers/quantization/CUSTOM-TORCH-TO-COMFYUI.md index 94f6c4c315b..72d16755a03 100644 --- a/examples/diffusers/quantization/CUSTOM-TORCH-TO-COMFYUI.md +++ b/examples/diffusers/quantization/CUSTOM-TORCH-TO-COMFYUI.md @@ -1,4 +1,4 @@ -# API guide: custom PyTorch to ComfyUI NVFP4 +# API guide: custom PyTorch to ComfyUI FP8 or NVFP4 The [`quantize.py`](./quantize.py) example loads a Hugging Face Diffusers pipeline and exports it with `export_hf_checkpoint`. This note instead assumes that the model remains a plain PyTorch `torch.nn.Module`: it @@ -25,10 +25,11 @@ def forward_loop(model): model(**sample) -mtq.quantize(model, copy.deepcopy(mtq.NVFP4_DEFAULT_CFG), forward_loop) +quant_cfg = mtq.NVFP4_DEFAULT_CFG # Use mtq.FP8_DEFAULT_CFG for FP8. +mtq.quantize(model, copy.deepcopy(quant_cfg), forward_loop) # Optional: save a resumable ModelOpt checkpoint before deployment packing. -mto.save(model, "model_nvfp4_modelopt.pt") +mto.save(model, "model_quantized_modelopt.pt") ``` ## Which export API to call @@ -38,7 +39,7 @@ mto.save(model, "model_nvfp4_modelopt.pt") | A Diffusers `ModelMixin` or `DiffusionPipeline` | Public `modelopt.torch.export.export_hf_checkpoint` | | A plain `torch.nn.Module` | Internal `_export_diffusers_checkpoint` shown below | -The internal Diffusers exporter already accepts a plain `nn.Module` and writes packed NVFP4 tensors to +The internal Diffusers exporter already accepts a plain `nn.Module` and writes packed deployment tensors to `model.safetensors`. This is an intermediate packed checkpoint; the adapter in the next section is what makes it compatible with a particular ComfyUI loader. @@ -61,17 +62,20 @@ the model itself remains a plain `nn.Module`. It packs recognized quantized modu disposable model instance after `mto.save`. To avoid the Diffusers dependency, build the same staging file by calling `_process_quantized_modules` and `hide_quantizers_from_state_dict` directly. +For FP8, the export call is unchanged. FP8 does not use the NVFP4 `weight_scale_2`, padding, or scale-swizzle +steps, so only the loader-specific key mapping, optional component merge, and metadata steps below remain. + ## Make the packed file ComfyUI-specific There is no universal ComfyUI key layout. After packing, add a small adapter for the exact ComfyUI loader or custom node that will consume the model: 1. Rename every tensor key to the loader's expected namespace. Rename `.weight`, `.weight_scale`, - `.weight_scale_2`, `.input_scale`, and bias keys consistently. + `.input_scale`, bias, and NVFP4 `.weight_scale_2` keys consistently. 2. If the loader expects a full checkpoint, merge the required VAE, text encoder, or other tensors from a base safetensors file. -3. Optionally call `pad_nvfp4_weights`, followed by `swizzle_nvfp4_scales`, only when the target runtime expects - the cuBLAS/comfy_kitchen scale layout. +3. For NVFP4 only, optionally call `pad_nvfp4_weights`, followed by `swizzle_nvfp4_scales`, when the target + runtime expects the cuBLAS/comfy_kitchen scale layout. Skip both for FP8. 4. Read `quantization_config` from the generated `packed_model/config.json`. After the final key mapping, call `build_layerwise_quant_metadata` and write both values into the single-file safetensors header. From 18a3e58f3eaf8829c59749e923dd774d3547c836 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Fri, 17 Jul 2026 15:51:56 -0700 Subject: [PATCH 3/3] Document direct custom checkpoint export Signed-off-by: Jingyu Xin --- .../quantization/CUSTOM-TORCH-TO-COMFYUI.md | 159 ++++++++++++------ 1 file changed, 104 insertions(+), 55 deletions(-) diff --git a/examples/diffusers/quantization/CUSTOM-TORCH-TO-COMFYUI.md b/examples/diffusers/quantization/CUSTOM-TORCH-TO-COMFYUI.md index 72d16755a03..a63e291f2c2 100644 --- a/examples/diffusers/quantization/CUSTOM-TORCH-TO-COMFYUI.md +++ b/examples/diffusers/quantization/CUSTOM-TORCH-TO-COMFYUI.md @@ -1,10 +1,13 @@ # API guide: custom PyTorch to ComfyUI FP8 or NVFP4 The [`quantize.py`](./quantize.py) example loads a Hugging Face Diffusers pipeline and exports it with -`export_hf_checkpoint`. This note instead assumes that the model remains a plain PyTorch `torch.nn.Module`: it -does not inherit from Diffusers and its source checkpoint is a regular torch state dict. +`export_hf_checkpoint`. This note covers a different case: the checkpoint belongs to an arbitrary Python class, +which may contain several attached `torch.nn.Module` objects but is not itself an `nn.Module`, Hugging Face model, +or Diffusers pipeline. -For a custom checkpoint, first instantiate its exact model class, load the state dict, and calibrate it: +Explicitly collect the components that should appear in the checkpoint into a temporary `nn.ModuleDict`. This is +only an export view; the original class does not need to change. Modules stored in ordinary Python lists or dicts +must also be added explicitly. ```python import copy @@ -14,75 +17,121 @@ import torch import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq -model = build_model() -state_dict = torch.load("model.pt", map_location="cpu", weights_only=True) -model.load_state_dict(state_dict, strict=True) -model.cuda().eval() +owner = build_custom_object() +owner.load_checkpoint(torch.load("model.pt", map_location="cpu", weights_only=True)) +modelopt_root = torch.nn.ModuleDict( + { + "denoiser": owner.denoiser, + "text_encoder": owner.text_encoder, + "vae": owner.vae, + } +).cuda().eval() -def forward_loop(model): + +def forward_loop(_): for sample in calibration_data: - model(**sample) + # Calling the owner's custom inference path is fine because it uses the same modules. + # It must exercise every component selected for quantization. + owner.run_calibration(sample) -quant_cfg = mtq.NVFP4_DEFAULT_CFG # Use mtq.FP8_DEFAULT_CFG for FP8. -mtq.quantize(model, copy.deepcopy(quant_cfg), forward_loop) +quant_cfg = mtq.FP8_DEFAULT_CFG # Or mtq.NVFP4_DEFAULT_CFG. +mtq.quantize(modelopt_root, copy.deepcopy(quant_cfg), forward_loop) -# Optional: save a resumable ModelOpt checkpoint before deployment packing. -mto.save(model, "model_quantized_modelopt.pt") +# Save a resumable ModelOpt checkpoint before the destructive deployment packing step. +mto.save(modelopt_root, "model_quantized_modelopt.pt") ``` -## Which export API to call +This example quantizes the supported leaves in `modelopt_root` that match the preset. Narrow `quant_cfg` by +component prefix when only a subset should be quantized. -| Model object | API | -| --- | --- | -| A Diffusers `ModelMixin` or `DiffusionPipeline` | Public `modelopt.torch.export.export_hf_checkpoint` | -| A plain `torch.nn.Module` | Internal `_export_diffusers_checkpoint` shown below | +## Export without Hugging Face or Diffusers -The internal Diffusers exporter already accepts a plain `nn.Module` and writes packed deployment tensors to -`model.safetensors`. This is an intermediate packed checkpoint; the adapter in the next section is what makes it -compatible with a particular ComfyUI loader. +There is currently no public one-call exporter for this shape of object. The minimal path is to call +`_process_quantized_modules` on the temporary root, hide ModelOpt quantizer state, adapt the tensor/config names to +the target ComfyUI loader, and write safetensors directly: ```python +import json from pathlib import Path -from modelopt.torch.export.unified_export_hf import _export_diffusers_checkpoint - -_export_diffusers_checkpoint( - pipe=model, - dtype=torch.bfloat16, - export_dir=Path("packed_model"), - components=None, - max_shard_size="1000GB", # Keep the deployment checkpoint in one file. +import torch +from safetensors.torch import save_file + +from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format +from modelopt.torch.export.diffusers_utils import ( + build_layerwise_quant_metadata, + hide_quantizers_from_state_dict, + pad_nvfp4_weights, + swizzle_nvfp4_scales, ) -``` - -This is a private, version-sensitive API. It currently requires the `diffusers` package to be installed, although -the model itself remains a plain `nn.Module`. It packs recognized quantized modules in place, so call it on a -disposable model instance after `mto.save`. To avoid the Diffusers dependency, build the same staging file by -calling `_process_quantized_modules` and `hide_quantizers_from_state_dict` directly. - -For FP8, the export call is unchanged. FP8 does not use the NVFP4 `weight_scale_2`, padding, or scale-swizzle -steps, so only the loader-specific key mapping, optional component merge, and metadata steps below remain. +from modelopt.torch.export.quant_utils import ( + get_quant_config, + has_quantized_modules, + sync_tied_input_amax, +) +from modelopt.torch.export.unified_export_hf import _process_quantized_modules -## Make the packed file ComfyUI-specific -There is no universal ComfyUI key layout. After packing, add a small adapter for the exact ComfyUI loader or -custom node that will consume the model: +def export_comfyui_checkpoint( + modelopt_root, + output_path, + adapt_for_loader, + *, + dtype=torch.bfloat16, + padding_strategy=None, + swizzle=False, +): + if not has_quantized_modules(modelopt_root): + raise ValueError("No ModelOpt-quantized modules found") + + modelopt_config = get_quant_config(modelopt_root) + quant_config = convert_hf_quant_config_format(modelopt_config) + quant_algo = quant_config.get("quant_algo") + if quant_algo not in {"FP8", "NVFP4"}: + raise ValueError(f"Expected FP8 or NVFP4, got {quant_algo!r}") + + sync_tied_input_amax(modelopt_root) + _process_quantized_modules(modelopt_root, dtype) + + with hide_quantizers_from_state_dict(modelopt_root): + state_dict = dict(modelopt_root.state_dict()) + + # Rename/filter tensor keys, merge external tensors, and update any config layer names. + state_dict, quant_config = adapt_for_loader(state_dict, quant_config) + state_dict = { + key: tensor.detach().to(device="cpu", copy=True).contiguous() + for key, tensor in state_dict.items() + } + + if quant_algo == "NVFP4": + if padding_strategy is not None: + pad_nvfp4_weights(state_dict, padding_strategy) + if swizzle: + swizzle_nvfp4_scales(state_dict) + elif padding_strategy is not None or swizzle: + raise ValueError("Padding and swizzling apply only to NVFP4") + + metadata = { + "quantization_config": json.dumps(quant_config), + "_quantization_metadata": build_layerwise_quant_metadata(state_dict, quant_config), + } + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + save_file(state_dict, str(output_path), metadata=metadata) +``` -1. Rename every tensor key to the loader's expected namespace. Rename `.weight`, `.weight_scale`, - `.input_scale`, bias, and NVFP4 `.weight_scale_2` keys consistently. -2. If the loader expects a full checkpoint, merge the required VAE, text encoder, or other tensors from a base - safetensors file. -3. For NVFP4 only, optionally call `pad_nvfp4_weights`, followed by `swizzle_nvfp4_scales`, when the target - runtime expects the cuBLAS/comfy_kitchen scale layout. Skip both for FP8. -4. Read `quantization_config` from the generated `packed_model/config.json`. After the final key mapping, call - `build_layerwise_quant_metadata` and write both values into the single-file safetensors header. +`adapt_for_loader` is model-specific: there is no universal ComfyUI key layout. Attached modules already included +in `modelopt_root` are present in its state dict; use the callback to merge only tensors outside that root. Preserve +the `.weight_scale` and `.weight_scale_2` suffixes, or build equivalent layer metadata yourself. If keys are renamed, +also rename layer references in `quant_config`, such as `ignore` or `quantized_layers`. -These post-processing helpers live in -[`modelopt/torch/export/diffusers_utils.py`](../../../modelopt/torch/export/diffusers_utils.py). The existing LTX-2 -implementation, `_merge_ltx2`, is a useful example, but its `model.diffusion_model.` prefix and component merge -rules are LTX-2-specific and should not be copied blindly. +FP8 is the simpler case: it has `.weight`, `.weight_scale`, and usually `.input_scale`, but no +`.weight_scale_2`, padding, or scale swizzling. NVFP4 adds those optional runtime-specific steps. -No Diffusers wrapper is required for this path. Until ModelOpt exposes a public plain-module deployment exporter, -pin the ModelOpt revision when using `_export_diffusers_checkpoint` or promote that path to a public API. +Despite their source filenames, the calls above do not construct or require a Hugging Face or Diffusers object; +Diffusers imports in `diffusers_utils.py` are optional. This minimal helper targets `FP8_DEFAULT_CFG` and +`NVFP4_DEFAULT_CFG` on ModelOpt-supported quantizable leaves such as `nn.Linear`. A custom weight-bearing leaf needs +a ModelOpt conversion/export plugin. These helpers are private and version-sensitive, and packing mutates the +selected modules in place.