Spaces:
Running on Zero
Running on Zero
File size: 4,542 Bytes
855e1f5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | """Export toolkit-quantized modules into ComfyUI ``comfy_quant`` checkpoints —
the inverse of toolkit/util/comfy_quant_import.py.
Every quantized OstrisLinear whose backend has a comfy storage format emits
its quantized tensors plus the ``<prefix>.comfy_quant`` uint8 JSON marker:
- convrot8 (int8_tensorwise + convrot): weight int8 [out, in], fp32
weight_scale [out, 1] (comfy_kitchen's per-channel convention)
- nvfp4: high-nibble-first packed fp4 pairs, e4m3 block scales re-swizzled
to the cuBLAS 128x4 tile layout, fp32 weight_scale_2 per-tensor scale and
optional AWQ pre_quant_scale
- float8_e4m3fn: fp8_e4m3 weight + one fp32 per-tensor weight_scale
- convrotcomfyw4a4: via convrot_quant.export_comfy_convrot_w4a4
Biases are NOT emitted here — they are ordinary parameters and flow through
the regular state_dict path.
"""
import json
from typing import Dict, List, Tuple
import torch
from toolkit.util.ostris_quant import OstrisLinear
def comfy_quant_marker(conf: dict) -> torch.Tensor:
return torch.tensor(list(json.dumps(conf).encode("utf-8")), dtype=torch.uint8)
@torch.no_grad()
def export_comfy_quantized_layers(
root: torch.nn.Module,
) -> Tuple[Dict[str, torch.Tensor], List[str], List[str]]:
"""Comfy-format state-dict entries for every quantized OstrisLinear in
``root``. Returns ``(entries, exported_names, unexportable_names)`` —
entries are keyed by module path in root's layout (run
convert_state_dict_on_save afterwards for the checkpoint layout);
unexportable_names lists quantized modules whose backend has no comfy
storage format (the caller decides whether to dequantize instead)."""
from toolkit.util.convrot_quant import (
ConvRotComfyW4A4Quantizer,
export_comfy_convrot_w4a4,
)
from toolkit.util.nvfp4_quant import swap_nvfp4_nibbles, swizzle_nvfp4_scales
entries: Dict[str, torch.Tensor] = {}
exported: List[str] = []
unexportable: List[str] = []
for name, module in root.named_modules():
if not isinstance(module, OstrisLinear):
continue
if hasattr(module, "cr8_qdata"):
rot = int(getattr(module, "cr8_rot_size", 1) or 1)
conf = {"format": "int8_tensorwise"}
if rot > 1:
conf.update({"convrot": True, "convrot_groupsize": rot})
entries[f"{name}.weight"] = module.cr8_qdata.detach().cpu().contiguous()
entries[f"{name}.weight_scale"] = (
module.cr8_scales.view(torch.float32)
.detach()
.cpu()
.reshape(module.out_features, 1)
.contiguous()
)
entries[f"{name}.comfy_quant"] = comfy_quant_marker(conf)
elif hasattr(module, "nv4_qdata"):
entries[f"{name}.weight"] = swap_nvfp4_nibbles(
module.nv4_qdata.detach().cpu()
)
scales = module.nv4_scales.view(torch.float8_e4m3fn).detach().cpu()
entries[f"{name}.weight_scale"] = swizzle_nvfp4_scales(
scales.reshape(module.out_features, module.in_features // 16)
).view(torch.float8_e4m3fn)
entries[f"{name}.weight_scale_2"] = (
module.nv4_pts.view(torch.float32).detach().cpu().reshape(())
)
if hasattr(module, "nv4_pre_scale"):
entries[f"{name}.pre_quant_scale"] = (
module.nv4_pre_scale.view(torch.float32).detach().cpu().contiguous()
)
entries[f"{name}.comfy_quant"] = comfy_quant_marker({"format": "nvfp4"})
elif hasattr(module, "f8_qdata"):
entries[f"{name}.weight"] = module.f8_qdata.detach().cpu().contiguous()
entries[f"{name}.weight_scale"] = (
module.f8_scale.view(torch.float32).detach().cpu().reshape(())
)
entries[f"{name}.comfy_quant"] = comfy_quant_marker(
{"format": "float8_e4m3fn", "full_precision_matrix_mult": True}
)
elif isinstance(module.ostris_quantizer, ConvRotComfyW4A4Quantizer):
layer_entries = export_comfy_convrot_w4a4(module, f"{name}.")
layer_entries.pop(f"{name}.bias", None)
entries.update(
{k: v.detach().cpu() if torch.is_tensor(v) else v for k, v in layer_entries.items()}
)
else:
unexportable.append(name)
continue
exported.append(name)
return entries, exported, unexportable
|