30 lines
1.4 KiB
Python
30 lines
1.4 KiB
Python
|
"""Replay the direct token refiner against a Comfy boundary trace."""
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
||
|
|
from h3_blackwell_runtime.token_refiner import H3TokenRefiner
|
||
|
|
|
||
|
|
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument("--trace-dir", type=Path, required=True)
|
||
|
|
parser.add_argument("--checkpoint", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
refiner = H3TokenRefiner(H3Checkpoint(args.checkpoint), attention_backend="sage2").eval()
|
||
|
|
hidden = torch.load(args.trace_dir / "refiner_input.pt", map_location="cuda", weights_only=False)
|
||
|
|
|
||
|
|
with torch.inference_mode():
|
||
|
|
for index, block in enumerate(refiner.blocks):
|
||
|
|
hidden = block(hidden)
|
||
|
|
expected = torch.load(args.trace_dir / f"refiner_block{index}.pt", map_location="cuda", weights_only=False)
|
||
|
|
delta = (hidden.float() - expected.float()).abs()
|
||
|
|
print(f"block={index} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|
||
|
|
output = refiner.final_norm
|
||
|
|
output = torch.nn.functional.rms_norm(hidden, output.shape, weight=output.to(hidden), eps=1e-5).unsqueeze(0)
|
||
|
|
expected = torch.load(args.trace_dir / "refiner_output.pt", map_location="cuda", weights_only=False)
|
||
|
|
delta = (output.float() - expected.float()).abs()
|
||
|
|
print(f"output max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}")
|