25 lines
1.3 KiB
Python
25 lines
1.3 KiB
Python
|
"""Compare ComfyUI's effective AdaLN dispatch tensors with direct FP32 linear."""
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torch.nn.functional as functional
|
||
|
|
|
||
|
|
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
||
|
|
|
||
|
|
|
||
|
|
capture_dir = "/artifacts/capture"
|
||
|
|
inputs = torch.load(f"{capture_dir}/input.pt", map_location="cuda", weights_only=False)
|
||
|
|
expected = torch.load(f"{capture_dir}/block0_norm1_adaln.pt", map_location="cuda", weights_only=False)
|
||
|
|
weight = expected["effective_weight"]
|
||
|
|
bias = expected["effective_bias"]
|
||
|
|
curve = H3Checkpoint("/models/minimax_h3_ref2va_pruned_nvfp4.safetensors").tensor("adaln_t_table", dtype=torch.float32)
|
||
|
|
position = inputs["timesteps"].float().clamp(0, 1) * (curve.shape[0] - 1)
|
||
|
|
lower = position.floor().long().clamp(max=curve.shape[0] - 2)
|
||
|
|
embedding = torch.lerp(curve[lower], curve[lower + 1], (position - lower).unsqueeze(1))
|
||
|
|
values = functional.linear(embedding.to(weight.dtype), weight, bias).float().view(embedding.shape[0] * 3, 6 * expected["shift"].shape[-1])
|
||
|
|
shift, scale, *_ = values.chunk(6, dim=-1)
|
||
|
|
|
||
|
|
print(f"effective_weight_dtype={weight.dtype} effective_bias_dtype={bias.dtype}")
|
||
|
|
for name, actual, reference in (("shift", shift, expected["shift"]), ("scale", scale, expected["scale"])):
|
||
|
|
delta = (actual.float() - reference.float()).abs()
|
||
|
|
print(f"{name} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|