diff --git a/PARITY.md b/PARITY.md index e75feaf..98e14f3 100644 --- a/PARITY.md +++ b/PARITY.md @@ -181,7 +181,7 @@ gate and is exact. | Component | Implemented | Known limitation | | --- | --- | --- | | Text-only Qwen | Yes | No vision encoder, MRoPE, image/video expansion, reference labels, or modality tags | -| Token refiner | Yes | Architecture ported; no completed direct-versus-Comfy refiner tensor comparison | +| Token refiner | Yes | Bit-exact from captured 5376-wide refiner input through both blocks and final RMSNorm; Qwen-to-refiner projection boundary is still not separately captured | | Prompt-only FL2VA packer | Yes | No keyframe/reference condition rows; all preview work is text-only FL2VA | | H3 DiT | Yes | Strict all-block numeric parity not achieved | | Beta/RES sampler | Yes | Exact sigma/state/update parity not demonstrated; direct preview has wrong initial audio state | @@ -194,27 +194,24 @@ gate and is exact. Only these are outstanding. Do not recapture or revisit rows marked complete unless the checkpoint, Comfy version, prompt, or backend changes. -1. **Token-refiner capture and comparison.** Capture exact Comfy pre-refiner, - block 0, block 1, and final refined states once, then compare direct using - the matched Qwen layer-50 tensor. -2. **H3 assembled input comparison.** Use the existing FL2VA input capture to +1. **H3 assembled input comparison.** Use the existing FL2VA input capture to compare text rows, audio/video rows, positions, times, and segments. This isolates packing from DiT error. -3. **H3 first-divergence repair.** Reuse the existing block trace. Start at +2. **H3 first-divergence repair.** Reuse the existing block trace. Start at block 0 sublayers, repair the first mismatch, then use Comfy block output as the next direct input to distinguish local error from accumulation. Existing evidence says block 0 is close but accumulation is unresolved. -4. **Final rows and video latent.** Compare final packed video rows and +3. **Final rows and video latent.** Compare final packed video rows and `unpatchify_video` output with Comfy before entering the sampler. -5. **Sampler replay using Comfy state.** Feed `initial.pt` and captured +4. **Sampler replay using Comfy state.** Feed `initial.pt` and captured `sigmas` directly to the sampler. Compare video and audio denoised/update states for all 12 saved steps. Do not test seed equality until this passes. -6. **VAE pixel gate.** Decode the identical captured final video latent in +5. **VAE pixel gate.** Decode the identical captured final video latent in direct and Comfy and compare pixels before ffmpeg. This decides whether any residual grid comes from decoder behavior. -7. **End-to-end FL2VA preview.** Generate the reference dragon only after +6. **End-to-end FL2VA preview.** Generate the reference dragon only after gates 1-8 pass. Compare its raw final latent first, then video. -8. **Feature/performance work.** Only then add audio, Ref2VA/reference paths, +7. **Feature/performance work.** Only then add audio, Ref2VA/reference paths, Sage3, CUDA graphs, and multi-GPU execution. ## Existing Tools And Their Intended Gate diff --git a/tools/compare_refiner_trace.py b/tools/compare_refiner_trace.py new file mode 100644 index 0000000..88cd6ed --- /dev/null +++ b/tools/compare_refiner_trace.py @@ -0,0 +1,29 @@ +"""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}") diff --git a/tools/patch_comfy_h3_refiner_trace.py b/tools/patch_comfy_h3_refiner_trace.py new file mode 100644 index 0000000..7afa050 --- /dev/null +++ b/tools/patch_comfy_h3_refiner_trace.py @@ -0,0 +1,34 @@ +"""Capture Comfy H3 token-refiner boundaries once for direct parity replay.""" + +from pathlib import Path + + +path = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py") +source = path.read_text(encoding="utf-8") +if "import os\n" not in source: + source = source.replace("import math\n", "import math\nimport os\n", 1) + +old = ( + " def forward(self, x, transformer_options={}):\n" + " for block in self.blocks:\n" + " x = block(x, transformer_options=transformer_options)\n" + " return self.final_norm(x)\n" +) +new = ( + " def forward(self, x, transformer_options={}):\n" + " capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n" + " if capture_dir:\n" + " torch.save(x.detach().cpu(), os.path.join(capture_dir, \"refiner_input.pt\"))\n" + " for index, block in enumerate(self.blocks):\n" + " x = block(x, transformer_options=transformer_options)\n" + " if capture_dir:\n" + " torch.save(x.detach().cpu(), os.path.join(capture_dir, f\"refiner_block{index}.pt\"))\n" + " output = self.final_norm(x)\n" + " if capture_dir:\n" + " torch.save(output.detach().cpu(), os.path.join(capture_dir, \"refiner_output.pt\"))\n" + " return output\n" +) +if source.count(old) != 1: + raise RuntimeError("Unable to locate TokenRefiner.forward.") +path.write_text(source.replace(old, new), encoding="utf-8") +print("Applied H3 token-refiner trace patch.")