42 lines
2.2 KiB
Python
42 lines
2.2 KiB
Python
"""Compare one direct H3 block's intermediates with a matching Comfy capture."""
|
|
|
|
import argparse
|
|
|
|
import torch
|
|
|
|
from h3_blackwell_runtime.attention import rms_norm
|
|
from h3_blackwell_runtime.block import gate_segments, modulate_segments
|
|
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
|
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
|
|
from h3_blackwell_runtime.rope import h3_rope_rotation
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--capture-dir", required=True)
|
|
parser.add_argument("--model", required=True)
|
|
parser.add_argument("--block", type=int, required=True)
|
|
args = parser.parse_args()
|
|
|
|
inputs = torch.load(f"{args.capture_dir}/input.pt", map_location="cuda", weights_only=False)
|
|
model = H3PackedDenoiser.from_checkpoint(H3Checkpoint(args.model)).eval()
|
|
rotation = h3_rope_rotation(inputs["position_ids"], model.backbone.inv_freq, inputs["hidden"].dtype)
|
|
hidden = inputs["hidden"]
|
|
|
|
with torch.inference_mode():
|
|
for index in range(args.block):
|
|
block = model.backbone.blocks[index]
|
|
hidden = block(hidden, rotation, *model.backbone.adaln[index](inputs["timesteps"]), inputs["segments"])
|
|
|
|
block = model.backbone.blocks[args.block]
|
|
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = model.backbone.adaln[args.block](inputs["timesteps"])
|
|
norm1 = modulate_segments(rms_norm(hidden, block.norm1_weight, block.norm_eps), shift_msa, scale_msa, inputs["segments"])
|
|
attention = block.attention(norm1, rotation)
|
|
post_attention = gate_segments(hidden, attention, gate_msa, inputs["segments"])
|
|
norm2 = modulate_segments(rms_norm(post_attention, block.norm2_weight, block.norm_eps), shift_mlp, scale_mlp, inputs["segments"])
|
|
mlp = block.mlp(norm2)
|
|
post_mlp = gate_segments(post_attention, mlp, gate_mlp, inputs["segments"])
|
|
|
|
for name, actual in (("norm1", norm1), ("attention", attention), ("post_attention", post_attention), ("norm2", norm2), ("mlp", mlp), ("post_mlp", post_mlp)):
|
|
expected = torch.load(f"{args.capture_dir}/block{args.block}_{name}.pt", map_location="cuda", weights_only=False)
|
|
delta = (actual.float() - expected.float()).abs()
|
|
print(f"{name} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|