46 lines
2.7 KiB
Python
46 lines
2.7 KiB
Python
"""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}")
|