From 807b54ca0e09aa3734fb9c4956cfda1ed4c0cda8 Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Thu, 13 Aug 2026 02:59:46 +0700 Subject: [PATCH] Trace H3 final AdaLN parity --- src/h3_blackwell_runtime/final.py | 17 ++++++++-------- tools/patch_comfy_h3_capture.py | 4 ++-- tools/patch_comfy_h3_final_capture.py | 10 +++++++-- tools/patch_comfy_h3_temb_trace.py | 29 +++++++++++++++++++++++++++ 4 files changed, 47 insertions(+), 13 deletions(-) create mode 100644 tools/patch_comfy_h3_temb_trace.py diff --git a/src/h3_blackwell_runtime/final.py b/src/h3_blackwell_runtime/final.py index c0d16d4..ed89e5e 100644 --- a/src/h3_blackwell_runtime/final.py +++ b/src/h3_blackwell_runtime/final.py @@ -42,11 +42,11 @@ class H3FinalLayer(nn.Module): return cls( checkpoint.tensor("adaln_t_table", dtype=torch.float32), checkpoint.tensor("final_layer.norm.weight", dtype=output_dtype), - checkpoint.tensor("final_layer.adaln_proj.linear.weight", dtype=torch.float32), - checkpoint.tensor("final_layer.adaln_proj.linear.bias", dtype=torch.float32), - checkpoint.tensor("final_layer.video_out.weight", dtype=torch.float32), + checkpoint.tensor("final_layer.adaln_proj.linear.weight", dtype=torch.bfloat16), + checkpoint.tensor("final_layer.adaln_proj.linear.bias", dtype=torch.float16), + checkpoint.tensor("final_layer.video_out.weight", dtype=torch.bfloat16), checkpoint.tensor("final_layer.video_out.bias", dtype=torch.float32), - checkpoint.tensor("final_layer.audio_out.weight", dtype=torch.float32), + checkpoint.tensor("final_layer.audio_out.weight", dtype=torch.bfloat16), checkpoint.tensor("final_layer.audio_out.bias", dtype=torch.float32), ) @@ -64,10 +64,9 @@ class H3FinalLayer(nn.Module): video_start, video_stop, video_row = video_segment audio_start, audio_stop, audio_row = audio_segment - normalized = rms_norm(hidden, self.norm_weight, self.eps) - video_hidden = normalized[video_start:video_stop] * (1 + scale[video_row].to(hidden.dtype)) + shift[video_row].to(hidden.dtype) - audio_hidden = normalized[audio_start:audio_stop] * (1 + scale[audio_row].to(hidden.dtype)) + shift[audio_row].to(hidden.dtype) + video_hidden = (rms_norm(hidden[video_start:video_stop], self.norm_weight, self.eps) * (1.0 + scale[video_row]) + shift[video_row]).to(torch.float32) + audio_hidden = (rms_norm(hidden[audio_start:audio_stop], self.norm_weight, self.eps) * (1.0 + scale[audio_row]) + shift[audio_row]).to(torch.float32) return ( - functional.linear(video_hidden.float(), self.video_weight, self.video_bias), - functional.linear(audio_hidden.float(), self.audio_weight, self.audio_bias), + functional.linear(video_hidden, self.video_weight, self.video_bias), + functional.linear(audio_hidden, self.audio_weight, self.audio_bias), ) diff --git a/tools/patch_comfy_h3_capture.py b/tools/patch_comfy_h3_capture.py index ffb22b5..20f0bab 100644 --- a/tools/patch_comfy_h3_capture.py +++ b/tools/patch_comfy_h3_capture.py @@ -24,14 +24,14 @@ 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, \"video_x\": video_x.detach().cpu(), \"audio_x\": audio_x.detach().cpu()}, os.path.join(capture_dir, \"input.pt\"))\n\n" + " torch.save({\"hidden\": h.detach().cpu(), \"timesteps\": t_vals.detach().cpu(), \"t_emb\": t_emb.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( model, " return [-video_out.to(video_x.dtype), -audio_out.to(audio_x.dtype)]\n", " if capture_dir and H3_CAPTURE_ACTIVE:\n" - " torch.save({\"video\": video_out.detach().cpu(), \"audio\": audio_out.detach().cpu(), \"video_segment\": video_seg, \"audio_segment\": audio_seg}, os.path.join(capture_dir, \"output.pt\"))\n" + " torch.save({\"video\": video_out.detach().cpu(), \"audio\": audio_out.detach().cpu(), \"video_segment\": video_seg, \"audio_segment\": audio_seg, \"adaln_t_table\": self.adaln_t_table.detach().cpu()}, os.path.join(capture_dir, \"output.pt\"))\n" " H3_CAPTURE_ACTIVE = False\n" " return [-video_out.to(video_x.dtype), -audio_out.to(audio_x.dtype)]\n", ) diff --git a/tools/patch_comfy_h3_final_capture.py b/tools/patch_comfy_h3_final_capture.py index aec2e74..7488436 100644 --- a/tools/patch_comfy_h3_final_capture.py +++ b/tools/patch_comfy_h3_final_capture.py @@ -14,7 +14,11 @@ old = ( " return self.video_out(hv), self.audio_out(ha)\n" ) new = ( - " shift, scale = self.adaln_proj(t_emb)\n" + " capture_input = t_emb.clone() if H3_CAPTURE_ACTIVE else None\n" + " adaln_output = self.adaln_proj(t_emb)\n" + " capture_replay = self.adaln_proj.linear(capture_input.clone()) if H3_CAPTURE_ACTIVE else None\n" + " capture_forward = self.adaln_proj.linear._forward(capture_input.clone(), self.adaln_proj.linear.weight, self.adaln_proj.linear.bias) if H3_CAPTURE_ACTIVE else None\n" + " shift, scale = adaln_output\n" " va, vb, vrow = video_seg\n" " aa, ab, arow = audio_seg\n" " norm_v = self.norm(x[va:vb])\n" @@ -25,7 +29,9 @@ new = ( " audio = self.audio_out(ha)\n" " capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if H3_CAPTURE_ACTIVE else None\n" " if capture_dir:\n" - " torch.save({\"hidden\": x.detach().cpu(), \"norm_v\": norm_v.detach().cpu(), \"norm_a\": norm_a.detach().cpu(), \"shift\": shift.detach().cpu(), \"scale\": scale.detach().cpu(), \"video_hidden\": hv.detach().cpu(), \"audio_hidden\": ha.detach().cpu(), \"video\": video.detach().cpu(), \"audio\": audio.detach().cpu()}, os.path.join(capture_dir, \"final.pt\"))\n" + " torch.cuda.synchronize(x.device)\n" + " linear = self.adaln_proj.linear\n" + " torch.save({\"hidden\": x.detach().cpu(), \"adaln_input\": capture_input.detach().cpu(), \"adaln_output\": torch.cat((shift, scale), dim=-1).detach().cpu(), \"adaln_replay\": capture_replay.detach().cpu(), \"adaln_forward\": capture_forward.detach().cpu(), \"t_emb\": t_emb.detach().cpu(), \"norm_v\": norm_v.detach().cpu(), \"norm_a\": norm_a.detach().cpu(), \"shift\": shift.detach().cpu(), \"scale\": scale.detach().cpu(), \"adaln_weight\": linear.weight.detach().cpu(), \"adaln_bias\": linear.bias.detach().cpu(), \"adaln_metadata\": {\"class\": f\"{type(linear).__module__}.{type(linear).__qualname__}\", \"weight_class\": f\"{type(linear.weight).__module__}.{type(linear.weight).__qualname__}\", \"quant_format\": getattr(linear, \"quant_format\", None), \"layout_type\": getattr(linear, \"layout_type\", None), \"full_precision_mm\": getattr(linear, \"_full_precision_mm\", None), \"input_scale\": getattr(linear, \"input_scale\", None)}, \"video_hidden\": hv.detach().cpu(), \"audio_hidden\": ha.detach().cpu(), \"video_weight\": self.video_out.weight.detach().cpu(), \"video_bias\": self.video_out.bias.detach().cpu(), \"audio_weight\": self.audio_out.weight.detach().cpu(), \"audio_bias\": self.audio_out.bias.detach().cpu(), \"video\": video.detach().cpu(), \"audio\": audio.detach().cpu(), \"matmul\": {\"allow_tf32\": torch.backends.cuda.matmul.allow_tf32, \"precision\": torch.get_float32_matmul_precision()}}, os.path.join(capture_dir, \"final.pt\"))\n" " return video, audio\n" ) if source.count(old) == 1: diff --git a/tools/patch_comfy_h3_temb_trace.py b/tools/patch_comfy_h3_temb_trace.py new file mode 100644 index 0000000..7f8fcc7 --- /dev/null +++ b/tools/patch_comfy_h3_temb_trace.py @@ -0,0 +1,29 @@ +"""Capture H3 curve embeddings after every transformer block.""" + +from pathlib import Path + + +model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py") +source = model.read_text(encoding="utf-8") +old = ( + " else:\n" + " h = block(h, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options)\n" + " if prefetch_queue is not None:\n" +) +new = ( + " else:\n" + " if capture_dir and H3_CAPTURE_ACTIVE and i == 0:\n" + " torch.save({\"t_emb\": t_emb.clone().detach().cpu(), \"timesteps\": t_vals.clone().detach().cpu()}, os.path.join(capture_dir, \"t_emb_before_block0.pt\"))\n" + " h = block(h, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options)\n" + " if capture_dir and H3_CAPTURE_ACTIVE:\n" + " temb_dir = os.path.join(capture_dir, \"t_emb\")\n" + " os.makedirs(temb_dir, exist_ok=True)\n" + " torch.save({\"value\": t_emb.detach().cpu(), \"class\": f\"{type(block.adaln_proj.linear).__module__}.{type(block.adaln_proj.linear).__qualname__}\", \"data_ptr\": t_emb.data_ptr(), \"stride\": t_emb.stride()}, os.path.join(temb_dir, f\"{i:02d}.pt\"))\n" + " if prefetch_queue is not None:\n" +) +if source.count(old) == 1: + source = source.replace(old, new) +elif new not in source: + raise RuntimeError("Unable to locate the H3 transformer block loop.") +model.write_text(source, encoding="utf-8") +print("Applied H3 t_emb trace patch.")