28 lines
1.2 KiB
Python
28 lines
1.2 KiB
Python
|
"""Find the AdaLN linear precision boundary that matches ComfyUI."""
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torch.nn.functional as functional
|
||
|
|
|
||
|
|
|
||
|
|
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)
|
||
|
|
embedding = inputs["t_emb"]
|
||
|
|
weight = expected["weight"]
|
||
|
|
bias = expected["bias"]
|
||
|
|
hidden = expected["shift"].shape[-1]
|
||
|
|
|
||
|
|
variants = {
|
||
|
|
"fp32": (embedding.float(), weight.float(), bias.float()),
|
||
|
|
"bf16_weight_fp32": (embedding.float(), weight.bfloat16().float(), bias.float()),
|
||
|
|
"bf16_all": (embedding.bfloat16(), weight.bfloat16(), bias.bfloat16()),
|
||
|
|
"bf16_input_fp32": (embedding.bfloat16().float(), weight.float(), bias.float()),
|
||
|
|
}
|
||
|
|
|
||
|
|
for name, (x, w, b) in variants.items():
|
||
|
|
values = functional.linear(x, w, b).float().view(x.shape[0] * 3, 6 * hidden)
|
||
|
|
shift, scale, *_ = values.chunk(6, dim=-1)
|
||
|
|
shift_delta = (shift - expected["shift"].float()).abs()
|
||
|
|
scale_delta = (scale - expected["scale"].float()).abs()
|
||
|
|
print(f"{name} shift_mean={shift_delta.mean().item():.6g} shift_max={shift_delta.max().item():.6g} scale_mean={scale_delta.mean().item():.6g} scale_max={scale_delta.max().item():.6g}")
|