diff --git a/PARITY.md b/PARITY.md index 98e14f3..c763a34 100644 --- a/PARITY.md +++ b/PARITY.md @@ -182,7 +182,7 @@ gate and is exact. | --- | --- | --- | | Text-only Qwen | Yes | No vision encoder, MRoPE, image/video expansion, reference labels, or modality tags | | 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 | +| Prompt-only FL2VA packer | Yes | Bit-exact for the coherent captured text-only FL2VA DiT input; no keyframe/reference condition rows | | 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 | | Video VAE decoder | Yes | No direct-versus-Comfy same-latent pixel comparison yet | @@ -194,24 +194,21 @@ 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. **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. -2. **H3 first-divergence repair.** Reuse the existing block trace. Start at +1. **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. -3. **Final rows and video latent.** Compare final packed video rows and +2. **Final rows and video latent.** Compare final packed video rows and `unpatchify_video` output with Comfy before entering the sampler. -4. **Sampler replay using Comfy state.** Feed `initial.pt` and captured +3. **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. -5. **VAE pixel gate.** Decode the identical captured final video latent in +4. **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. -6. **End-to-end FL2VA preview.** Generate the reference dragon only after +5. **End-to-end FL2VA preview.** Generate the reference dragon only after gates 1-8 pass. Compare its raw final latent first, then video. -7. **Feature/performance work.** Only then add audio, Ref2VA/reference paths, +6. **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/src/h3_blackwell_runtime/packing.py b/src/h3_blackwell_runtime/packing.py index 5a5d639..e34f9b1 100644 --- a/src/h3_blackwell_runtime/packing.py +++ b/src/h3_blackwell_runtime/packing.py @@ -20,7 +20,7 @@ def patchify_video(latent: torch.Tensor) -> torch.Tensor: def pack_audio(latent: torch.Tensor) -> torch.Tensor: - return latent[0].permute(1, 2, 0).reshape(-1, latent.shape[1]) + return latent[0].permute(1, 2, 0).reshape(-1, latent.shape[1]).transpose(0, 1).contiguous().transpose(0, 1) def unpatchify_video(rows: torch.Tensor, frames: int, latent_height: int, latent_width: int) -> torch.Tensor: @@ -58,10 +58,10 @@ class H3PromptPacker: """Build `[text | audio | video]` tokens for prompt-only H3 T2V.""" def __init__(self, checkpoint): - self.video_weight = checkpoint.tensor("video_patch_proj.weight", dtype=torch.float32) - self.video_bias = checkpoint.tensor("video_patch_proj.bias", dtype=torch.float32) - self.audio_weight = checkpoint.tensor("audio_patch_proj.weight", dtype=torch.float32) - self.audio_bias = checkpoint.tensor("audio_patch_proj.bias", dtype=torch.float32) + self.video_weight = checkpoint.tensor("video_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32) + self.video_bias = checkpoint.tensor("video_patch_proj.bias", dtype=torch.bfloat16).to(torch.float32) + self.audio_weight = checkpoint.tensor("audio_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32) + self.audio_bias = checkpoint.tensor("audio_patch_proj.bias", dtype=torch.bfloat16).to(torch.float32) self.text_weight = checkpoint.tensor("condition_proj.weight", dtype=torch.bfloat16) self.text_bias = checkpoint.tensor("condition_proj.bias", dtype=torch.bfloat16) @@ -84,7 +84,7 @@ class H3PromptPacker: row = {value: index for index, value in enumerate(unique_times)} video_row, audio_row = row[video_time] * 3, row[audio_time] * 3 times = torch.tensor(unique_times, device=hidden.device, dtype=torch.float32) - positions = torch.cat((torch.stack((torch.arange(text_length, dtype=torch.float64), torch.zeros(text_length), torch.zeros(text_length)), dim=-1), _audio_positions(audio.shape[-1], float(text_length), video.shape[-1], video.shape[-2]), _video_positions(video.shape[2], video.shape[-2], video.shape[-1], float(text_length) + audio.shape[-1]))) + positions = torch.cat((torch.stack((torch.arange(text_length, dtype=torch.float64), torch.zeros(text_length), torch.zeros(text_length)), dim=-1), _audio_positions(audio.shape[-1], float(text_length), video.shape[-1], video.shape[-2]), _video_positions(video.shape[2], video.shape[-2], video.shape[-1], float(text_length)))) block_video_segment = (text_length + audio_length, hidden.shape[0], video_row) block_audio_segment = (text_length, text_length + audio_length, audio_row + 2) final_video_segment = (text_length + audio_length, hidden.shape[0], row[video_time]) diff --git a/tools/compare_h3_assembled_input.py b/tools/compare_h3_assembled_input.py new file mode 100644 index 0000000..db1e88f --- /dev/null +++ b/tools/compare_h3_assembled_input.py @@ -0,0 +1,49 @@ +"""Compare direct prompt-only FL2VA packing with an immutable Comfy DiT input.""" + +import argparse +from pathlib import Path + +import torch +import torch.nn.functional as functional + +from h3_blackwell_runtime.checkpoint import H3Checkpoint +from h3_blackwell_runtime.packing import H3PromptPacker, pack_audio, patchify_video +from h3_blackwell_runtime.sampler import _audio_sigma + + +parser = argparse.ArgumentParser() +parser.add_argument("--capture", type=Path, required=True) +parser.add_argument("--refiner-trace", type=Path, required=True) +parser.add_argument("--checkpoint", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors") +args = parser.parse_args() + +captured = torch.load(args.capture, map_location="cuda", weights_only=False) +text = torch.load(args.refiner_trace / "refiner_output.pt", map_location="cuda", weights_only=False) +if text.ndim == 2: + text = text.unsqueeze(0) +checkpoint = H3Checkpoint(args.checkpoint) +packer = H3PromptPacker(checkpoint) + +video_shape = (1, 24, 7, 12, 20) +audio_shape = (1, 32, 2, 37) +video = captured["video_x"].to("cuda") +audio = captured["audio_x"].to("cuda") +sigma = 1 - captured["timesteps"].min() + +with torch.inference_mode(): + hidden, times, segments, positions, _, _ = packer(text, video, audio, float(sigma)) + text_rows = text[0].to(torch.bfloat16) + audio_rows = functional.linear(pack_audio(audio).float(), packer.audio_weight, packer.audio_bias).to(torch.bfloat16) + video_rows = functional.linear(patchify_video(video).float(), packer.video_weight, packer.video_bias).to(torch.bfloat16) + +for name, actual, expected in ( + ("text", text_rows, captured["hidden"][:17]), + ("audio", audio_rows, captured["hidden"][17:91]), + ("video", video_rows, captured["hidden"][91:]), + ("hidden", hidden, captured["hidden"]), + ("times", times, captured["timesteps"]), + ("positions", positions, captured["position_ids"]), +): + delta = (actual.float() - expected.to(actual.device).float()).abs() + print(f"{name} shape={tuple(actual.shape)} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}") +print(f"segments direct={segments} comfy={captured['segments']}") diff --git a/tools/compare_h3_loaded_patch_projection.py b/tools/compare_h3_loaded_patch_projection.py new file mode 100644 index 0000000..56c7d40 --- /dev/null +++ b/tools/compare_h3_loaded_patch_projection.py @@ -0,0 +1,30 @@ +"""Compare direct patch GEMMs with saved outputs from loaded Comfy modules.""" + +import argparse +from pathlib import Path + +import torch +import torch.nn.functional as functional + +from h3_blackwell_runtime.checkpoint import H3Checkpoint + + +parser = argparse.ArgumentParser() +parser.add_argument("--capture", type=Path, required=True) +parser.add_argument("--checkpoint", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors") +args = parser.parse_args() + +captured = torch.load(args.capture, map_location="cuda", weights_only=False) +checkpoint = H3Checkpoint(args.checkpoint) +for name, data in captured.items(): + rows = data["rows"].to("cuda") + weight = checkpoint.tensor(f"{name}_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32) + bias = checkpoint.tensor(f"{name}_patch_proj.bias", dtype=torch.float32) + with torch.inference_mode(): + direct = functional.linear(rows, weight, bias) + row_delta = (rows.float() - data["rows"].to(rows.device).float()).abs() + output_delta = (direct.float() - data["output"].to(direct.device).float()).abs() + forward_delta = (data["forward_cast"].float() - data["output"].float()).abs() + print(f"{name}.rows stride={tuple(rows.stride())} comfy_stride={data['rows_stride']} max_abs={row_delta.max().item():.6g}") + print(f"{name}.direct max_abs={output_delta.max().item():.6g} mean_abs={output_delta.mean().item():.6g}") + print(f"{name}.comfy_forward_cast max_abs={forward_delta.max().item():.6g} mean_abs={forward_delta.mean().item():.6g}") diff --git a/tools/compare_h3_patch_projections.py b/tools/compare_h3_patch_projections.py new file mode 100644 index 0000000..a2897be --- /dev/null +++ b/tools/compare_h3_patch_projections.py @@ -0,0 +1,37 @@ +"""Compare direct FP32 patch GEMMs under Comfy's captured matmul policy.""" + +import argparse +from pathlib import Path + +import torch +import torch.nn.functional as functional + +from h3_blackwell_runtime.checkpoint import H3Checkpoint + + +parser = argparse.ArgumentParser() +parser.add_argument("--capture", type=Path, required=True) +parser.add_argument("--checkpoint", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors") +args = parser.parse_args() + +captured = torch.load(args.capture, map_location="cuda", weights_only=False) +checkpoint = H3Checkpoint(args.checkpoint) +previous_tf32 = torch.backends.cuda.matmul.allow_tf32 +previous_precision = torch.get_float32_matmul_precision() +torch.backends.cuda.matmul.allow_tf32 = captured["matmul"]["allow_tf32"] +torch.set_float32_matmul_precision(captured["matmul"]["precision"]) +try: + for name in ("video", "audio"): + rows = captured[f"{name}_rows"].to("cuda") + weight = checkpoint.tensor(f"{name}_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32) + bias = checkpoint.tensor(f"{name}_patch_proj.bias", dtype=torch.bfloat16).to(torch.float32) + output = functional.linear(rows, weight, bias) + for label, actual, expected in (("weight", weight, captured[f"{name}_weight"]), ("bias", bias, captured[f"{name}_bias"]), ("fp32", output, captured[f"{name}_embed_fp32"])): + delta = (actual.float() - expected.to(actual.device).float()).abs() + print(f"{name}.{label} max_abs={delta.max().item():.6g} mean_abs={delta.mean().item():.6g}") + bf16_delta = (output.to(torch.bfloat16).float() - captured[f"{name}_embed_fp32"].to(torch.bfloat16).float()).abs() + print(f"{name}.bf16 max_abs={bf16_delta.max().item():.6g} mean_abs={bf16_delta.mean().item():.6g}") + print(f"matmul={captured['matmul']}") +finally: + torch.backends.cuda.matmul.allow_tf32 = previous_tf32 + torch.set_float32_matmul_precision(previous_precision) diff --git a/tools/patch_comfy_h3_capture.py b/tools/patch_comfy_h3_capture.py index afbe4ef..ffb22b5 100644 --- a/tools/patch_comfy_h3_capture.py +++ b/tools/patch_comfy_h3_capture.py @@ -24,7 +24,7 @@ replace_once( " if capture_dir and not H3_CAPTURE_ACTIVE:\n" " H3_CAPTURE_ACTIVE = True\n" " os.makedirs(capture_dir, exist_ok=True)\n" - " torch.save({\"hidden\": h.detach().cpu(), \"timesteps\": t_vals.detach().cpu(), \"position_ids\": layout.position_ids, \"segments\": mod_segments}, os.path.join(capture_dir, \"input.pt\"))\n\n" + " torch.save({\"hidden\": h.detach().cpu(), \"timesteps\": t_vals.detach().cpu(), \"position_ids\": layout.position_ids, \"segments\": mod_segments, \"video_x\": video_x.detach().cpu(), \"audio_x\": audio_x.detach().cpu()}, os.path.join(capture_dir, \"input.pt\"))\n\n" " # blocks\n patches_replace = transformer_options.get(\"patches_replace\", {})\n", ) replace_once( diff --git a/tools/patch_comfy_h3_loaded_patch_projection_probe.py b/tools/patch_comfy_h3_loaded_patch_projection_probe.py new file mode 100644 index 0000000..e2ed6dc --- /dev/null +++ b/tools/patch_comfy_h3_loaded_patch_projection_probe.py @@ -0,0 +1,33 @@ +"""Probe loaded Comfy H3 patch projections on their captured FP32 rows.""" + +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 = ( + " video_embed = self.video_patch_proj(all_video_rows).to(dtype)\n" + " audio_embed = self.audio_patch_proj(all_audio_rows).to(dtype)\n" +) +new = ( + " video_embed_fp32 = self.video_patch_proj(all_video_rows)\n" + " audio_embed_fp32 = self.audio_patch_proj(all_audio_rows)\n" + " capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n" + " if capture_dir and not os.path.exists(os.path.join(capture_dir, \"h3_loaded_patch_projection_probe.pt\")):\n" + " def probe(layer, rows, output):\n" + " weight, bias, state = comfy.ops.cast_bias_weight(layer, rows, offloadable=True)\n" + " try:\n" + " cast_output = torch.nn.functional.linear(rows, weight, bias)\n" + " return {\"rows\": rows.detach().cpu(), \"output\": output.detach().cpu(), \"forward_cast\": layer.forward_comfy_cast_weights(rows).detach().cpu(), \"effective_weight\": weight.detach().cpu(), \"effective_bias\": None if bias is None else bias.detach().cpu(), \"effective_weight_dtype\": str(weight.dtype), \"effective_bias_dtype\": None if bias is None else str(bias.dtype), \"cast_output\": cast_output.detach().cpu(), \"rows_dtype\": str(rows.dtype), \"rows_stride\": tuple(rows.stride()), \"weight_dtype\": str(layer.weight.dtype), \"weight_stride\": tuple(layer.weight.stride()), \"bias_dtype\": None if layer.bias is None else str(layer.bias.dtype), \"bias_stride\": None if layer.bias is None else tuple(layer.bias.stride()), \"autocast\": torch.is_autocast_enabled(), \"force_cast\": getattr(layer, \"comfy_force_cast_weights\", False), \"weight_functions\": len(layer.weight_function), \"bias_functions\": len(layer.bias_function)}\n" + " finally:\n" + " comfy.ops.uncast_bias_weight(layer, weight, bias, state)\n" + " torch.save({\"video\": probe(self.video_patch_proj, all_video_rows, video_embed_fp32), \"audio\": probe(self.audio_patch_proj, all_audio_rows, audio_embed_fp32)}, os.path.join(capture_dir, \"h3_loaded_patch_projection_probe.pt\"))\n" + " video_embed = video_embed_fp32.to(dtype)\n" + " audio_embed = audio_embed_fp32.to(dtype)\n" +) +if source.count(old) != 1: + raise RuntimeError("Unable to locate H3 patch-projection calls.") +path.write_text(source.replace(old, new), encoding="utf-8") +print("Applied loaded H3 patch-projection probe.") diff --git a/tools/patch_comfy_h3_patch_projection_probe.py b/tools/patch_comfy_h3_patch_projection_probe.py new file mode 100644 index 0000000..9b6c827 --- /dev/null +++ b/tools/patch_comfy_h3_patch_projection_probe.py @@ -0,0 +1,24 @@ +"""Capture Comfy H3 FP32 patch-projection inputs, outputs, and matmul policy.""" + +from pathlib import Path + + +path = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py") +source = path.read_text(encoding="utf-8") +old = ( + " video_embed = self.video_patch_proj(all_video_rows).to(dtype)\n" + " audio_embed = self.audio_patch_proj(all_audio_rows).to(dtype)\n" +) +new = ( + " video_embed_fp32 = self.video_patch_proj(all_video_rows)\n" + " audio_embed_fp32 = self.audio_patch_proj(all_audio_rows)\n" + " capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n" + " if capture_dir and not os.path.exists(os.path.join(capture_dir, \"h3_patch_projections.pt\")):\n" + " torch.save({\"video_rows\": all_video_rows.detach().cpu(), \"audio_rows\": all_audio_rows.detach().cpu(), \"video_embed_fp32\": video_embed_fp32.detach().cpu(), \"audio_embed_fp32\": audio_embed_fp32.detach().cpu(), \"video_weight\": self.video_patch_proj.weight.detach().cpu(), \"video_bias\": self.video_patch_proj.bias.detach().cpu(), \"audio_weight\": self.audio_patch_proj.weight.detach().cpu(), \"audio_bias\": self.audio_patch_proj.bias.detach().cpu(), \"matmul\": {\"allow_tf32\": torch.backends.cuda.matmul.allow_tf32, \"precision\": torch.get_float32_matmul_precision()}}, os.path.join(capture_dir, \"h3_patch_projections.pt\"))\n" + " video_embed = video_embed_fp32.to(dtype)\n" + " audio_embed = audio_embed_fp32.to(dtype)\n" +) +if source.count(old) != 1: + raise RuntimeError("Unable to locate H3 patch-projection calls.") +path.write_text(source.replace(old, new), encoding="utf-8") +print("Applied H3 patch-projection probe.")