daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Trace H3 final AdaLN parity

Browse source
This commit is contained in:
Daniel Maddern 2026-08-13 02:59:46 +07:00
parent b69903b7e6
commit 807b54ca0e
4 changed files with 47 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
src/h3_blackwell_runtime/final.py
View file

@ -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),
)

4
tools/patch_comfy_h3_capture.py
View file

@ -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",
)

10
tools/patch_comfy_h3_final_capture.py
View file

@ -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:

29
tools/patch_comfy_h3_temb_trace.py Normal file
View file

@ -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.")