22 lines
1.1 KiB
Python
22 lines
1.1 KiB
Python
"""Identify the ComfyUI-compatible block-0 RMSNorm precision path."""
|
|
|
|
import torch
|
|
|
|
from h3_blackwell_runtime.block import modulate_segments
|
|
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
|
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
|
|
|
|
|
|
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.pt", map_location="cuda", weights_only=False)
|
|
checkpoint = H3Checkpoint("/models/minimax_h3_ref2va_pruned_nvfp4.safetensors")
|
|
model = H3PackedDenoiser.from_checkpoint(checkpoint).eval()
|
|
shift, scale, *_ = model.backbone.adaln[0](inputs["timesteps"])
|
|
raw_weight = checkpoint.tensor("blocks.0.norm1.weight")
|
|
|
|
for epsilon in (1e-6, 1e-5, 1e-4):
|
|
normalized = torch.nn.functional.rms_norm(inputs["hidden"], (inputs["hidden"].shape[-1],), raw_weight, epsilon)
|
|
actual = modulate_segments(normalized, shift, scale, inputs["segments"])
|
|
delta = (actual.float() - expected.float()).abs()
|
|
print(f"epsilon={epsilon:.0e} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|