From ab10de7521a6af8cd9a0cea641596f797f11993a Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Wed, 12 Aug 2026 23:59:05 +0700 Subject: [PATCH] Match Comfy Qwen FP32 projections --- PARITY.md | 26 +++++++++ src/h3_blackwell_runtime/nvfp4.py | 6 +- src/h3_blackwell_runtime/qwen3vl_text.py | 4 +- tools/compare_qwen0_quantized_dispatch.py | 46 +++++++++++++++ ...patch_comfy_qwen_quantized_tensor_probe.py | 58 +++++++++++++++++++ 5 files changed, 136 insertions(+), 4 deletions(-) create mode 100644 tools/compare_qwen0_quantized_dispatch.py create mode 100644 tools/patch_comfy_qwen_quantized_tensor_probe.py diff --git a/PARITY.md b/PARITY.md index 6a1ebcc..5b08590 100644 --- a/PARITY.md +++ b/PARITY.md @@ -149,6 +149,32 @@ investigation must capture/compare the effective `QuantizedTensor` parameters and direct packed-weight wrapper at dispatch time. Do not revisit attention, sampler, or MLP ordering until that exact module-dispatch difference is found. +### Effective Weight Root Cause And Resolution + +The storage-safe dispatch probe found the active difference. Comfy's text +encoder calls Qwen with `dtype=torch.float32`; its full-precision NVFP4 branch +casts each projection to an effective FP32 matrix before `F.linear`. Direct had +constructed the Qwen encoder with BF16, causing its full-precision NVFP4 branch +to dequantize the same packed data into BF16. + +The probe proved that packed qdata, tensor scale, and block scale are exact for +all Q/K/V/gate/up/down projections. With direct Qwen switched to FP32: + +| Projection check | Result | +| --- | --- | +| Effective Q/K/V/gate/up/down FP32 weights | Bit-exact to Comfy | +| Q/K/V/gate/up linear outputs | Bit-exact to Comfy | +| Layer-0 O-projected attention | Mean/max `0` / `0` | +| Layer-0 post-attention residual | Mean/max `0` / `0` | +| Layer-0 post-attention RMSNorm | Mean/max `0` / `0` | +| Layer-0 MLP output | Mean/max `0` / `0` | +| Layer-0 final output | Mean/max `0` / `0` | + +The special `down` dispatch sub-probe that supplied Comfy's already-activated +input directly to `F.linear` is intentionally invalid because it bypasses the +module's required `pre_quant_scale`; the full layer replay above is the valid +gate and is exact. + ## Current Runtime Scope | Component | Implemented | Known limitation | diff --git a/src/h3_blackwell_runtime/nvfp4.py b/src/h3_blackwell_runtime/nvfp4.py index 4587cd6..ec9afde 100644 --- a/src/h3_blackwell_runtime/nvfp4.py +++ b/src/h3_blackwell_runtime/nvfp4.py @@ -67,8 +67,8 @@ class Nvfp4Linear(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: if x.shape[-1] != self.in_features: raise ValueError(f"Expected feature width {self.in_features}, received {x.shape[-1]}.") - if x.dtype not in (torch.float16, torch.bfloat16): - raise ValueError("NVFP4 linear accepts FP16 or BF16 activations.") + if x.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise ValueError("NVFP4 linear accepts FP16, BF16, or FP32 activations.") from comfy_kitchen.tensor import QuantizedTensor @@ -82,6 +82,8 @@ class Nvfp4Linear(nn.Module): weight = packed_weight.dequantize().to(flat_x) output = functional.linear(flat_x, weight, bias) return output.reshape(*original_shape, self.out_features) + if x.dtype == torch.float32: + raise ValueError("Quantized NVFP4 activation GEMM requires FP16 or BF16 activations.") packed_x = QuantizedTensor.from_float(flat_x, "TensorCoreNVFP4Layout") output = functional.linear(packed_x, packed_weight, bias) return output[:flat_x.shape[0], :self.out_features].reshape(*original_shape, self.out_features) diff --git a/src/h3_blackwell_runtime/qwen3vl_text.py b/src/h3_blackwell_runtime/qwen3vl_text.py index d060d65..0388086 100644 --- a/src/h3_blackwell_runtime/qwen3vl_text.py +++ b/src/h3_blackwell_runtime/qwen3vl_text.py @@ -133,7 +133,7 @@ class Qwen3VL32BTextEncoder(nn.Module): """Mounted-checkpoint Qwen3-VL prompt conditioner returning layer-50 states.""" config = Qwen3VL32BTextConfig() - def __init__(self, checkpoint_path: str | Path, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.bfloat16): + def __init__(self, checkpoint_path: str | Path, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.float32): super().__init__() self.checkpoint_path = str(checkpoint_path) self.device_name = str(device) @@ -195,7 +195,7 @@ class Qwen3VL32BTextEncoder(nn.Module): class Qwen3VLPromptConditioner: """Tokenize raw H3 prompt text and produce Qwen layer-50 conditioning.""" - def __init__(self, checkpoint_path: str | Path, tokenizer_dir: str | Path | None = None, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.bfloat16): + def __init__(self, checkpoint_path: str | Path, tokenizer_dir: str | Path | None = None, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.float32): from .conditioning import H3PromptTokenizer tokenizer_dir = tokenizer_dir or Path(__file__).with_name("qwen25_tokenizer") diff --git a/tools/compare_qwen0_quantized_dispatch.py b/tools/compare_qwen0_quantized_dispatch.py new file mode 100644 index 0000000..0be8ee0 --- /dev/null +++ b/tools/compare_qwen0_quantized_dispatch.py @@ -0,0 +1,46 @@ +"""Compare direct NVFP4 wrapper and effective Comfy BF16 weights.""" + +import argparse + +import torch +import torch.nn.functional as functional + +from h3_blackwell_runtime.qwen3vl_text import Qwen3VL32BTextEncoder + + +parser = argparse.ArgumentParser() +parser.add_argument("--capture-dir", required=True) +parser.add_argument("--checkpoint", default="/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors") +parser.add_argument("--dtype", choices=("bfloat16", "float32"), default="bfloat16") +args = parser.parse_args() + +encoder = Qwen3VL32BTextEncoder(args.checkpoint, dtype=getattr(torch, args.dtype)).eval() +layer = encoder.layers[0] +qkv = torch.load(f"{args.capture_dir}/qwen0_quantized_qkv.pt", map_location="cuda", weights_only=False) +mlp = torch.load(f"{args.capture_dir}/qwen0_quantized_mlp.pt", map_location="cuda", weights_only=False) + +for group, capture, modules, inputs in ( + ("qkv", qkv, {"q": layer.q_proj, "k": layer.k_proj, "v": layer.v_proj}, {"q": qkv["input"], "k": qkv["input"], "v": qkv["input"]}), + ("mlp", mlp, {"gate": layer.gate_proj, "up": layer.up_proj, "down": layer.down_proj}, {"gate": mlp["input"], "up": mlp["input"], "down": mlp["activated"]}), +): + for name, module in modules.items(): + expected = capture["modules"][name] + packed = module._packed_weight() + params = packed._params + qdata, scale, block_scale = packed.layout_cls.get_plain_tensors(packed) + reference_qdata = expected["qdata"].as_subclass(torch.Tensor).cpu() + print(f"{group}.{name}.qdata exact={torch.equal(qdata.cpu(), reference_qdata)} shape={tuple(qdata.shape)}") + for field, actual in (("scale", scale), ("block_scale", block_scale)): + reference = expected[field].to(actual.device) + print(f"{group}.{name}.{field} exact={torch.equal(actual, reference)} shape={tuple(actual.shape)} dtype={actual.dtype}") + value = inputs[name].to(encoder.dtype) + with torch.inference_mode(): + dequantized = packed.dequantize() + direct = module(value) + effective = expected["effective_weight"].to("cuda") + comfy_effective = functional.linear(value.to(effective.dtype), effective, None) + print(f"{group}.{name}.effective dtype={expected['effective_dtype']} shape={expected['effective_shape']} direct_dtype={encoder.dtype}") + for variant, actual in (("weight", dequantized), ("linear", direct), ("effective_linear", comfy_effective)): + reference = expected["effective_weight"].to(actual.device) if variant == "weight" else expected["output"].to(actual.device) + delta = (actual.float() - reference.float()).abs() + print(f"{group}.{name}.{variant}_vs_comfy max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}") diff --git a/tools/patch_comfy_qwen_quantized_tensor_probe.py b/tools/patch_comfy_qwen_quantized_tensor_probe.py new file mode 100644 index 0000000..34d68f1 --- /dev/null +++ b/tools/patch_comfy_qwen_quantized_tensor_probe.py @@ -0,0 +1,58 @@ +"""Capture loaded Comfy Qwen QuantizedTensor fields and dispatch variants.""" + +from pathlib import Path + + +path = Path("/opt/ComfyUI/comfy/text_encoders/llama.py") +source = path.read_text(encoding="utf-8") +helper = ( + "def _h3_owned_cpu(tensor):\n" + " return tensor.detach().to(device=\"cpu\", copy=True).contiguous().clone()\n\n" +) +if "def _h3_owned_cpu(tensor):\n" not in source: + source = helper + source +old = " xq = self.q_proj(hidden_states)\n xk = self.k_proj(hidden_states)\n xv = self.v_proj(hidden_states)\n" +new = old + ( + " capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n" + " if capture_dir:\n" + " modules = {\"q\": self.q_proj, \"k\": self.k_proj, \"v\": self.v_proj}\n" + " def probe(module):\n" + " stored = module.weight\n" + " qdata, scale, block_scale = stored.layout_cls.get_plain_tensors(stored)\n" + " weight, bias, offload_stream = comfy.ops.cast_bias_weight(module, hidden_states, offloadable=True)\n" + " try:\n" + " output = torch.nn.functional.linear(hidden_states, weight, bias)\n" + " return {\"qdata\": _h3_owned_cpu(qdata), \"scale\": _h3_owned_cpu(scale), \"block_scale\": _h3_owned_cpu(block_scale), \"layout_class\": stored._layout_cls, \"orig_dtype\": stored._params.orig_dtype, \"orig_shape\": tuple(stored._params.orig_shape), \"effective_weight\": _h3_owned_cpu(weight), \"effective_dtype\": str(weight.dtype), \"effective_shape\": tuple(weight.shape), \"output\": _h3_owned_cpu(output)}\n" + " finally:\n" + " comfy.ops.uncast_bias_weight(module, weight, bias, offload_stream)\n" + " torch.save({\"input\": hidden_states.detach().cpu(), \"modules\": {name: probe(module) for name, module in modules.items()}}, os.path.join(capture_dir, \"qwen0_quantized_qkv.pt\"))\n" +) +if source.count(old) != 1: + raise RuntimeError("Unable to locate Qwen QKV projections.") +source = source.replace(old, new) + +old = " return self.down_proj(self.activation(self.gate_proj(x)) * self.up_proj(x))\n" +new = ( + " gate = self.gate_proj(x)\n" + " up = self.up_proj(x)\n" + " activated = self.activation(gate) * up\n" + " output = self.down_proj(activated)\n" + " capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_trace_index\", -1) == 0 else None\n" + " if capture_dir:\n" + " modules = {\"gate\": self.gate_proj, \"up\": self.up_proj, \"down\": self.down_proj}\n" + " def probe(module, value):\n" + " stored = module.weight\n" + " qdata, scale, block_scale = stored.layout_cls.get_plain_tensors(stored)\n" + " weight, bias, offload_stream = comfy.ops.cast_bias_weight(module, value, offloadable=True)\n" + " try:\n" + " output = torch.nn.functional.linear(value, weight, bias)\n" + " return {\"qdata\": _h3_owned_cpu(qdata), \"scale\": _h3_owned_cpu(scale), \"block_scale\": _h3_owned_cpu(block_scale), \"layout_class\": stored._layout_cls, \"orig_dtype\": stored._params.orig_dtype, \"orig_shape\": tuple(stored._params.orig_shape), \"effective_weight\": _h3_owned_cpu(weight), \"effective_dtype\": str(weight.dtype), \"effective_shape\": tuple(weight.shape), \"output\": _h3_owned_cpu(output)}\n" + " finally:\n" + " comfy.ops.uncast_bias_weight(module, weight, bias, offload_stream)\n" + " torch.save({\"input\": x.detach().cpu(), \"activated\": activated.detach().cpu(), \"modules\": {\"gate\": probe(self.gate_proj, x), \"up\": probe(self.up_proj, x), \"down\": probe(self.down_proj, activated)}}, os.path.join(capture_dir, \"qwen0_quantized_mlp.pt\"))\n" + " return output\n" +) +if source.count(old) != 1: + raise RuntimeError("Unable to locate Qwen MLP projections.") +path.write_text(source.replace(old, new), encoding="utf-8") +print("Applied Qwen QuantizedTensor dispatch probe.")