WIP PTQ tool

Author: contentis

tools/ptq/checkpoint_merger.py

@@ -0,0 +1,255 @@
+import argparse
+import logging
+import sys
+import yaml
+import re
+from typing import Dict, Tuple
+import torch
+from safetensors.torch import save_file
+import json
+
+# Add comfyui to path if needed
+import os
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
+
+import comfy.utils
+from comfy.ops import QUANT_FORMAT_MIXINS
+from comfy.quant_ops import F8_E4M3_MAX, F4_E2M1_MAX
+
+class QuantizationConfig:
+    def __init__(self, config_path: str):
+        with open(config_path, 'r') as f:
+            self.config = yaml.safe_load(f)
+
+        # Compile disable list patterns
+        self.disable_patterns = []
+        for pattern in self.config.get('disable_list', []):
+            # Convert glob-style patterns to regex
+            regex_pattern = pattern.replace('*', '.*')
+            self.disable_patterns.append(re.compile(regex_pattern))
+
+        # Parse per-layer dtype config
+        self.per_layer_dtype = self.config.get('per_layer_dtype', {})
+        self.dtype_patterns = []
+        for pattern, dtype in self.per_layer_dtype.items():
+            regex_pattern = pattern.replace('*', '.*')
+            self.dtype_patterns.append((re.compile(regex_pattern), dtype))
+
+        logging.info(f"Loaded config with {len(self.disable_patterns)} disable patterns")
+        logging.info(f"Per-layer dtype rules: {self.per_layer_dtype}")
+
+    def should_quantize(self, layer_name: str) -> bool:
+        for pattern in self.disable_patterns:
+            if pattern.match(layer_name):
+                logging.debug(f"Layer {layer_name} disabled by pattern {pattern.pattern}")
+                return False
+        return True
+
+    def get_dtype(self, layer_name: str) -> str:
+        for pattern, dtype in self.dtype_patterns:
+            if pattern.match(layer_name):
+                return dtype
+        return None
+
+def load_amax_artefact(artefact_path: str) -> Dict:
+    logging.info(f"Loading amax artefact from {artefact_path}")
+
+    with open(artefact_path, 'r') as f:
+        data = json.load(f)
+
+    if 'amax_values' not in data:
+        raise ValueError("Invalid artefact format: missing 'amax_values' key")
+
+    metadata = data.get('metadata', {})
+    amax_values = data['amax_values']
+
+    logging.info(f"Loaded {len(amax_values)} amax values from artefact")
+    logging.info(f"Artefact metadata: {metadata}")
+
+    return data
+
+def get_scale_fp8(amax: float, dtype: torch.dtype) -> torch.Tensor:
+    scale = amax / torch.finfo(dtype).max
+    scale_tensor = torch.tensor(scale, dtype=torch.float32)
+    return scale_tensor
+
+def get_scale_nvfp4(amax: float, dtype: torch.dtype) -> torch.Tensor:
+    scale = amax / (F8_E4M3_MAX * F4_E2M1_MAX)
+    scale_tensor = torch.tensor(scale, dtype=torch.float32)
+    return scale_tensor
+
+def get_scale(amax: float, dtype: torch.dtype):
+    if dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:
+        return get_scale_fp8(amax, dtype)
+    elif dtype in [torch.float4_e2m1fn_x2]:
+        return get_scale_nvfp4(amax, dtype)
+    else:
+        raise ValueError(f"Unsupported dtype {dtype} ")
+
+def apply_quantization(
+    checkpoint: Dict,
+    amax_values: Dict[str, float],
+    config: QuantizationConfig
+) -> Tuple[Dict, Dict]:
+    quantized_dict = {}
+    layer_metadata = {}
+
+    for key, amax in amax_values.items():
+        if key.endswith(".input_quantizer"):
+            continue
+
+        layer_name = ".".join(key.split(".")[:-1])
+
+        if not config.should_quantize(layer_name):
+            logging.debug(f"Layer {layer_name} disabled by config")
+            continue
+
+        dtype_str = config.get_dtype(layer_name)
+        dtype = getattr(torch, dtype_str)
+        device = torch.device("cuda") # Required for NVFP4
+
+        weight = checkpoint.pop(f"{layer_name}.weight").to(device)
+        scale_tensor = get_scale(amax, dtype)
+
+        input_amax = amax_values.get(f"{layer_name}.input_quantizer", None)
+        if input_amax is not None:
+            input_scale = get_scale(input_amax, dtype)
+            quantized_dict[f"{layer_name}.input_scale"] = input_scale.clone()
+
+        # logging.info(f"Quantizing {layer_name}: amax={amax}, scale={scale_tensor:.6f}")
+        tensor_layout = QUANT_FORMAT_MIXINS[dtype_str]["layout_type"]
+        quantized_weight, layout_params = tensor_layout.quantize(
+            weight,
+            scale=scale_tensor,
+            dtype=dtype
+        )
+        quantized_dict[f"{layer_name}.weight_scale"] = scale_tensor.clone()
+        quantized_dict[f"{layer_name}.weight"] = quantized_weight.clone()
+
+        if "block_scale" in layout_params:
+            quantized_dict[f"{layer_name}.weight_block_scale"] = layout_params["block_scale"].clone()
+
+        # Build metadata
+        layer_metadata[layer_name] = {
+            "format": dtype_str,
+            "params": {}
+        }
+
+    logging.info(f"Quantized {len(layer_metadata)} layers")
+
+    quantized_dict = quantized_dict | checkpoint
+
+    metadata_dict = {
+        "_quantization_metadata": json.dumps({
+            "format_version": "1.0",
+            "layers": layer_metadata
+        })
+    }
+    return quantized_dict, metadata_dict
+
+
+def main():
+    """Main entry point for checkpoint merger."""
+
+    parser = argparse.ArgumentParser(
+        description="Merge calibration artifacts with checkpoint to create quantized model",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+
+    parser.add_argument(
+        "--artefact",
+        required=True,
+        help="Path to calibration artefact JSON file (amax values)"
+    )
+    parser.add_argument(
+        "--checkpoint",
+        required=True,
+        help="Path to original checkpoint to quantize"
+    )
+    parser.add_argument(
+        "--config",
+        required=True,
+        help="Path to YAML quantization config file"
+    )
+    parser.add_argument(
+        "--output",
+        required=True,
+        help="Output path for quantized checkpoint"
+    )
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        help="Enable debug logging"
+    )
+
+    args = parser.parse_args()
+
+    # Configure logging
+    if args.debug:
+        logging.basicConfig(
+            level=logging.DEBUG,
+            format='[%(levelname)s] %(name)s: %(message)s'
+        )
+    else:
+        logging.basicConfig(
+            level=logging.INFO,
+            format='[%(levelname)s] %(message)s'
+        )
+
+    # Print header
+
+    # Step 1: Load calibration artefact
+    logging.info("[1/5] Loading calibration artefact...")
+    try:
+        artefact_data = load_amax_artefact(args.artefact)
+        amax_values = artefact_data['amax_values']
+    except Exception as e:
+        logging.error(f"Failed to load artefact: {e}")
+        sys.exit(1)
+
+    # Step 2: Load quantization config
+    logging.info("[2/5] Loading quantization config...")
+    try:
+        config = QuantizationConfig(args.config)
+    except Exception as e:
+        logging.error(f"Failed to load config: {e}")
+        sys.exit(1)
+
+    # Step 3: Load checkpoint
+    logging.info("[3/5] Loading checkpoint...")
+    try:
+        checkpoint = comfy.utils.load_torch_file(args.checkpoint)
+        logging.info(f"Loaded checkpoint with {len(checkpoint)} keys")
+    except Exception as e:
+        logging.error(f"Failed to load checkpoint: {e}")
+        sys.exit(1)
+
+    # Step 4: Apply quantization
+    logging.info("[4/5] Applying quantization...")
+    try:
+        quantized_dict, metadata_json = apply_quantization(
+            checkpoint,
+            amax_values,
+            config
+        )
+    except Exception as e:
+        logging.error(f"Failed to apply quantization: {e}")
+        import traceback
+        traceback.print_exc()
+        sys.exit(1)
+
+    # Step 5: Export quantized checkpoint
+    logging.info("[5/5] Exporting quantized checkpoint...")
+    try:
+        save_file(quantized_dict, args.output, metadata=metadata_json)
+
+    except Exception as e:
+        logging.error(f"Failed to export checkpoint: {e}")
+        import traceback
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()
+

tools/ptq/configs/flux_fp8.yml

@@ -0,0 +1,23 @@
+# FLUX Quantization Config: Transformer Blocks Only
+#
+# Quantize only double and single transformer blocks,
+# leave input/output projections in higher precision.
+
+disable_list: [
+  # Disable input projections
+  "*img_in*",
+  "*txt_in*",
+  "*time_in*",
+  "*vector_in*",
+  "*guidance_in*",
+  
+  # Disable output layers
+  "*final_layer*",
+
+  # Disable positional embeddings
+  "*pe_embedder*",
+]
+
+per_layer_dtype: {
+  "*": "float8_e4m3fn",
+}

tools/ptq/configs/flux_nvfp4.yml

@@ -0,0 +1,27 @@
+# FLUX Quantization Config: Transformer Blocks Only
+#
+# Quantize only double and single transformer blocks,
+# leave input/output projections in higher precision.
+
+disable_list: [
+  # Disable input projections
+  "*img_in*",
+  "*txt_in*",
+  "*time_in*",
+  "*vector_in*",
+  "*guidance_in*",
+  
+  # Disable output layers
+  "*final_layer*",
+
+  # Disable positional embeddings
+  "*pe_embedder*",
+
+  "*modulation*",
+  "*txt_mod*",
+  "*img_mod*",
+]
+
+per_layer_dtype: {
+  "*": "float4_e2m1fn_x2",
+}