37 lines
1.9 KiB
Python
37 lines
1.9 KiB
Python
"""Compare direct FP32 patch GEMMs under Comfy's captured matmul policy."""
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import torch.nn.functional as functional
|
|
|
|
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--capture", type=Path, required=True)
|
|
parser.add_argument("--checkpoint", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
|
|
args = parser.parse_args()
|
|
|
|
captured = torch.load(args.capture, map_location="cuda", weights_only=False)
|
|
checkpoint = H3Checkpoint(args.checkpoint)
|
|
previous_tf32 = torch.backends.cuda.matmul.allow_tf32
|
|
previous_precision = torch.get_float32_matmul_precision()
|
|
torch.backends.cuda.matmul.allow_tf32 = captured["matmul"]["allow_tf32"]
|
|
torch.set_float32_matmul_precision(captured["matmul"]["precision"])
|
|
try:
|
|
for name in ("video", "audio"):
|
|
rows = captured[f"{name}_rows"].to("cuda")
|
|
weight = checkpoint.tensor(f"{name}_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32)
|
|
bias = checkpoint.tensor(f"{name}_patch_proj.bias", dtype=torch.bfloat16).to(torch.float32)
|
|
output = functional.linear(rows, weight, bias)
|
|
for label, actual, expected in (("weight", weight, captured[f"{name}_weight"]), ("bias", bias, captured[f"{name}_bias"]), ("fp32", output, captured[f"{name}_embed_fp32"])):
|
|
delta = (actual.float() - expected.to(actual.device).float()).abs()
|
|
print(f"{name}.{label} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|
|
bf16_delta = (output.to(torch.bfloat16).float() - captured[f"{name}_embed_fp32"].to(torch.bfloat16).float()).abs()
|
|
print(f"{name}.bf16 max_abs={bf16_delta.max().item():.6g} mean_abs={bf16_delta.mean().item():.6g}")
|
|
print(f"matmul={captured['matmul']}")
|
|
finally:
|
|
torch.backends.cuda.matmul.allow_tf32 = previous_tf32
|
|
torch.set_float32_matmul_precision(previous_precision)
|